From b430a15ecea9d2ee0a2abfd5aa4624f766374b04 Mon Sep 17 00:00:00 2001
From: Brian Bjarke Jensen <bbj@skov.dk>
Date: Sun, 31 Mar 2024 11:24:28 +0200
Subject: [PATCH] updated database classes

---
 core/database/classes/__init__.py             |  1 +
 core/database/classes/dataset.py              | 15 +++++++++++++++
 core/database/classes/exceptions.py           |  3 ++-
 core/database/classes/visual_communication.py | 11 +++++++++--
 4 files changed, 27 insertions(+), 3 deletions(-)
 mode change 100644 => 100755 core/database/classes/__init__.py
 mode change 100644 => 100755 core/database/classes/dataset.py
 mode change 100644 => 100755 core/database/classes/exceptions.py
 mode change 100644 => 100755 core/database/classes/visual_communication.py

diff --git a/core/database/classes/__init__.py b/core/database/classes/__init__.py
old mode 100644
new mode 100755
index 295e6f2..9e260cd
--- a/core/database/classes/__init__.py
+++ b/core/database/classes/__init__.py
@@ -1,3 +1,4 @@
+"""Database classes module content."""
 from __future__ import annotations
 
 from .dataset import Dataset
diff --git a/core/database/classes/dataset.py b/core/database/classes/dataset.py
old mode 100644
new mode 100755
index ff3f4aa..0c3fa1a
--- a/core/database/classes/dataset.py
+++ b/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)
diff --git a/core/database/classes/exceptions.py b/core/database/classes/exceptions.py
old mode 100644
new mode 100755
index 56d1194..5372031
--- a/core/database/classes/exceptions.py
+++ b/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."""
diff --git a/core/database/classes/visual_communication.py b/core/database/classes/visual_communication.py
old mode 100644
new mode 100755
index 2f6b44b..78e9e61
--- a/core/database/classes/visual_communication.py
+++ b/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):