mirror of
https://github.com/ynput/ayon-core.git
synced 2026-01-01 16:34:53 +01:00
448 lines
14 KiB
Python
448 lines
14 KiB
Python
import logging
|
|
import contextlib
|
|
import collections
|
|
|
|
from avalon.vendor import qtawesome
|
|
from Qt import QtWidgets, QtCore, QtGui
|
|
from avalon import style, io
|
|
|
|
from .model import (
|
|
TreeModel,
|
|
Item,
|
|
RecursiveSortFilterProxyModel,
|
|
DeselectableTreeView
|
|
)
|
|
|
|
log = logging.getLogger(__name__)
|
|
|
|
|
|
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_expanded_rows(tree_view,
|
|
column=0,
|
|
role=QtCore.Qt.DisplayRole):
|
|
"""Preserves expanded row in QTreeView by column's data role.
|
|
|
|
This function is created to maintain the expand vs collapse status of
|
|
the model items. When refresh is triggered the items which are expanded
|
|
will stay expanded and vice versa.
|
|
|
|
Arguments:
|
|
tree_view (QWidgets.QTreeView): the tree view which is
|
|
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
|
|
|
|
"""
|
|
|
|
model = tree_view.model()
|
|
|
|
expanded = set()
|
|
|
|
for index in _iter_model_rows(model,
|
|
column=column,
|
|
include_root=False):
|
|
if tree_view.isExpanded(index):
|
|
value = index.data(role)
|
|
expanded.add(value)
|
|
|
|
try:
|
|
yield
|
|
finally:
|
|
if not expanded:
|
|
return
|
|
|
|
for index in _iter_model_rows(model,
|
|
column=column,
|
|
include_root=False):
|
|
value = index.data(role)
|
|
state = value in expanded
|
|
if state:
|
|
tree_view.expand(index)
|
|
else:
|
|
tree_view.collapse(index)
|
|
|
|
|
|
@contextlib.contextmanager
|
|
def preserve_selection(tree_view,
|
|
column=0,
|
|
role=QtCore.Qt.DisplayRole,
|
|
current_index=True):
|
|
"""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 vice 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
|
|
|
|
"""
|
|
|
|
model = tree_view.model()
|
|
selection_model = tree_view.selectionModel()
|
|
flags = selection_model.Select | selection_model.Rows
|
|
|
|
if current_index:
|
|
current_index_value = tree_view.currentIndex().data(role)
|
|
else:
|
|
current_index_value = None
|
|
|
|
selected_rows = selection_model.selectedRows()
|
|
if not selected_rows:
|
|
yield
|
|
return
|
|
|
|
selected = set(row.data(role) for row in selected_rows)
|
|
try:
|
|
yield
|
|
finally:
|
|
if not selected:
|
|
return
|
|
|
|
# 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(role)
|
|
state = value in selected
|
|
if state:
|
|
tree_view.scrollTo(index) # Ensure item is visible
|
|
selection_model.select(index, flags)
|
|
|
|
if current_index_value and value == current_index_value:
|
|
tree_view.setCurrentIndex(index)
|
|
|
|
|
|
class AssetModel(TreeModel):
|
|
"""A model listing assets in the silo in the active project.
|
|
|
|
The assets are displayed in a treeview, they are visually parented by
|
|
a `visualParent` field in the database containing an `_id` to a parent
|
|
asset.
|
|
|
|
"""
|
|
|
|
Columns = ["label"]
|
|
Name = 0
|
|
Deprecated = 2
|
|
ObjectId = 3
|
|
|
|
DocumentRole = QtCore.Qt.UserRole + 2
|
|
ObjectIdRole = QtCore.Qt.UserRole + 3
|
|
|
|
def __init__(self, parent=None):
|
|
super(AssetModel, self).__init__(parent=parent)
|
|
self.refresh()
|
|
|
|
def _add_hierarchy(self, assets, parent=None, silos=None):
|
|
"""Add the assets that are related to the parent as children items.
|
|
|
|
This method does *not* query the database. These instead are queried
|
|
in a single batch upfront as an optimization to reduce database
|
|
queries. Resulting in up to 10x speed increase.
|
|
|
|
Args:
|
|
assets (dict): All assets in the currently active silo stored
|
|
by key/value
|
|
|
|
Returns:
|
|
None
|
|
|
|
"""
|
|
if silos:
|
|
# WARNING: Silo item "_id" is set to silo value
|
|
# mainly because GUI issue with preserve selection and expanded row
|
|
# and because of easier hierarchy parenting (in "assets")
|
|
for silo in silos:
|
|
item = Item({
|
|
"_id": silo,
|
|
"name": silo,
|
|
"label": silo,
|
|
"type": "silo"
|
|
})
|
|
self.add_child(item, parent=parent)
|
|
self._add_hierarchy(assets, parent=item)
|
|
|
|
parent_id = parent["_id"] if parent else None
|
|
current_assets = assets.get(parent_id, list())
|
|
|
|
for asset in current_assets:
|
|
# get label from data, otherwise use name
|
|
data = asset.get("data", {})
|
|
label = data.get("label", asset["name"])
|
|
tags = data.get("tags", [])
|
|
|
|
# store for the asset for optimization
|
|
deprecated = "deprecated" in tags
|
|
|
|
item = Item({
|
|
"_id": asset["_id"],
|
|
"name": asset["name"],
|
|
"label": label,
|
|
"type": asset["type"],
|
|
"tags": ", ".join(tags),
|
|
"deprecated": deprecated,
|
|
"_document": asset
|
|
})
|
|
self.add_child(item, parent=parent)
|
|
|
|
# Add asset's children recursively if it has children
|
|
if asset["_id"] in assets:
|
|
self._add_hierarchy(assets, parent=item)
|
|
|
|
def refresh(self):
|
|
"""Refresh the data for the model."""
|
|
|
|
self.clear()
|
|
self.beginResetModel()
|
|
|
|
# Get all assets in current silo sorted by name
|
|
db_assets = io.find({"type": "asset"}).sort("name", 1)
|
|
silos = db_assets.distinct("silo") or None
|
|
# if any silo is set to None then it's expected it should not be used
|
|
if silos and None in silos:
|
|
silos = None
|
|
|
|
# Group the assets by their visual parent's id
|
|
assets_by_parent = collections.defaultdict(list)
|
|
for asset in db_assets:
|
|
parent_id = (
|
|
asset.get("data", {}).get("visualParent") or
|
|
asset.get("silo")
|
|
)
|
|
assets_by_parent[parent_id].append(asset)
|
|
|
|
# Build the hierarchical tree items recursively
|
|
self._add_hierarchy(
|
|
assets_by_parent,
|
|
parent=None,
|
|
silos=silos
|
|
)
|
|
|
|
self.endResetModel()
|
|
|
|
def flags(self, index):
|
|
return QtCore.Qt.ItemIsEnabled | QtCore.Qt.ItemIsSelectable
|
|
|
|
def data(self, index, role):
|
|
|
|
if not index.isValid():
|
|
return
|
|
|
|
item = index.internalPointer()
|
|
if role == QtCore.Qt.DecorationRole: # icon
|
|
|
|
column = index.column()
|
|
if column == self.Name:
|
|
|
|
# Allow a custom icon and custom icon color to be defined
|
|
data = item.get("_document", {}).get("data", {})
|
|
icon = data.get("icon", None)
|
|
if icon is None and item.get("type") == "silo":
|
|
icon = "database"
|
|
color = data.get("color", style.colors.default)
|
|
|
|
if icon is None:
|
|
# Use default icons if no custom one is specified.
|
|
# If it has children show a full folder, otherwise
|
|
# show an open folder
|
|
has_children = self.rowCount(index) > 0
|
|
icon = "folder" if has_children else "folder-o"
|
|
|
|
# Make the color darker when the asset is deprecated
|
|
if item.get("deprecated", False):
|
|
color = QtGui.QColor(color).darker(250)
|
|
|
|
try:
|
|
key = "fa.{0}".format(icon) # font-awesome key
|
|
icon = qtawesome.icon(key, color=color)
|
|
return icon
|
|
except Exception as exception:
|
|
# Log an error message instead of erroring out completely
|
|
# when the icon couldn't be created (e.g. invalid name)
|
|
log.error(exception)
|
|
|
|
return
|
|
|
|
if role == QtCore.Qt.ForegroundRole: # font color
|
|
if "deprecated" in item.get("tags", []):
|
|
return QtGui.QColor(style.colors.light).darker(250)
|
|
|
|
if role == self.ObjectIdRole:
|
|
return item.get("_id", None)
|
|
|
|
if role == self.DocumentRole:
|
|
return item.get("_document", None)
|
|
|
|
return super(AssetModel, self).data(index, role)
|
|
|
|
|
|
class AssetWidget(QtWidgets.QWidget):
|
|
"""A Widget to display a tree of assets with filter
|
|
|
|
To list the assets of the active project:
|
|
>>> # widget = AssetWidget()
|
|
>>> # widget.refresh()
|
|
>>> # widget.show()
|
|
|
|
"""
|
|
|
|
assets_refreshed = QtCore.Signal() # on model refresh
|
|
selection_changed = QtCore.Signal() # on view selection change
|
|
current_changed = QtCore.Signal() # on view current index change
|
|
|
|
def __init__(self, parent=None):
|
|
super(AssetWidget, self).__init__(parent=parent)
|
|
self.setContentsMargins(0, 0, 0, 0)
|
|
|
|
layout = QtWidgets.QVBoxLayout(self)
|
|
layout.setContentsMargins(0, 0, 0, 0)
|
|
layout.setSpacing(4)
|
|
|
|
# Tree View
|
|
model = AssetModel(self)
|
|
proxy = RecursiveSortFilterProxyModel()
|
|
proxy.setSourceModel(model)
|
|
proxy.setFilterCaseSensitivity(QtCore.Qt.CaseInsensitive)
|
|
|
|
view = DeselectableTreeView()
|
|
view.setIndentation(15)
|
|
view.setContextMenuPolicy(QtCore.Qt.CustomContextMenu)
|
|
view.setHeaderHidden(True)
|
|
view.setModel(proxy)
|
|
|
|
# Header
|
|
header = QtWidgets.QHBoxLayout()
|
|
|
|
icon = qtawesome.icon("fa.refresh", color=style.colors.light)
|
|
refresh = QtWidgets.QPushButton(icon, "")
|
|
refresh.setToolTip("Refresh items")
|
|
|
|
filter = QtWidgets.QLineEdit()
|
|
filter.textChanged.connect(proxy.setFilterFixedString)
|
|
filter.setPlaceholderText("Filter assets..")
|
|
|
|
header.addWidget(filter)
|
|
header.addWidget(refresh)
|
|
|
|
# Layout
|
|
layout.addLayout(header)
|
|
layout.addWidget(view)
|
|
|
|
# Signals/Slots
|
|
selection = view.selectionModel()
|
|
selection.selectionChanged.connect(self.selection_changed)
|
|
selection.currentChanged.connect(self.current_changed)
|
|
refresh.clicked.connect(self.refresh)
|
|
|
|
self.refreshButton = refresh
|
|
self.model = model
|
|
self.proxy = proxy
|
|
self.view = view
|
|
|
|
def _refresh_model(self):
|
|
with preserve_expanded_rows(
|
|
self.view, column=0, role=self.model.ObjectIdRole
|
|
):
|
|
with preserve_selection(
|
|
self.view, column=0, role=self.model.ObjectIdRole
|
|
):
|
|
self.model.refresh()
|
|
|
|
self.assets_refreshed.emit()
|
|
|
|
def refresh(self):
|
|
self._refresh_model()
|
|
|
|
def get_active_asset(self):
|
|
"""Return the asset id the current asset."""
|
|
current = self.view.currentIndex()
|
|
return current.data(self.model.ItemRole)
|
|
|
|
def get_active_index(self):
|
|
return self.view.currentIndex()
|
|
|
|
def get_selected_assets(self):
|
|
"""Return the assets' ids that are selected."""
|
|
selection = self.view.selectionModel()
|
|
rows = selection.selectedRows()
|
|
return [row.data(self.model.ObjectIdRole) for row in rows]
|
|
|
|
def select_assets(self, assets, expand=True, key="name"):
|
|
"""Select assets by name.
|
|
|
|
Args:
|
|
assets (list): List of asset names
|
|
expand (bool): Whether to also expand to the asset in the view
|
|
|
|
Returns:
|
|
None
|
|
|
|
"""
|
|
# TODO: Instead of individual selection optimize for many assets
|
|
|
|
if not isinstance(assets, (tuple, list)):
|
|
assets = [assets]
|
|
assert isinstance(
|
|
assets, (tuple, list)
|
|
), "Assets must be list or tuple"
|
|
|
|
# convert to list - tuple cant be modified
|
|
assets = list(assets)
|
|
|
|
# Clear selection
|
|
selection_model = self.view.selectionModel()
|
|
selection_model.clearSelection()
|
|
|
|
# Select
|
|
mode = selection_model.Select | selection_model.Rows
|
|
for index in iter_model_rows(
|
|
self.proxy, column=0, include_root=False
|
|
):
|
|
# stop iteration if there are no assets to process
|
|
if not assets:
|
|
break
|
|
|
|
value = index.data(self.model.ItemRole).get(key)
|
|
if value not in assets:
|
|
continue
|
|
|
|
# Remove processed asset
|
|
assets.pop(assets.index(value))
|
|
|
|
selection_model.select(index, mode)
|
|
|
|
if expand:
|
|
# Expand parent index
|
|
self.view.expand(self.proxy.parent(index))
|
|
|
|
# Set the currently active index
|
|
self.view.setCurrentIndex(index)
|