PyCS/pycs/database/Database.py

import sqlite3
from contextlib import closing
from time import time
from typing import Optional, Iterator

from pycs import app
from pycs import db
from pycs.database.Collection import Collection
from pycs.database.File import File
from pycs.database.LabelProvider import LabelProvider
from pycs.database.Model import Model
from pycs.database.Project import Project
from pycs.database.Result import Result


class Database:
    """
    opens an sqlite database and allows to access several objects
    """

    def __init__(self, path: str = ':memory:', initialization=True, discovery=True):
        """
        opens or creates a given sqlite database and creates all required tables

        :param path: path to sqlite database
        """
        # save properties
        self.path = path

        if discovery:
            # run discovery modules
            Model.discover("models/")
            LabelProvider.discover("labels/")

    def __enter__(self):
        app.logger.warning("REMOVE ME!")
        return self

    def __exit__(self, exc_type, exc_val, exc_tb):
        app.logger.warning("REMOVE ME!")

        if exc_type is None:
            db.session.commit()
        else:
            app.logger.info("Rolling back a transaction!")
            db.session.rollback()

    def copy(self):
        return Database(self.path, initialization=False, discovery=False)

    def models(self) -> Iterator[Model]:
        """
        get a list of all available models

        :return: iterator of models
        """
        return Model.query.all()

    def model(self, identifier: int) -> Optional[Model]:
        """
        get a model using its unique identifier

        :param identifier: unique identifier
        :return: model
        """
        return Model.query.get(identifier)

    def label_providers(self) -> Iterator[LabelProvider]:
        """
        get a list of all available label providers

        :return: iterator over label providers
        """
        return LabelProvider.query.all()

    def label_provider(self, identifier: int) -> Optional[LabelProvider]:
        """
        get a label provider using its unique identifier

        :param identifier: unique identifier
        :return: label provider
        """
        return LabelProvider.query.get(identifier)

    def projects(self) -> Iterator[Project]:
        """
        get a list of all available projects

        :return: iterator over projects
        """
        return Project.query.all()

    def project(self, identifier: int) -> Optional[Project]:
        """
        get a project using its unique identifier

        :param identifier: unique identifier
        :return: project
        """
        return Project.query.get(identifier)

    def collection(self, identifier: int) -> Optional[Collection]:
        """
        get a collection using its unique identifier

        :param identifier: unique identifier
        :return: collection
        """
        return Collection.query.get(identifier)

    def file(self, identifier) -> Optional[File]:
        """
        get a file using its unique identifier

        :param identifier: unique identifier
        :return: file
        """
        return File.query.get(identifier)

    def result(self, identifier) -> Optional[Result]:
        """
        get a result using its unique identifier

        :param identifier: unique identifier
        :return: result
        """
        return Result.query.get(identifier)

    def create_project(self,
                       name: str,
                       description: str,
                       model: Model,
                       label_provider: Optional[LabelProvider],
                       root_folder: str,
                       external_data: bool,
                       data_folder: str):
        """
        insert a project into the database

        :param name: project name
        :param description: project description
        :param model: used model
        :param label_provider: used label provider (optional)
        :param root_folder: path to project folder
        :param external_data: whether an external data directory is used
        :param data_folder: path to data folder
        :return: created project
        """
        # prepare some values

        return Project.new(commit=True,
            name=name,
            description=description,
            model=model,
            label_provider=label_provider,
            root_folder=root_folder,
            external_data=external_data,
            data_folder=data_folder
        )