PyCS/pycs/interfaces/MediaFileList.py

from typing import List, Iterator

from pycs.database.Project import Project
from pycs.frontend.notifications.NotificationList import NotificationList
from pycs.interfaces.MediaFile import MediaFile
from pycs.interfaces.MediaLabel import MediaLabel


class MediaFileList:
    """
    helper class to filter and receive MediaFile elements
    """

    def __init__(self, project: Project, notifications: NotificationList):
        self.__project = project
        self.__notifications = notifications
        self.__collection = None
        self.__label = None

    def filter_collection(self, collection_reference: str):
        """
        only include results matching the collection

        :param collection_reference: reference string
        :return: MediaFileList
        """
        self.__collection = collection_reference
        return self

    def filter_label(self, label: MediaLabel):
        """
        only include results matching the label

        :param label: label
        :return: MediaFileList
        """
        self.__label = label
        return self

    def iter(self) -> Iterator[MediaFile]:
        """
        receive an iterator of files

        :return: iterator of files
        """
        if self.__collection is not None:
            source = self.__project.collection_by_reference(self.__collection)
        else:
            source = self.__project

        if self.__label is None:
            for file in source.files.all():
                yield MediaFile(file, self.__notifications)
        else:
            for file in source.files.all():
                for result in file.results.all():
                    if result.label == self.__label:
                        yield MediaFile(file, self.__notifications)
                        break

    def list(self) -> List[MediaFile]:
        """
        receive a list of files

        :return: list of files
        """
        return list(self.iter())