src/database/utils.py

@@ -1,57 +1,58 @@
-from pymongo.collection import Collection
+from __future__ import annotations
+
 import logging
 
-from .classes import (
-    VisualCommunication,
-    NoDocumentFoundException,
-    ModelOutputs
-)
+from pymongo.collection import Collection
+
+from .classes import ModelOutputs
+from .classes import NoDocumentFoundException
+from .classes import VisualCommunication
 
 
 def total_documents(
-    collection: Collection
+    collection: Collection,
 ) -> int:
     """Get total number of documents in database."""
     return collection.count_documents(filter={})
 
 
 def total_annotated(
-    collection: Collection
+    collection: Collection,
 ) -> int:
     """Get total number of annotated documents in database."""
     query = {
-        "annotation": {
-            "$ne": None
-        }
+        'annotation': {
+            '$ne': None,
+        },
     }
     return collection.count_documents(filter=query)
 
 
 def get_visual_communication(
     collection: Collection,
-    with_annotation: bool = False
+    with_annotation: bool = False,
 ) -> VisualCommunication:
     """Get a random visual communication from the database."""
     query = {}
     if with_annotation:
-        query["annotation"] = {"$ne": None}
+        query['annotation'] = {'$ne': None}
     else:
-        query["annotation"] = {"$eq": None}
+        query['annotation'] = {'$eq': None}
     data = collection.aggregate([
         {
-            "$match": query     # find using filters
+            '$match': query,     # find using filters
         },
         {
-            "$sample": {
-                "size": 1   # get one random
-            }
-        }
+            '$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")
+        logging.error('failed getting visual communication')
         raise NoDocumentFoundException()
-    logging.info("finished")
+    logging.info('finished')
     return VisualCommunication.model_validate(data_list[0])
 
 
@@ -62,20 +63,20 @@ def upsert_predictions(
 ) -> None:
     """Upsert prediction data in the database."""
     query = {
-        "name": vis_com_name
+        'name': vis_com_name,
     }
     update = {
-        "$set": {
-            "prediction": predictions.model_dump()
-        }
+        '$set': {
+            'prediction': predictions.model_dump(),
+        },
     }
     res = collection.update_one(
         filter=query,
         update=update,
         upsert=True,
     )
-    logging.debug("upserted document: %s", res)
-    logging.info("finished")
+    logging.debug('upserted document: %s', res)
+    logging.info('finished')
 
 
 def upsert_annotations(
@@ -85,17 +86,35 @@ def upsert_annotations(
 ) -> None:
     """Upserts annotation data in the database."""
     query = {
-        "name": vis_com_name
+        'name': vis_com_name,
     }
     update = {
-        "$set": {
-            "annotation": annotations.model_dump()
-        }
+        '$set': {
+            'annotation': annotations.model_dump(),
+        },
     }
     res = collection.update_one(
         filter=query,
         update=update,
         upsert=True,
     )
-    logging.info("upserted document: %s", res)
-    logging.info("finished")
+    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