diff --git a/core/__init__.py b/core/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/core/database/__init__.py b/core/database/__init__.py index ef9599c..a1321c1 100644 --- a/core/database/__init__.py +++ b/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 diff --git a/core/database/utils/get_dataset.py b/core/database/utils/get_dataset.py new file mode 100755 index 0000000..5440df1 --- /dev/null +++ b/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 [] diff --git a/core/database/utils/save_dataset.py b/core/database/utils/save_dataset.py new file mode 100755 index 0000000..afb2a64 --- /dev/null +++ b/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 + # )