nmreval/src/gui_qt/io/hdfreader.py

from nmreval.io.hdfreader import HdfReader

from ..Qt import QtGui, QtCore, QtWidgets
from .._py.hdftree import Ui_Hdf_Dialog


class QHdfViewer(QtWidgets.QDialog, Ui_Hdf_Dialog):
    data_read = QtCore.pyqtSignal(list)
    file_ext = ['.h5', '.hdf', '.hdf5']

    def __init__(self, fname: str = None, parent=None):
        super().__init__(parent=parent)
        self.setupUi(self)

        self.treewidget = HDFTreeWidget(self)
        self.verticalLayout_3.addWidget(self.treewidget)

        self.variables = VarTable(self)
        self.widget_2.addWidget(self.variables)
        self.widget_2.setText('Variables')

        self._reader = HdfReader()

        if fname is not None:
            self.__call__(fname)

    def __call__(self, fname, *args, **kwargs):
        self.treewidget.clear()
        for cb in [self.comboBox, self.comboBox_2]:
            cb.clear()
            cb.addItem('Automatic for the people')
            cb.show()

        self._reader(filename=fname)
        self._fill_boxes()
        self._populate_tree(self._reader, self.treewidget.invisibleRootItem())

        if self.comboBox.count() == 1:
            self.comboBox.hide()
            self.comboBox_2.hide()

        self.treewidget.expandToDepth(0)
        self.treewidget.resizeColumnToContents(1)

        return self

    def _populate_tree(self, node, item):
        self.treewidget.blockSignals(True)

        # add varied parameter to comboboxes
        for key, value in node.title_parameter[1].items():
            if isinstance(value, list):
                idx = self.comboBox.findText(key)
                if idx == -1:
                    self.comboBox.addItem(key)
                    self.comboBox_2.addItem(key)

        if node.children is not None:
            for child in node.children.values():
                label_item = QtWidgets.QTreeWidgetItem([child.name])
                label_item.setFlags(QtCore.Qt.ItemIsEnabled | QtCore.Qt.ItemIsUserCheckable | QtCore.Qt.ItemIsSelectable)
                label_item.setCheckState(0, QtCore.Qt.Unchecked)
                if child.type == 'signal':
                    label_item.setBackground(0, QtGui.QBrush(QtGui.QColor(31, 119, 180)))
                    label_item.setCheckState(1, QtCore.Qt.Unchecked)
                elif child.type == 'points':
                    label_item.setBackground(0, QtGui.QBrush(QtGui.QColor(255, 127, 14)))
                item.addChild(label_item)
                self._populate_tree(child, label_item)

        self.treewidget.blockSignals(False)

    def _fill_boxes(self):
        for k, v in self._reader.parameter.items():
            if isinstance(v, list):
                idx = self.comboBox.findText(k)
                if idx == -1:
                    self.comboBox.addItem(k)
                    self.comboBox_2.addItem(k)

        self.variables.populate(self._reader.parameter)

    @QtCore.pyqtSlot()
    def get_selected(self):

        if self.comboBox.currentIndex() == 0:
            value = None
        else:
            value = self.comboBox.currentText()

        if self.comboBox_2.currentIndex() == 0:
            group = None
        else:
            group = self.comboBox_2.currentText()

        iterator = QtWidgets.QTreeWidgetItemIterator(self.treewidget, QtWidgets.QTreeWidgetItemIterator.NoChildren)
        selected = []
        while iterator.value():
            item = iterator.value()
            iterator += 1
            if item.checkState(0) == QtCore.Qt.Checked:
                path = item.text(0)
                parent = item.parent()
                while parent is not None:
                    path = parent.text(0) + '/' + path
                    parent = parent.parent()

                if item.checkState(1) == QtCore.Qt.Checked:
                    selected.extend(self._reader.get_selected(path, flag='spectrum', value=value, group=group))
                else:
                    selected.extend(self._reader.get_selected(path, value=value, group=group))

        self.data_read.emit(selected)

    def accept(self):
        self.get_selected()
        self.close()


class VarTable(QtWidgets.QTableWidget):
    def __init__(self, parent=None):
        super().__init__(parent=parent)

        self.setColumnCount(2)
        self.setHorizontalHeaderLabels(['Key', 'Value'])
        self.horizontalHeader().setStretchLastSection(True)
        self.verticalHeader().setVisible(False)
        self.horizontalHeader().setVisible(True)

    def populate(self, vars: dict):
        self.setHorizontalHeaderLabels(['Key', 'Value'])
        self.setRowCount(len(vars))
        for i, (k, v) in enumerate(vars.items()):
            key_item = QtWidgets.QTableWidgetItem(k)
            key_item.setFlags(QtCore.Qt.NoItemFlags)
            key_item.setForeground(QtGui.QBrush(QtGui.QColor(0, 0, 0)))
            if isinstance(v, list):
                val = f'<{len(v)} values>'
                if len(v) < 40:
                    tip = '\n'.join(str(p) for p in v)
                else:
                    tip = '\n'.join(str(p) for p in v[:40])
                    tip += '\n...'
            else:
                val = str(v)
                tip = val
            value_item = QtWidgets.QTableWidgetItem(val)
            value_item.setFlags(QtCore.Qt.NoItemFlags)
            value_item.setForeground(QtGui.QBrush(QtGui.QColor(0, 0, 0)))
            value_item.setToolTip(tip)
            self.setItem(i, 0, key_item)
            self.setItem(i, 1, value_item)


class HDFTreeWidget(QtWidgets.QTreeWidget):
    def __init__(self, parent):
        super().__init__(parent=parent)

        self.setHeaderLabels(['Data', 'Spectrum?'])
        self.header().setSectionResizeMode(0, QtWidgets.QHeaderView.Stretch)
        self.header().setSectionResizeMode(1, QtWidgets.QHeaderView.Fixed)
        self.header().setStretchLastSection(False)
        self.header().setCascadingSectionResizes(True)
        self.setSelectionMode(QtWidgets.QAbstractItemView.ExtendedSelection)
        self.setSelectionBehavior(QtWidgets.QAbstractItemView.SelectItems)

        self.itemChanged.connect(self.change_item)

    def keyPressEvent(self, evt: QtGui.QKeyEvent):
        if evt.key() == QtCore.Qt.Key_Space:
            for idx in self.selectedIndexes():
                item = self.itemFromIndex(idx)
                col = idx.column()
                cs = item.checkState(col)
                item.setCheckState(col, QtCore.Qt.Unchecked if cs == QtCore.Qt.Checked else QtCore.Qt.Checked)
        else:
            super().keyPressEvent(evt)

    @staticmethod
    def change_item(item: QtWidgets.QTreeWidgetItem, column: int):
        state = item.checkState(column)
        for i in range(item.childCount()):
            child = item.child(i)
            child.setCheckState(column, state)