moved functions into separate files

core/database/utils.py

@@ -1,161 +0,0 @@
-from __future__ import annotations
-
-import logging
-
-from pymongo.collection import Collection
-
-from .classes import NoDocumentFoundException
-from .classes import VisualCommunication
-from core.dto import ModelData
-
-
-def count_documents(
-    collection: Collection,
-    has_annotation: bool = False,
-) -> int:
-    """
-    Get the total number of documents
-    in database that matches the filters.
-    """
-    assert isinstance(collection, Collection)
-    assert isinstance(has_annotation, bool)
-    # build query
-    query = {}
-    if has_annotation:
-        query['annotation'] = {'$ne': None}
-    return collection.count_documents(filter=query)
-
-
-def list_names(
-    collection: Collection,
-    has_annotation: bool = True,
-) -> list[str]:
-    """List the names of entries that match the filters."""
-    assert isinstance(collection, Collection)
-    assert isinstance(has_annotation, bool)
-    # build query
-    query = {}
-    if has_annotation:
-        query['annotation'] = {'$ne': None}
-    # execute query
-    res_list = collection.find(
-        filter=query,
-        projection={
-            '_id': False,
-            'name': True,
-        },
-    )
-    # extract information
-    name_list = [elem['name'] for elem in res_list]
-    return name_list
-
-
-def total_documents(
-    collection: Collection,
-) -> int:
-    """Get total number of documents in database."""
-    return collection.count_documents(filter={})
-
-
-def total_annotated(
-    collection: Collection,
-) -> int:
-    """Get total number of annotated documents in database."""
-    query = {
-        'annotation': {
-            '$ne': None,
-        },
-    }
-    return collection.count_documents(filter=query)
-
-
-def get_visual_communication(
-    collection: Collection,
-    with_annotation: bool = False,
-) -> VisualCommunication:
-    """Get a random visual communication from the database."""
-    query = {}
-    if with_annotation:
-        query['annotation'] = {'$ne': None}
-    else:
-        query['annotation'] = {'$eq': None}
-    data = collection.aggregate([
-        {
-            '$match': query,     # find using filters
-        },
-        {
-            '$sample': {
-                'size': 1,   # get one random
-            },
-        },
-    ])
-    data_list = list(data)   # read data from cursor object
-    if len(data_list) == 0:
-        logging.error('failed getting visual communication')
-        raise NoDocumentFoundException()
-    logging.info('finished')
-    return VisualCommunication.model_validate(data_list[0])
-
-
-def upsert_predictions(
-    collection: Collection,
-    vis_com_name: str,
-    predictions: ModelData,
-) -> None:
-    """Upsert prediction data in the database."""
-    query = {
-        'name': vis_com_name,
-    }
-    update = {
-        '$set': {
-            'prediction': predictions.model_dump(),
-        },
-    }
-    res = collection.update_one(
-        filter=query,
-        update=update,
-        upsert=True,
-    )
-    logging.debug('upserted document: %s', res)
-    logging.info('finished')
-
-
-def upsert_annotations(
-    collection: Collection,
-    vis_com_name: str,
-    annotations: ModelData,
-) -> None:
-    """Upserts annotation data in the database."""
-    query = {
-        'name': vis_com_name,
-    }
-    update = {
-        '$set': {
-            'annotation': annotations.model_dump(),
-        },
-    }
-    res = collection.update_one(
-        filter=query,
-        update=update,
-        upsert=True,
-    )
-    logging.info('upserted document: %s', res)
-    logging.info('finished')
-
-
-def upsert_visual_communication(
-    collection: Collection,
-    visual_communication_list: list[VisualCommunication],
-) -> bool:
-    """
-    Upsert VisualCommunication object in the database.
-    Returns bool stating success.
-    """
-    response = collection.insert_many(
-        [
-            vis_com.model_dump()
-            for vis_com
-            in visual_communication_list
-        ],
-    )
-    return response.acknowledged

core/database/utils/__init__.py

@@ -0,0 +1,8 @@
+from __future__ import annotations
+
+from .count_documents import count_documents
+from .get_visual_communication import get_visual_communication
+from .list_names import list_names
+from .upsert_annotation import upsert_annotations
+from .upsert_prediction import upsert_predictions
+from .upsert_visual_communication import upsert_visual_communication

core/database/utils/count_documents.py

@@ -0,0 +1,20 @@
+from __future__ import annotations
+
+from pymongo.collection import Collection
+
+
+def count_documents(
+    collection: Collection,
+    has_annotation: bool = False,
+) -> int:
+    """
+    Get the total number of documents
+    in database that matches the filters.
+    """
+    assert isinstance(collection, Collection)
+    assert isinstance(has_annotation, bool)
+    # build query
+    query = {}
+    if has_annotation:
+        query['annotation'] = {'$ne': None}
+    return collection.count_documents(filter=query)

core/database/utils/get_visual_communication.py

@@ -0,0 +1,36 @@
+from __future__ import annotations
+
+import logging
+
+from pymongo.collection import Collection
+
+from core.database import NoDocumentFoundException
+from core.database import VisualCommunication
+
+
+def get_visual_communication(
+    collection: Collection,
+    with_annotation: bool = False,
+) -> VisualCommunication:
+    """Get a random visual communication from the database."""
+    query = {}
+    if with_annotation:
+        query['annotation'] = {'$ne': None}
+    else:
+        query['annotation'] = {'$eq': None}
+    data = collection.aggregate([
+        {
+            '$match': query,     # find using filters
+        },
+        {
+            '$sample': {
+                'size': 1,   # get one random
+            },
+        },
+    ])
+    data_list = list(data)   # read data from cursor object
+    if len(data_list) == 0:
+        logging.error('failed getting visual communication')
+        raise NoDocumentFoundException()
+    logging.info('finished')
+    return VisualCommunication.model_validate(data_list[0])

core/database/utils/list_names.py

@@ -0,0 +1,27 @@
+from __future__ import annotations
+
+from pymongo.collection import Collection
+
+
+def list_names(
+    collection: Collection,
+    has_annotation: bool = True,
+) -> list[str]:
+    """List the names of entries that match the filters."""
+    assert isinstance(collection, Collection)
+    assert isinstance(has_annotation, bool)
+    # build query
+    query = {}
+    if has_annotation:
+        query['annotation'] = {'$ne': None}
+    # execute query
+    res_list = collection.find(
+        filter=query,
+        projection={
+            '_id': False,
+            'name': True,
+        },
+    )
+    # extract information
+    name_list = [elem['name'] for elem in res_list]
+    return name_list

core/database/utils/upsert_annotation.py

@@ -0,0 +1,30 @@
+from __future__ import annotations
+
+import logging
+
+from pymongo.collection import Collection
+
+from core.dto import ModelData
+
+
+def upsert_annotations(
+    collection: Collection,
+    vis_com_name: str,
+    annotations: ModelData,
+) -> None:
+    """Upserts annotation data in the database."""
+    query = {
+        'name': vis_com_name,
+    }
+    update = {
+        '$set': {
+            'annotation': annotations.model_dump(),
+        },
+    }
+    res = collection.update_one(
+        filter=query,
+        update=update,
+        upsert=True,
+    )
+    logging.info('upserted document: %s', res)
+    logging.info('finished')

core/database/utils/upsert_prediction.py

@@ -0,0 +1,30 @@
+from __future__ import annotations
+
+import logging
+
+from pymongo.collection import Collection
+
+from core.dto import ModelData
+
+
+def upsert_predictions(
+    collection: Collection,
+    vis_com_name: str,
+    predictions: ModelData,
+) -> None:
+    """Upsert prediction data in the database."""
+    query = {
+        'name': vis_com_name,
+    }
+    update = {
+        '$set': {
+            'prediction': predictions.model_dump(),
+        },
+    }
+    res = collection.update_one(
+        filter=query,
+        update=update,
+        upsert=True,
+    )
+    logging.debug('upserted document: %s', res)
+    logging.info('finished')

core/database/utils/upsert_visual_communication.py

@@ -0,0 +1,23 @@
+from __future__ import annotations
+
+from pymongo.collection import Collection
+
+from core.database import VisualCommunication
+
+
+def upsert_visual_communication(
+    collection: Collection,
+    visual_communication_list: list[VisualCommunication],
+) -> bool:
+    """
+    Upsert VisualCommunication object in the database.
+    Returns bool stating success.
+    """
+    response = collection.insert_many(
+        [
+            vis_com.model_dump()
+            for vis_com
+            in visual_communication_list
+        ],
+    )
+    return response.acknowledged