From 821990662a496ad33eb59b12bc339b582fe3a9d4 Mon Sep 17 00:00:00 2001 From: iLLiCiTiT Date: Wed, 17 Nov 2021 21:54:07 +0100 Subject: [PATCH] copied and reimplemented asset model --- openpype/tools/utils/assets_widget.py | 238 ++++++++++++++++++++++++++ 1 file changed, 238 insertions(+) diff --git a/openpype/tools/utils/assets_widget.py b/openpype/tools/utils/assets_widget.py index 3a501d40d1..1b62e0b2c5 100644 --- a/openpype/tools/utils/assets_widget.py +++ b/openpype/tools/utils/assets_widget.py @@ -1,8 +1,14 @@ +import collections + import Qt from Qt import QtWidgets, QtCore, QtGui +from avalon import style +from avalon.vendor import qtawesome + from openpype.style import get_objected_colors +from .lib import DynamicQThread from .views import ( TreeViewSpinner, DeselectableTreeView @@ -220,3 +226,235 @@ class UnderlinesAssetDelegate(QtWidgets.QItemDelegate): ) painter.restore() + + +class AssetModel(QtGui.QStandardItemModel): + """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. + + """ + + _doc_fetched = QtCore.Signal() + refreshed = QtCore.Signal(bool) + + # Asset document projection + _asset_projection = { + "type": 1, + "schema": 1, + "name": 1, + "parent": 1, + "data.visualParent": 1, + "data.label": 1, + "data.tags": 1, + "data.icon": 1, + "data.color": 1, + "data.deprecated": 1 + } + + def __init__(self, dbcon, parent=None): + super(AssetModel, self).__init__(parent=parent) + self.dbcon = dbcon + + self._doc_fetching_thread = None + self._doc_fetching_stop = False + self._doc_payload = [] + + self._doc_fetched.connect(self._on_docs_fetched) + + self._items_with_color_by_id = {} + self._items_by_asset_id = {} + + def get_index_by_asset_id(self, asset_id): + item = self._items_by_asset_id.get(asset_id) + if item is not None: + return item.index() + return QtCore.QModelIndex() + + def get_indexes_by_asset_ids(self, asset_ids): + return [ + self.get_index_by_asset_id(asset_id) + for asset_id in asset_ids + ] + + def get_index_by_asset_name(self, asset_name): + indexes = self.get_indexes_by_asset_names([asset_name]) + for index in indexes: + if index.isValid(): + return index + return indexes[0] + + def get_indexes_by_asset_names(self, asset_names): + asset_ids_by_name = { + asset_name: None + for asset_name in asset_names + } + + for asset_id, item in self._items_by_asset_id.items(): + asset_name = item.data(ASSET_NAME_ROLE) + if asset_name in asset_ids_by_name: + asset_ids_by_name[asset_name] = asset_id + + asset_ids = [ + asset_ids_by_name[asset_name] + for asset_name in asset_names + ] + + return self.get_indexes_by_asset_ids(asset_ids) + + def refresh(self, force=False): + """Refresh the data for the model.""" + # Skip fetch if there is already other thread fetching documents + if self._doc_fetching_thread is not None: + if not force: + return + self._stop_fetch_thread() + + # Fetch documents from mongo + # Restart payload + self._doc_payload = [] + self._doc_fetching_stop = False + self._doc_fetching_thread = DynamicQThread(self._threaded_fetch) + self._doc_fetching_thread.start() + + def clear_underlines(self): + for asset_id in tuple(self._items_with_color_by_id.keys()): + item = self._items_with_color_by_id.pop(asset_id) + item.setData(None, ASSET_UNDERLINE_COLORS_ROLE) + + def set_underline_colors(self, colors_by_asset_id): + self.clear_underlines() + + for asset_id, colors in colors_by_asset_id.items(): + item = self._items_by_asset_id.get(asset_id) + if item is None: + continue + item.setData(colors, ASSET_UNDERLINE_COLORS_ROLE) + + def _on_docs_fetched(self): + asset_docs = self._doc_payload + + asset_ids = set() + asset_docs_by_id = {} + asset_ids_by_parents = collections.defaultdict(set) + for asset_doc in asset_docs: + asset_id = asset_doc["_id"] + asset_data = asset_doc.get("data") or {} + parent_id = asset_data.get("visualParent") + asset_ids.add(asset_id) + asset_docs_by_id[asset_id] = asset_doc + asset_ids_by_parents[parent_id].add(asset_id) + + root_item = self.invisibleRootItem() + asset_items_queue = collections.deque() + asset_items_queue.append((None, root_item)) + + removed_asset_ids = set() + while asset_items_queue: + parent_id, parent_item = asset_items_queue.popleft() + children_ids = asset_ids_by_parents[parent_id] + if not children_ids: + continue + + for row in reversed(range(parent_item.rowCount())): + child_item = parent_item.child(row, 0) + asset_id = child_item.data(ASSET_ID_ROLE) + if asset_id not in children_ids: + parent_item.removeRow(row) + if asset_id not in asset_docs_by_id: + removed_asset_ids.add(asset_id) + continue + + children_ids.remove(asset_id) + asset_items_queue.append((asset_id, child_item)) + + new_items = [] + for asset_id in children_ids: + item = QtGui.QStandardItem() + item.setEditable(False) + item.setData(asset_id, ASSET_ID_ROLE) + new_items.append(item) + self._items_by_asset_id[asset_id] = item + asset_items_queue.append((asset_id, item)) + + if new_items: + parent_item.appendRows(new_items) + + for asset_id in removed_asset_ids: + self._items_by_asset_id.pop(asset_id) + if asset_id in self._items_with_color_by_id: + self._items_with_color_by_id.pop(asset_id) + + # Refresh data + for asset_id, item in self._items_by_asset_id.items(): + asset_doc = asset_docs_by_id[asset_id] + + asset_name = asset_doc["name"] + if item.data(ASSET_NAME_ROLE) != asset_name: + item.setData(asset_name, ASSET_NAME_ROLE) + + asset_data = asset_doc.get("data") or {} + asset_label = asset_data.get("label") or asset_name + if item.data(ASSET_LABEL_ROLE) != asset_label: + item.setData(asset_label, QtCore.Qt.DisplayRole) + item.setData(asset_label, ASSET_LABEL_ROLE) + + icon_color = asset_data.get("color") or style.colors.default + icon_name = asset_data.get("icon") + if not icon_name: + # Use default icons if no custom one is specified. + # If it has children show a full folder, otherwise + # show an open folder + if item.rowCount() > 0: + icon_name = "folder" + else: + icon_name = "folder-o" + + try: + # font-awesome key + full_icon_name = "fa.{0}".format(icon_name) + icon = qtawesome.icon(full_icon_name, color=icon_color) + item.setData(icon, QtCore.Qt.DecorationRole) + + except Exception as exception: + pass + + self.refreshed.emit(bool(self._items_by_asset_id)) + + self._stop_fetch_thread() + + def _threaded_fetch(self): + asset_docs = self._fetch_asset_docs() or [] + if self._doc_fetching_stop: + return + + self._doc_payload = asset_docs + + # Emit doc fetched only if was not stopped + self._doc_fetched.emit() + + def _fetch_asset_docs(self): + if not self.dbcon.Session.get("AVALON_PROJECT"): + return + + project_doc = self.dbcon.find_one( + {"type": "project"}, + {"_id": True} + ) + if not project_doc: + return + + # Get all assets sorted by name + return list(self.dbcon.find( + {"type": "asset"}, + self._asset_projection + )) + + def _stop_fetch_thread(self): + if self._doc_fetching_thread is not None: + self._doc_fetching_stop = True + while self._doc_fetching_thread.isRunning(): + time.sleep(0.01) + self._doc_fetching_thread = None