fixed imports

core/database/__init__.py

@@ -1,10 +1,10 @@
 """Database module content."""
 from __future__ import annotations
 
-from .classes import Dataset
-from .classes import NoDocumentFoundException
-from .classes import VisualCommunication
-from .utils import connect
+from .classes.dataset import Dataset
+from .classes.exceptions import NoDocumentFoundException
+from .classes.visual_communication import VisualCommunication
+from .utils.connect import connect
 from .utils.count_documents import count_documents
 from .utils.get_visual_communication import get_visual_communication
 from .utils.list_names import list_names

core/database/utils/get_dataset.py

@@ -0,0 +1,13 @@
+from __future__ import annotations
+
+from typing import Literal
+
+from pymongo.collection import Collection
+
+
+def get_dataset(
+    collection: Collection,
+    type: Literal['train', 'test', 'validation'],
+) -> list[str]:
+    """Get list of data names for the corresponding type."""
+    return []

core/database/utils/save_dataset.py

@@ -0,0 +1,36 @@
+from __future__ import annotations
+
+import logging
+
+from pymongo.collection import Collection
+
+from core.database import Dataset
+
+
+def save_dataset(
+    collection: Collection,
+    dataset: Dataset,
+) -> None:
+    """Save dataset to database."""
+    res = collection.insert_one(
+        document=dataset.model_dump(),
+    )
+    logging.debug('inserted document: %s', res)
+
+
+if __name__ == '__main__':
+    from dotenv import load_dotenv
+    load_dotenv('local.env')
+    from core.database import list_names, connect
+    # connect to database
+    collection, db, client = connect()
+    print(client.server_info())
+
+    name_list = list_names(collection=collection, only_with_annotation=True)
+    ds = Dataset.new_from_name_list(name_list=name_list)
+
+    print(ds)
+    # save_dataset(
+    #     collection=collection,
+    #     dataset=ds
+    # )