added database util functions

core/database/__init__.py

@@ -5,6 +5,7 @@ from .classes import NoDocumentFoundException
 from .classes import VisualCommunication
 from .database import connect
 from .utils import get_visual_communication
+from .utils import list_names
 from .utils import total_annotated
 from .utils import total_documents
 from .utils import upsert_annotations

core/database/utils.py

@@ -9,6 +9,44 @@ from .classes import NoDocumentFoundException
 from .classes import VisualCommunication
 
 
+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}
+    res_list = collection.find(
+        filter=query,
+        projection={
+            '_id': False,
+            'name': True,
+        },
+    )
+    return list(res_list)
+
+
 def total_documents(
     collection: Collection,
 ) -> int: