PyCS/pycs/database/Result.py

import typing as T

from contextlib import closing
from json import dumps, loads

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

class Result(BaseModel):

    file_id = db.Column(
        db.Integer,
        db.ForeignKey("file.id", ondelete="CASCADE"),
        nullable=False)

    origin = db.Column(db.String, nullable=False)
    type = db.Column(db.String, nullable=False)

    label_id = db.Column(
        db.Integer,
        db.ForeignKey("label.id", ondelete="SET NULL"),
        nullable=True)

    data_encoded = db.Column(db.String)

    serialize_only = BaseModel.serialize_only + (
        "file_id",
        "origin",
        "type",
        "label_id",
        "data",
    )

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

    @property
    def data(self):
        return None if self.data_encoded is None else json.loads(self.data_encoded)

    @data.setter
    def data(self, value):
        if isinstance(value, str) or value is None:
            self.data_encoded = value

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

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

    @commit_on_return
    def set_origin(self, origin: str):
        """
        set this results origin

        :param origin: either 'user' or 'pipeline'
        :return:
        """
        self.origin = origin


    @commit_on_return
    def set_label(self, label: int):
        """
        set this results origin

        :param label: label ID
        :return:
        """
        self.label_id = label