mirror of
https://github.com/ynput/ayon-core.git
synced 2025-12-24 21:04:40 +01:00
94 lines
2.9 KiB
Python
94 lines
2.9 KiB
Python
import contextlib
|
|
from Qt import QtCore
|
|
|
|
|
|
def _iter_model_rows(
|
|
model, column, include_root=False
|
|
):
|
|
"""Iterate over all row indices in a model"""
|
|
indices = [QtCore.QModelIndex()] # start iteration at root
|
|
|
|
for index in indices:
|
|
# Add children to the iterations
|
|
child_rows = model.rowCount(index)
|
|
for child_row in range(child_rows):
|
|
child_index = model.index(child_row, column, index)
|
|
indices.append(child_index)
|
|
|
|
if not include_root and not index.isValid():
|
|
continue
|
|
|
|
yield index
|
|
|
|
|
|
@contextlib.contextmanager
|
|
def preserve_states(
|
|
tree_view, column=0, role=None,
|
|
preserve_expanded=True, preserve_selection=True,
|
|
expanded_role=QtCore.Qt.DisplayRole, selection_role=QtCore.Qt.DisplayRole
|
|
|
|
):
|
|
"""Preserves row selection in QTreeView by column's data role.
|
|
|
|
This function is created to maintain the selection status of
|
|
the model items. When refresh is triggered the items which are expanded
|
|
will stay expanded and vise versa.
|
|
|
|
tree_view (QWidgets.QTreeView): the tree view nested in the application
|
|
column (int): the column to retrieve the data from
|
|
role (int): the role which dictates what will be returned
|
|
|
|
Returns:
|
|
None
|
|
|
|
"""
|
|
# When `role` is set then override both expanded and selection roles
|
|
if role:
|
|
expanded_role = role
|
|
selection_role = role
|
|
|
|
model = tree_view.model()
|
|
selection_model = tree_view.selectionModel()
|
|
flags = selection_model.Select | selection_model.Rows
|
|
|
|
expanded = set()
|
|
|
|
if preserve_expanded:
|
|
for index in _iter_model_rows(
|
|
model, column=column, include_root=False
|
|
):
|
|
if tree_view.isExpanded(index):
|
|
value = index.data(expanded_role)
|
|
expanded.add(value)
|
|
|
|
selected = None
|
|
|
|
if preserve_selection:
|
|
selected_rows = selection_model.selectedRows()
|
|
if selected_rows:
|
|
selected = set(row.data(selection_role) for row in selected_rows)
|
|
|
|
try:
|
|
yield
|
|
finally:
|
|
if expanded:
|
|
for index in _iter_model_rows(
|
|
model, column=0, include_root=False
|
|
):
|
|
value = index.data(expanded_role)
|
|
is_expanded = value in expanded
|
|
# skip if new index was created meanwhile
|
|
if is_expanded is None:
|
|
continue
|
|
tree_view.setExpanded(index, is_expanded)
|
|
|
|
if selected:
|
|
# Go through all indices, select the ones with similar data
|
|
for index in _iter_model_rows(
|
|
model, column=column, include_root=False
|
|
):
|
|
value = index.data(selection_role)
|
|
state = value in selected
|
|
if state:
|
|
tree_view.scrollTo(index) # Ensure item is visible
|
|
selection_model.select(index, flags)
|