PyCS/pycs/interfaces/MediaStorage.py

from typing import List

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


class MediaStorage:
    """
    helper class for pipelines to interact with database entities
    """

    def __init__(self, project_id: int, notifications: NotificationList = None):
        self.__project_id = project_id
        self.__notifications = notifications

        self.__project = Project.query.get(self.__project_id)
        # this one is not used anywhere
        # self.__collections = self.__project.collections.all()

    def labels(self) -> List[MediaLabel]:
        """
        receive a list of labels

        :return: list of labels
        """
        label_list = self.__project.labels.all()
        label_dict = {la.id: MediaLabel(la) for la in label_list}
        result = []

        for label in label_list:
            medial_label = label_dict[label.id]

            if label.parent_id is not None:
                medial_label.parent = label_dict[label.parent_id]
                medial_label.parent.children.append(medial_label)

            result.append(medial_label)

        return result

    def labels_tree(self) -> List[MediaLabel]:
        """
        receive a tree of labels

        :return: list of root-level labels (parent is None)
        """
        return [label for label in self.labels() if label.parent is None]

    def files(self) -> MediaFileList:
        """
        get a FileList object to filter and receive MediaFile elements

        :return: FileList
        """
        return MediaFileList(self.__project, self.__notifications)