diff --git a/core/database/__init__.py b/core/database/__init__.py index 0990ebc..e964218 100644 --- a/core/database/__init__.py +++ b/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 diff --git a/core/database/utils.py b/core/database/utils.py index d43fa77..581b807 100644 --- a/core/database/utils.py +++ b/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: