PyCS/pycs/database/Result.py

import typing as T
import json

from contextlib import closing

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

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 = db.Column(db.String)

    serialize_only = (
        "id",
        "file_id",
        "origin",
        "type",
        "label_id",
        "data",
    )


    def serialize(self):
        result = super().serialize()
        if result["data"] is not None:
            result["data"] = json.loads(result["data"])
        return result

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

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


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

        :param label: label ID
        :return:
        """
        self.label_id = label
        self.commit()

    def set_data(self, data: T.Optional[dict]):
        """
        set this results data object

        :param data: data object
        :return:
        """
        data = data if data is None else json.dumps(data)

        self.data = data
        self.commit()