PyCS/pycs/database/File.py

from contextlib import closing
from json import dumps
from typing import List, Optional

from pycs.database.Result import Result


class File:
    """
    database class for files
    """

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

        self.identifier = row[0]
        self.uuid = row[1]
        self.project_id = row[2]
        self.collection_id = row[3]
        self.type = row[4]
        self.name = row[5]
        self.extension = row[6]
        self.size = row[7]
        self.created = row[8]
        self.path = row[9]
        self.frames = row[10]
        self.fps = row[11]

    def project(self):
        """
        get the project associated with this file

        :return: project
        """
        return self.database.project(self.project_id)

    def collection(self):
        """
        get the collection associated with this file

        :return: collection
        """
        if self.collection_id is None:
            return None

        return self.database.collection(self.collection_id)

    def set_collection(self, collection_id: Optional[int]):
        """
        set this file's collection

        :param collection_id: new collection
        :return:
        """
        with closing(self.database.con.cursor()) as cursor:
            cursor.execute('UPDATE files SET collection = ? WHERE id = ?',
                           (collection_id, self.identifier))
            self.collection_id = collection_id

    def set_collection_by_reference(self, collection_reference: Optional[str]):
        """
        set this file's collection

        :param collection_reference: collection reference
        :return:
        """
        if collection_reference is None:
            self.set_collection(None)
            return

        with closing(self.database.con.cursor()) as cursor:
            cursor.execute('SELECT id FROM collections WHERE reference = ?', [collection_reference])
            row = cursor.fetchone()

        self.set_collection(row[0] if row is not None else None)

    def remove(self) -> None:
        """
        remove this file from the database

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

    def previous(self):
        """
        get the predecessor of this file

        :return: another file
        """
        with closing(self.database.con.cursor()) as cursor:
            cursor.execute('''
                SELECT * FROM files WHERE id < ? AND project = ? ORDER BY id DESC LIMIT 1
            ''', (self.identifier, self.project_id))
            row = cursor.fetchone()

            if row is not None:
                return File(self.database, row)

            return None

    def next(self):
        """
        get the successor of this file

        :return: another file
        """
        with closing(self.database.con.cursor()) as cursor:
            cursor.execute('''
                    SELECT * FROM files WHERE id > ? AND project = ? ORDER BY id ASC LIMIT 1
                ''', (self.identifier, self.project_id))
            row = cursor.fetchone()

            if row is not None:
                return File(self.database, row)

            return None

    def previous_in_collection(self):
        """
        get the predecessor of this file

        :return: another file
        """
        with closing(self.database.con.cursor()) as cursor:
            if self.collection_id is None:
                cursor.execute('''
                    SELECT * FROM files
                    WHERE id < ? AND project = ? AND collection IS NULL
                    ORDER BY id DESC
                    LIMIT 1
                ''', (self.identifier, self.project_id))
            else:
                cursor.execute('''
                    SELECT * FROM files
                    WHERE id < ? AND project = ? AND collection = ?
                    ORDER BY id DESC
                    LIMIT 1
                ''', (self.identifier, self.project_id, self.collection_id))

            row = cursor.fetchone()
            if row is not None:
                return File(self.database, row)

            return None

    def next_in_collection(self):
        """
        get the successor of this file

        :return: another file
        """
        with closing(self.database.con.cursor()) as cursor:
            if self.collection_id is None:
                cursor.execute('''
                    SELECT * FROM files
                    WHERE id > ? AND project = ? AND collection IS NULL
                    ORDER BY id ASC
                    LIMIT 1
                ''', (self.identifier, self.project_id))
            else:
                cursor.execute(''' 
                    SELECT * FROM files
                    WHERE id > ? AND project = ? AND collection = ?
                    ORDER BY id ASC
                    LIMIT 1
                ''', (self.identifier, self.project_id, self.collection_id))

            row = cursor.fetchone()
            if row is not None:
                return File(self.database, row)

            return None

    def results(self) -> List[Result]:
        """
        get a list of all results associated with this file

        :return: list of results
        """
        with closing(self.database.con.cursor()) as cursor:
            cursor.execute('SELECT * FROM results WHERE file = ?', [self.identifier])
            return list(map(
                lambda row: Result(self.database, row),
                cursor.fetchall()
            ))

    def result(self, identifier) -> Optional[Result]:
        """
        get a specific result using its unique identifier

        :param identifier: unique identifier
        :return: result
        """
        with closing(self.database.con.cursor()) as cursor:
            cursor.execute('''
                SELECT * FROM results WHERE id = ? AND file = ?
            ''', (identifier, self.identifier))
            row = cursor.fetchone()

            if row is not None:
                return Result(self.database, row)

            return None

    def create_result(self, origin, result_type, label, data=None):
        """
        create a result

        :param origin:
        :param result_type:
        :param label:
        :param data:
        :return:
        """
        if data is not None:
            data = dumps(data)

        with closing(self.database.con.cursor()) as cursor:
            cursor.execute('''
                INSERT INTO results (file, origin, type, label, data)
                VALUES              (   ?,      ?,    ?,     ?,    ?)
            ''', (self.identifier, origin, result_type, label, data))

            return self.result(cursor.lastrowid)

    def remove_results(self, origin='pipeline') -> List[Result]:
        with closing(self.database.con.cursor()) as cursor:
            cursor.execute('''
                SELECT * FROM results WHERE file = ? AND origin = ?
            ''', (self.identifier, origin))

            results = list(map(lambda row: Result(self.database, row), cursor.fetchall()))

            cursor.execute('''
                DELETE FROM results WHERE file = ? AND origin = ?
            ''', (self.identifier, origin))

            return results