PyCS/pycs/database/Model.py

from contextlib import closing
from json import loads, dumps


class Model:
    """
    database class for label providers
    """

    def __init__(self, database, row):
        self.database = database

        self.identifier = row[0]
        self.name = row[1]
        self.description = row[2]
        self.root_folder = row[3]
        self.supports = loads(row[4])

    def copy_to(self, name: str, root_folder: str):
        """
        copies the models database entry while changing name and root_folder

        :param name: copy name
        :param root_folder: copy root folder
        :return: copy
        """
        supports = dumps(self.supports)

        with closing(self.database.con.cursor()) as cursor:
            cursor.execute('''
                INSERT INTO models (name, description, root_folder, supports)
                VALUES (?, ?, ?, ?)
                ON CONFLICT (root_folder)
                DO UPDATE SET name = ?, description = ?, supports = ?
            ''', (name, self.description, root_folder, supports, name, self.description, supports))

            # lastrowid is 0 if on conflict clause applies.
            # If this is the case we do an extra query to receive the row id.
            if cursor.lastrowid > 0:
                row_id = cursor.lastrowid
                insert = True
            else:
                cursor.execute('SELECT id FROM models WHERE root_folder = ?', [root_folder])
                row_id = cursor.fetchone()[0]
                insert = False

        return self.database.model(row_id), insert

    def remove(self):
        """
        remove this model from the database

        :return:
        """
        with closing(self.database.con.cursor()) as cursor:
            cursor.execute('DELETE FROM models WHERE id = ?', [self.identifier])