PyCS/pycs/database/Model.py

import json

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):
    """
    database class 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 = NamedBaseModel.serialize_only + (
        "description",
        "root_folder",
    )


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

    @classmethod
    def discover(cls, root: Path, config_name: str = "configuration.json"):
        for folder in Path(root).glob("*"):
            with open(folder / config_name) as f:
                config = json.load(f)

            # 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

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

    @property
    def supports(self):
        return json.loads(self.supports_encoded)

    @supports.setter
    def supports(self, value):
        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, new_name: str, new_root_folder: str):

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

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

        return model, is_new