import typing
import PySide2
from PySide2.QtCore import Signal, QModelIndex, Qt, QItemSelection, QItemSelectionModel
from PySide2.QtGui import QColor
from PySide2.QtWidgets import QTreeView, QToolButton
from arthropod_describer.common.label_hierarchy import LabelHierarchy, Node
from arthropod_describer.common.label_tree_model import LabelTreeModel
from arthropod_describer.common.state import State
[docs]class LabelTreeView(QTreeView):
constraint_requested = Signal(int)
unset_constraint = Signal()
label_clicked = Signal(int)
label_color_change = Signal([int, QColor])
label_dbl_click = Signal(Node)
def __init__(self, state: State, parent: typing.Optional[PySide2.QtWidgets.QWidget] = None):
super().__init__(parent)
self._btn = QToolButton()
self._curr_index: QModelIndex = QModelIndex()
self._curr_label: int = -1
self._constraint_index: typing.Optional[QModelIndex] = QModelIndex()
self._constraint_label: int = -1
self._btn_constraint: typing.Optional[QToolButton] = None
self.state: State = state
self.state.label_hierarchy_changed.connect(self._handle_label_hierarchy_changed)
self.state.primary_label_changed.connect(self.choose_label)
# self.clicked.connect(self._handle_index_activated)
self.clicked.connect(self._index_click_handler)
self.setExpandsOnDoubleClick(False)
self.clicked_idx: QModelIndex = QModelIndex()
self.setStyleSheet("QTreeView::item { padding: 1px }")
self.setMouseTracking(True)
self.entered.connect(self._set_constraint_button)
self._can_be_self_constrained: bool = True
self.setContextMenuPolicy(Qt.CustomContextMenu)
self._constraint_mode: bool = False
self._label_tree_model: typing.Optional[LabelTreeModel] = None
def _handle_label_hierarchy_changed(self, lab_hier: LabelHierarchy):
self.reset()
self.resizeColumnToContents(0)
def _set_constraint_button(self, index: QModelIndex):
if not self._can_be_self_constrained:
return
if index.internalPointer().label == 0 or (index.internalPointer().label == self.state.constraint_label):
return
dst_index = index if index.column() == 1 else index.siblingAtColumn(1)
if self._curr_index.isValid():
self.setIndexWidget(self._curr_index, None)
self._curr_index = QModelIndex()
btn = QToolButton(text="Set constraint")
btn.clicked.connect(self.accept_constraint)
self.setIndexWidget(dst_index, btn)
self._curr_index = dst_index
def _handle_index_activated(self, idx: QModelIndex):
if idx.column() > 0:
return
if self.clicked_idx == idx:
return
# self.label_clicked.emit(idx.internalPointer().label)
if self.clicked_idx.isValid():
self.model().dataChanged.emit(self.clicked_idx, self.clicked_idx, Qt.BackgroundRole)
self.model().dataChanged.emit(idx, idx, Qt.BackgroundRole)
self.clicked_idx = idx
# TODO update the label tree view to show the selected label
#self.reset()
#self.expandAll()
[docs] def choose_label(self, label: typing.Union[int, QModelIndex]):
if isinstance(label, QModelIndex):
index = label
else:
index = self._label_tree_model.find_index(label)
self.selectionModel().select(index, QItemSelectionModel.SelectCurrent)
def _index_click_handler(self, index: QModelIndex):
label = index.internalPointer().label
self.state.primary_label = label
# def setModel(self, model: PySide2.QtCore.QAbstractItemModel):
[docs] def setModel(self, model: LabelTreeModel):
self._label_tree_model = model
self.setMouseTracking(True)
self.clicked_idx = QModelIndex()
self._can_be_self_constrained = False
if self.state.label_can_be_constrained:
if (possible_constraints := self.state.storage.label_img_info[self.state.current_label_name].can_constrain_to) is not None:
self._can_be_self_constrained = self.state.current_label_name in possible_constraints
super().setModel(model)
self._constraint_mode = self.model().columnCount() == 2
self._can_be_self_constrained = self._can_be_self_constrained and self._constraint_mode
self.expandAll()
self.resizeColumnToContents(0)
[docs] def accept_constraint(self):
if self._constraint_index is not None and self._constraint_index.isValid():
self.remove_unset_constraint_btn()
btn = QToolButton()
btn.setText("Remove constraint")
btn.clicked.connect(self._unset_constraint)
self.setIndexWidget(self._curr_index, btn)
self._btn_constraint = btn
self._constraint_index = self._curr_index
self._curr_index = QModelIndex()
self.state.current_constraint.label_name = self.state.current_label_name
self.state.constraint_label = self._constraint_index.internalPointer().label
[docs] def leaveEvent(self, event:PySide2.QtCore.QEvent):
if self._curr_index.internalPointer() is None:
self.setIndexWidget(self._curr_index, None)
return
if self._curr_index.isValid():
if self._constraint_index.internalPointer() is not None:
if self._constraint_index.internalPointer().label != self._curr_index.internalPointer().label:
self.setIndexWidget(self._curr_index, None)
self._curr_index = QModelIndex()
else:
self.setIndexWidget(self._curr_index, None)
self._curr_index = QModelIndex()
#self._delegate.destroyEditor(self._delegate.editor, self._curr_index)
self._curr_label = -1
def _unset_constraint(self, set_constraint_button: bool = True):
index = self._constraint_index
self.remove_unset_constraint_btn()
self.state.constraint_label = 0
self.state.current_constraint.label_name = None
self.state.current_constraint.label_node = None
if set_constraint_button:
self._set_constraint_button(index)
[docs] def remove_unset_constraint_btn(self):
self.setIndexWidget(self._constraint_index, None)
self._constraint_index = QModelIndex()
self._btn_constraint = None
# self.reset()
# self.expandAll()
[docs] def set_constraint(self, index: QModelIndex):
self._curr_index = index
self.accept_constraint()
[docs] def mouseDoubleClickEvent(self, event:PySide2.QtGui.QMouseEvent):
index = self.indexAt(event.pos())
node: Node = index.internalPointer()
self.label_dbl_click.emit(node)