updated database classes

core/database/classes/__init__.py

@@ -1,3 +1,4 @@
+"""Database classes module content."""
 from __future__ import annotations
 
 from .dataset import Dataset

core/database/classes/dataset.py

@@ -1,14 +1,18 @@
+"""Definition of database Dataset class."""
 from __future__ import annotations
 
+import logging
 import random
 from datetime import datetime
 from datetime import UTC
 
 from pydantic import BaseModel
 from pydantic import Field
+from pymongo.collection import Collection
 
 
 class Dataset(BaseModel):
+    """Database Dataset model."""
     create_time: datetime = Field(default_factory=lambda: datetime.now(UTC))
     train_names: list[str]
     test_names: list[str]
@@ -53,4 +57,15 @@ class Dataset(BaseModel):
             test_names=test_name_list,
             validation_names=validation_name_list,
         )
+        logging.debug('finished')
         return dataset
+
+    def save(
+        self,
+        collection: Collection,
+    ) -> None:
+        """Save dataset to database."""
+        res = collection.insert_one(
+            document=self.model_dump(),
+        )
+        logging.debug('inserted document: %s', res)

core/database/classes/exceptions.py

@@ -1,5 +1,6 @@
+"""Definition of database exception."""
 from __future__ import annotations
 
 
 class NoDocumentFoundException(Exception):
-    pass
+    """Database exception for when no documents are found."""

core/database/classes/visual_communication.py

@@ -1,3 +1,4 @@
+"""Definition of VisualCommunication model."""
 from __future__ import annotations
 
 import logging
@@ -15,12 +16,15 @@ from core.dto import ModelData
 
 
 class VisualCommunication(BaseModel):
+    """Visual communication model."""
+
     name: str
     image: Image.Image
     annotation: ModelData | None = None
     prediction: ModelData | None = None
 
     class Config:
+        """BaseModel configuration."""
         arbitrary_types_allowed = True
 
     @classmethod
@@ -38,12 +42,14 @@ class VisualCommunication(BaseModel):
 
     @classmethod
     def decode_image(cls, content: str) -> Image.Image:
-        """Decode image."""
+        """Extract image from webencoded content."""
         _, content_data = content.split(',')
         return Image.open(BytesIO(b64decode(content_data)))
 
     @field_serializer('image')
-    def serialize_image(image: Image.Image) -> bytes:   # type: ignore
+    @classmethod
+    def serialize_image(cls, image: Image.Image) -> bytes:   # type: ignore
+        """Convert image to bytes for storage in database."""
         buffer = BytesIO()
         image.save(buffer, format='JPEG')
         return buffer.getvalue()
@@ -54,6 +60,7 @@ class VisualCommunication(BaseModel):
         cls,
         image: Image.Image | BytesIO | bytes,
     ) -> Image.Image:
+        """Convert bytes input from database into image."""
         if isinstance(image, bytes):
             image = BytesIO(image)
         if isinstance(image, BytesIO):