PyCS/pycs/database/Collection.py

from contextlib import closing
from typing import Iterator

from pycs.database.File import File


class Collection:
    """
    database class for collections
    """

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

        self.identifier = row[0]
        self.project_id = row[1]
        self.reference = row[2]
        self.name = row[3]
        self.description = row[4]
        self.position = row[5]
        self.autoselect = row[6] > 0

    def set_name(self, name: str):
        """
        set this collection's name

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

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

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

    def count_files(self) -> int:
        """
        count files associated with this project

        :return: count
        """
        with closing(self.database.con.cursor()) as cursor:
            cursor.execute('SELECT COUNT(*) FROM files WHERE project = ? AND collection = ?',
                           (self.project_id, self.identifier))
            return cursor.fetchone()[0]

    def files(self, offset: int = 0, limit: int = -1) -> Iterator[File]:
        """
        get an iterator of files associated with this collection

        :param offset: file offset
        :param limit: file limit
        :return: iterator of files
        """
        with closing(self.database.con.cursor()) as cursor:
            cursor.execute('''
                SELECT * FROM files
                WHERE project = ? AND collection = ?
                ORDER BY id ASC LIMIT ? OFFSET ?
                ''', (self.project_id, self.identifier, limit, offset))

            for row in cursor:
                yield File(self.database, row)