Source code for SDF.GUI.browser.sdf_tree

from typing import Optional, List

from PyQt5.QtCore import QItemSelection
from PyQt5.QtWidgets import QTreeWidget, QTreeWidgetItem, QWidget

from SDF.GUI.browser.preview import DataPreviewWidget
from SDF.GUI.browser.summary import SummaryWidget
from SDF.data_model import ArrayDataset1D, ArrayDataset2D, ImageDataset, SDFObject, Workspace


class SDFTreeWidgetItem(QTreeWidgetItem):
    def __init__(self, parent: "SDFTreeWidgetItem", sdf: SDFObject):
        super().__init__(parent)
        self.sdf = sdf
        self.setText(0, sdf.name)
        self.setText(1, self.__create_summary(sdf))

    @staticmethod
    def __create_summary(sdf: SDFObject) -> str:
        """
        Create a summary text of an sdf_object to be displayed
        in the second column of the QTreeView.
        """
        ret = ""
        if isinstance(sdf, ArrayDataset1D):
            ret += f"single-column data of length {sdf.data.shape[0]}, "
        elif isinstance(sdf, ArrayDataset2D):
            ret += f"data array of shape {sdf.data.shape}, "
        elif isinstance(sdf, ImageDataset):
            ret += "image data, "
        elif isinstance(sdf, Workspace):
            ret += f"contains {len(sdf.workspaces)} workspaces, {len(sdf.datasets)} datasets, "
        else:
            raise TypeError(f"Expected SDFObject, got {sdf}")

        n_samples = len(sdf.samples)
        if n_samples:
            ret += f"{n_samples} sample{'s' if n_samples>1 else ''}, "

        n_instruments = len(sdf.instruments)
        if n_instruments:
            ret += f"{n_instruments} instrument{'s' if n_instruments>1 else ''}, "

        n_parameters = len(sdf.parameters)
        if n_parameters:
            ret += f"{n_parameters} parameter{'s' if n_parameters>1 else ''}"
        ret = ret.rstrip()
        ret = ret.rstrip(',')

        return ret


class SDFTreeWidget(QTreeWidget):
    def __init__(self, parent: QWidget, summary_widget: SummaryWidget, data_preview_widget: DataPreviewWidget):
        super().__init__(parent=parent)
        self.data_preview_widget = data_preview_widget
        self.summary_widget = summary_widget
        self.setColumnCount(2)
        self.setHeaderLabels(["Name", "Summary"])

    def add_sdf_object(self, sdf: SDFObject) -> None:
        self.addTopLevelItem(self.__sdf_to_item(sdf, parent=None))
        self.resizeColumnToContents(0)  # resize first column so that its content is fully displayed

    @staticmethod
    def __sdf_to_item(sdf: SDFObject, parent: Optional[SDFTreeWidgetItem]) -> QTreeWidgetItem:
        item = SDFTreeWidgetItem(parent, sdf)
        if isinstance(sdf, Workspace):
            for child in sdf:
                item.addChild(SDFTreeWidget.__sdf_to_item(child, parent=item))
        return item

    def selectionChanged(self, selected: QItemSelection, deselected: QItemSelection) -> None:
        selected_items: List[SDFTreeWidgetItem] = self.selectedItems()
        if selected_items:
            sdf = selected_items[0].sdf
            if sdf is None:
                # sometimes, the left arrow key can select empty elements of the TreeView
                return
            self.summary_widget.display_sdf_object(sdf)
            self.data_preview_widget.display_sdf_object(sdf)