from ..Qt import QtWidgets, QtGui, QtCore


class TreeWidget(QtWidgets.QTreeWidget):
    def keyPressEvent(self, evt: QtGui.QKeyEvent):
        if evt.key() == QtCore.Qt.Key_Space:
            sets = []
            from_parent = []

            for idx in self.selectedIndexes():
                if idx.column() != 0:
                    continue
                item = self.itemFromIndex(idx)

                if item.parent() is None:
                    is_selected = item.checkState(0)
                    self.blockSignals(True)
                    for i in range(item.childCount()):
                        child = item.child(i)
                        # child.setCheckState(0, is_selected)
                        from_parent.append(child)
                    self.blockSignals(False)
                    item.setCheckState(0, QtCore.Qt.Unchecked if is_selected == QtCore.Qt.Checked else QtCore.Qt.Checked)
                else:
                    sets.append(item)
            for it in sets:
                if it in from_parent:
                    continue
                it.setCheckState(0, QtCore.Qt.Unchecked if it.checkState(0) == QtCore.Qt.Checked else QtCore.Qt.Checked)
        else:
            super().keyPressEvent(evt)


class TableWidget(QtWidgets.QTableWidget):
    def keyPressEvent(self, evt: QtGui.QKeyEvent):
        if evt.key() == QtCore.Qt.Key.Key_Space:
            for idx in self.selectedIndexes():
                item = self.itemFromIndex(idx)
                cs = item.checkState()
                item.setCheckState(QtCore.Qt.CheckState.Unchecked if cs == QtCore.Qt.CheckState.Checked
                                   else QtCore.Qt.CheckState.Checked)
        else:
            super().keyPressEvent(evt)