PyCS/pycs/database/Model.py

import json
import typing as T

from pathlib import Path

from pycs import db
from pycs.database.base import NamedBaseModel
from pycs.database.util import commit_on_return


class Model(NamedBaseModel):
    """
    DB model for ML models
    """

    description = db.Column(db.String)
    root_folder = db.Column(db.String, nullable=False, unique=True)
    supports_encoded = db.Column(db.String, nullable=False)

    # relationships to other models
    projects = db.relationship("Project", backref="model", lazy="dynamic")

    serialize_only: tuple = NamedBaseModel.serialize_only + (
        "description",
        "root_folder",
    )

    def serialize(self):
        result = super().serialize()
        result["supports"] = self.supports
        return result

    @classmethod
    def discover(cls, root: T.Union[Path, str], config_name: str = "configuration.json"):
        """
            searches for models under the given path
            and stores them in the database
        """
        for folder in Path(root).glob("*"):
            with open(folder / config_name, encoding='utf8') as config_file:
                config = json.load(config_file)

            # extract data
            name = config['name']
            description = config.get('description', None)
            supports = config['supports']

            model, _ = cls.get_or_create(root_folder=str(folder))

            model.name = name
            model.description = description
            model.supports = supports

            model.flush()
        db.session.commit()

    @property
    def supports(self) -> dict:
        """ getter for the 'supports' attribute """
        return json.loads(self.supports_encoded)

    @supports.setter
    def supports(self, value):
        """
            setter for the 'supports' attribute.
            The attribute is encoded property before assigned to the object.
        """

        if isinstance(value, str):
            self.supports_encoded = value

        elif isinstance(value, (dict, list)):
            self.supports_encoded = json.dumps(value)

        else:
            raise ValueError(f"Not supported type: {type(value)}")

    @commit_on_return
    def copy_to(self, name: str, root_folder: str):
        """ copies current model to another folder and updates the name """

        model, is_new = Model.get_or_create(root_folder=root_folder)

        model.name = name
        model.description = self.description
        model.supports_encoded = self.supports_encoded

        return model, is_new