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)