import collections import logging from functools import partial from Qt import QtWidgets, QtCore from avalon import io, api, style from avalon.vendor import qtawesome from avalon.lib import HeroVersionType from avalon.tools import lib as tools_lib from openpype.modules import ModulesManager from .switch_dialog import SwitchAssetDialog from .model import InventoryModel DEFAULT_COLOR = "#fb9c15" log = logging.getLogger("SceneInventory") class SceneInvetoryView(QtWidgets.QTreeView): data_changed = QtCore.Signal() hierarchy_view_changed = QtCore.Signal(bool) def __init__(self, parent=None): super(SceneInvetoryView, self).__init__(parent=parent) # view settings self.setIndentation(12) self.setAlternatingRowColors(True) self.setSortingEnabled(True) self.setSelectionMode(self.ExtendedSelection) self.setContextMenuPolicy(QtCore.Qt.CustomContextMenu) self.customContextMenuRequested.connect(self._show_right_mouse_menu) self._hierarchy_view = False self._selected = None manager = ModulesManager() self.sync_server = manager.modules_by_name["sync_server"] self.sync_enabled = self.sync_server.enabled def _set_hierarchy_view(self, enabled): if enabled == self._hierarchy_view: return self._hierarchy_view = enabled self.hierarchy_view_changed.emit(enabled) def _enter_hierarchy(self, items): self._selected = set(i["objectName"] for i in items) self._set_hierarchy_view(True) self.data_changed.emit() self.expandToDepth(1) self.setStyleSheet(""" QTreeView { border-color: #fb9c15; } """) def _leave_hierarchy(self): self._set_hierarchy_view(False) self.data_changed.emit() self.setStyleSheet("QTreeView {}") def _build_item_menu_for_selection(self, items, menu): if not items: return repre_ids = [] for item in items: item_id = io.ObjectId(item["representation"]) if item_id not in repre_ids: repre_ids.append(item_id) repre_docs = io.find( { "type": "representation", "_id": {"$in": repre_ids} }, {"parent": 1} ) version_ids = [] for repre_doc in repre_docs: version_id = repre_doc["parent"] if version_id not in version_ids: version_ids.append(version_id) loaded_versions = io.find({ "_id": {"$in": version_ids}, "type": {"$in": ["version", "hero_version"]} }) loaded_hero_versions = [] versions_by_parent_id = collections.defaultdict(list) version_parents = [] for version in loaded_versions: if version["type"] == "hero_version": loaded_hero_versions.append(version) else: parent_id = version["parent"] versions_by_parent_id[parent_id].append(version) if parent_id not in version_parents: version_parents.append(parent_id) all_versions = io.find({ "type": {"$in": ["hero_version", "version"]}, "parent": {"$in": version_parents} }) hero_versions = [] versions = [] for version in all_versions: if version["type"] == "hero_version": hero_versions.append(version) else: versions.append(version) has_loaded_hero_versions = len(loaded_hero_versions) > 0 has_available_hero_version = len(hero_versions) > 0 has_outdated = False for version in versions: parent_id = version["parent"] current_versions = versions_by_parent_id[parent_id] for current_version in current_versions: if current_version["name"] < version["name"]: has_outdated = True break if has_outdated: break switch_to_versioned = None if has_loaded_hero_versions: def _on_switch_to_versioned(items): repre_ids = [] for item in items: item_id = io.ObjectId(item["representation"]) if item_id not in repre_ids: repre_ids.append(item_id) repre_docs = io.find( { "type": "representation", "_id": {"$in": repre_ids} }, {"parent": 1} ) version_ids = [] version_id_by_repre_id = {} for repre_doc in repre_docs: version_id = repre_doc["parent"] version_id_by_repre_id[repre_doc["_id"]] = version_id if version_id not in version_ids: version_ids.append(version_id) hero_versions = io.find( { "_id": {"$in": version_ids}, "type": "hero_version" }, {"version_id": 1} ) version_ids = set() for hero_version in hero_versions: version_id = hero_version["version_id"] version_ids.add(version_id) hero_version_id = hero_version["_id"] for _repre_id, current_version_id in ( version_id_by_repre_id.items() ): if current_version_id == hero_version_id: version_id_by_repre_id[_repre_id] = version_id version_docs = io.find( { "_id": {"$in": list(version_ids)}, "type": "version" }, {"name": 1} ) version_name_by_id = {} for version_doc in version_docs: version_name_by_id[version_doc["_id"]] = \ version_doc["name"] for item in items: repre_id = io.ObjectId(item["representation"]) version_id = version_id_by_repre_id.get(repre_id) version_name = version_name_by_id.get(version_id) if version_name is not None: try: api.update(item, version_name) except AssertionError: self._show_version_error_dialog( version_name, [item] ) log.warning("Update failed", exc_info=True) self.data_changed.emit() update_icon = qtawesome.icon( "fa.asterisk", color=DEFAULT_COLOR ) switch_to_versioned = QtWidgets.QAction( update_icon, "Switch to versioned", menu ) switch_to_versioned.triggered.connect( lambda: _on_switch_to_versioned(items) ) update_to_latest_action = None if has_outdated or has_loaded_hero_versions: # update to latest version def _on_update_to_latest(items): for item in items: try: api.update(item, -1) except AssertionError: self._show_version_error_dialog(None, [item]) log.warning("Update failed", exc_info=True) self.data_changed.emit() update_icon = qtawesome.icon( "fa.angle-double-up", color=DEFAULT_COLOR ) update_to_latest_action = QtWidgets.QAction( update_icon, "Update to latest", menu ) update_to_latest_action.triggered.connect( lambda: _on_update_to_latest(items) ) change_to_hero = None if has_available_hero_version: # change to hero version def _on_update_to_hero(items): for item in items: try: api.update(item, HeroVersionType(-1)) except AssertionError: self._show_version_error_dialog('hero', [item]) log.warning("Update failed", exc_info=True) self.data_changed.emit() # TODO change icon change_icon = qtawesome.icon( "fa.asterisk", color="#00b359" ) change_to_hero = QtWidgets.QAction( change_icon, "Change to hero", menu ) change_to_hero.triggered.connect( lambda: _on_update_to_hero(items) ) # set version set_version_icon = qtawesome.icon("fa.hashtag", color=DEFAULT_COLOR) set_version_action = QtWidgets.QAction( set_version_icon, "Set version", menu ) set_version_action.triggered.connect( lambda: self._show_version_dialog(items)) # switch asset switch_asset_icon = qtawesome.icon("fa.sitemap", color=DEFAULT_COLOR) switch_asset_action = QtWidgets.QAction( switch_asset_icon, "Switch Asset", menu ) switch_asset_action.triggered.connect( lambda: self._show_switch_dialog(items)) # remove remove_icon = qtawesome.icon("fa.remove", color=DEFAULT_COLOR) remove_action = QtWidgets.QAction(remove_icon, "Remove items", menu) remove_action.triggered.connect( lambda: self._show_remove_warning_dialog(items)) # add the actions if switch_to_versioned: menu.addAction(switch_to_versioned) if update_to_latest_action: menu.addAction(update_to_latest_action) if change_to_hero: menu.addAction(change_to_hero) menu.addAction(set_version_action) menu.addAction(switch_asset_action) menu.addSeparator() menu.addAction(remove_action) self._handle_sync_server(menu, repre_ids) def _handle_sync_server(self, menu, repre_ids): """ Adds actions for download/upload when SyncServer is enabled Args: menu (OptionMenu) repre_ids (list) of object_ids Returns: (OptionMenu) """ if not self.sync_enabled: return menu.addSeparator() download_icon = qtawesome.icon("fa.download", color=DEFAULT_COLOR) download_active_action = QtWidgets.QAction( download_icon, "Download", menu ) download_active_action.triggered.connect( lambda: self._add_sites(repre_ids, 'active_site')) upload_icon = qtawesome.icon("fa.upload", color=DEFAULT_COLOR) upload_remote_action = QtWidgets.QAction( upload_icon, "Upload", menu ) upload_remote_action.triggered.connect( lambda: self._add_sites(repre_ids, 'remote_site')) menu.addAction(download_active_action) menu.addAction(upload_remote_action) def _add_sites(self, repre_ids, side): """ (Re)sync all 'repre_ids' to specific site. It checks if opposite site has fully available content to limit accidents. (ReSync active when no remote >> losing active content) Args: repre_ids (list) side (str): 'active_site'|'remote_site' """ project_name = io.Session["AVALON_PROJECT"] active_site = self.sync_server.get_active_site(project_name) remote_site = self.sync_server.get_remote_site(project_name) repre_docs = io.find({ "type": "representation", "_id": {"$in": repre_ids} }) repre_docs_by_id = { repre_doc["_id"]: repre_doc for repre_doc in repre_docs } for repre_id in repre_ids: repre_doc = repre_docs_by_id.get(repre_id) if not repre_doc: continue progress = tools_lib.get_progress_for_repre( repre_doc, active_site, remote_site ) if side == "active_site": # check opposite from added site, must be 1 or unable to sync check_progress = progress[remote_site] site = active_site else: check_progress = progress[active_site] site = remote_site if check_progress == 1: self.sync_server.add_site( project_name, repre_id, site, force=True ) self.data_changed.emit() def _build_item_menu(self, items=None): """Create menu for the selected items""" if not items: items = [] menu = QtWidgets.QMenu(self) # add the actions self._build_item_menu_for_selection(items, menu) # These two actions should be able to work without selection # expand all items expandall_action = QtWidgets.QAction(menu, text="Expand all items") expandall_action.triggered.connect(self.expandAll) # collapse all items collapse_action = QtWidgets.QAction(menu, text="Collapse all items") collapse_action.triggered.connect(self.collapseAll) menu.addAction(expandall_action) menu.addAction(collapse_action) custom_actions = self._get_custom_actions(containers=items) if custom_actions: submenu = QtWidgets.QMenu("Actions", self) for action in custom_actions: color = action.color or DEFAULT_COLOR icon = qtawesome.icon("fa.%s" % action.icon, color=color) action_item = QtWidgets.QAction(icon, action.label, submenu) action_item.triggered.connect( partial(self._process_custom_action, action, items)) submenu.addAction(action_item) menu.addMenu(submenu) # go back to flat view if self._hierarchy_view: back_to_flat_icon = qtawesome.icon("fa.list", color=DEFAULT_COLOR) back_to_flat_action = QtWidgets.QAction( back_to_flat_icon, "Back to Full-View", menu ) back_to_flat_action.triggered.connect(self._leave_hierarchy) # send items to hierarchy view enter_hierarchy_icon = qtawesome.icon("fa.indent", color="#d8d8d8") enter_hierarchy_action = QtWidgets.QAction( enter_hierarchy_icon, "Cherry-Pick (Hierarchy)", menu ) enter_hierarchy_action.triggered.connect( lambda: self._enter_hierarchy(items)) if items: menu.addAction(enter_hierarchy_action) if self._hierarchy_view: menu.addAction(back_to_flat_action) return menu def _get_custom_actions(self, containers): """Get the registered Inventory Actions Args: containers(list): collection of containers Returns: list: collection of filter and initialized actions """ def sorter(Plugin): """Sort based on order attribute of the plugin""" return Plugin.order # Fedd an empty dict if no selection, this will ensure the compat # lookup always work, so plugin can interact with Scene Inventory # reversely. containers = containers or [dict()] # Check which action will be available in the menu Plugins = api.discover(api.InventoryAction) compatible = [p() for p in Plugins if any(p.is_compatible(c) for c in containers)] return sorted(compatible, key=sorter) def _process_custom_action(self, action, containers): """Run action and if results are returned positive update the view If the result is list or dict, will select view items by the result. Args: action (InventoryAction): Inventory Action instance containers (list): Data of currently selected items Returns: None """ result = action.process(containers) if result: self.data_changed.emit() if isinstance(result, (list, set)): self._select_items_by_action(result) if isinstance(result, dict): self._select_items_by_action( result["objectNames"], result["options"] ) def _select_items_by_action(self, object_names, options=None): """Select view items by the result of action Args: object_names (list or set): A list/set of container object name options (dict): GUI operation options. Returns: None """ options = options or dict() if options.get("clear", True): self.clearSelection() object_names = set(object_names) if ( self._hierarchy_view and not self._selected.issuperset(object_names) ): # If any container not in current cherry-picked view, update # view before selecting them. self._selected.update(object_names) self.data_changed.emit() model = self.model() selection_model = self.selectionModel() select_mode = { "select": selection_model.Select, "deselect": selection_model.Deselect, "toggle": selection_model.Toggle, }[options.get("mode", "select")] for item in tools_lib.iter_model_rows(model, 0): item = item.data(InventoryModel.ItemRole) if item.get("isGroupNode"): continue name = item.get("objectName") if name in object_names: self.scrollTo(item) # Ensure item is visible flags = select_mode | selection_model.Rows selection_model.select(item, flags) object_names.remove(name) if len(object_names) == 0: break def _show_right_mouse_menu(self, pos): """Display the menu when at the position of the item clicked""" globalpos = self.viewport().mapToGlobal(pos) if not self.selectionModel().hasSelection(): print("No selection") # Build menu without selection, feed an empty list menu = self._build_item_menu() menu.exec_(globalpos) return active = self.currentIndex() # index under mouse active = active.sibling(active.row(), 0) # get first column # move index under mouse indices = self.get_indices() if active in indices: indices.remove(active) indices.append(active) # Extend to the sub-items all_indices = self._extend_to_children(indices) items = [dict(i.data(InventoryModel.ItemRole)) for i in all_indices if i.parent().isValid()] if self._hierarchy_view: # Ensure no group item items = [n for n in items if not n.get("isGroupNode")] menu = self._build_item_menu(items) menu.exec_(globalpos) def get_indices(self): """Get the selected rows""" selection_model = self.selectionModel() return selection_model.selectedRows() def _extend_to_children(self, indices): """Extend the indices to the children indices. Top-level indices are extended to its children indices. Sub-items are kept as is. Args: indices (list): The indices to extend. Returns: list: The children indices """ def get_children(i): model = i.model() rows = model.rowCount(parent=i) for row in range(rows): child = model.index(row, 0, parent=i) yield child subitems = set() for i in indices: valid_parent = i.parent().isValid() if valid_parent and i not in subitems: subitems.add(i) if self._hierarchy_view: # Assume this is a group item for child in get_children(i): subitems.add(child) else: # is top level item for child in get_children(i): subitems.add(child) return list(subitems) def _show_version_dialog(self, items): """Create a dialog with the available versions for the selected file Args: items (list): list of items to run the "set_version" for Returns: None """ active = items[-1] # Get available versions for active representation representation_id = io.ObjectId(active["representation"]) representation = io.find_one({"_id": representation_id}) version = io.find_one({ "_id": representation["parent"] }) versions = list(io.find( { "parent": version["parent"], "type": "version" }, sort=[("name", 1)] )) hero_version = io.find_one({ "parent": version["parent"], "type": "hero_version" }) if hero_version: _version_id = hero_version["version_id"] for _version in versions: if _version["_id"] != _version_id: continue hero_version["name"] = HeroVersionType( _version["name"] ) hero_version["data"] = _version["data"] break # Get index among the listed versions current_item = None current_version = active["version"] if isinstance(current_version, HeroVersionType): current_item = hero_version else: for version in versions: if version["name"] == current_version: current_item = version break all_versions = [] if hero_version: all_versions.append(hero_version) all_versions.extend(reversed(versions)) if current_item: index = all_versions.index(current_item) else: index = 0 versions_by_label = dict() labels = [] for version in all_versions: is_hero = version["type"] == "hero_version" label = tools_lib.format_version(version["name"], is_hero) labels.append(label) versions_by_label[label] = version["name"] label, state = QtWidgets.QInputDialog.getItem( self, "Set version..", "Set version number to", labels, current=index, editable=False ) if not state: return if label: version = versions_by_label[label] for item in items: try: api.update(item, version) except AssertionError: self._show_version_error_dialog(version, [item]) log.warning("Update failed", exc_info=True) # refresh model when done self.data_changed.emit() def _show_switch_dialog(self, items): """Display Switch dialog""" dialog = SwitchAssetDialog(self, items) dialog.switched.connect(self.data_changed.emit) dialog.show() def _show_remove_warning_dialog(self, items): """Prompt a dialog to inform the user the action will remove items""" accept = QtWidgets.QMessageBox.Ok buttons = accept | QtWidgets.QMessageBox.Cancel state = QtWidgets.QMessageBox.question( self, "Are you sure?", "Are you sure you want to remove {} item(s)".format(len(items)), buttons=buttons, defaultButton=accept ) if state != accept: return for item in items: api.remove(item) self.data_changed.emit() def _show_version_error_dialog(self, version, items): """Shows QMessageBox when version switch doesn't work Args: version: str or int or None """ if not version: version_str = "latest" elif version == "hero": version_str = "hero" elif isinstance(version, int): version_str = "v{:03d}".format(version) else: version_str = version dialog = QtWidgets.QMessageBox() dialog.setIcon(QtWidgets.QMessageBox.Warning) dialog.setStyleSheet(style.load_stylesheet()) dialog.setWindowTitle("Update failed") switch_btn = dialog.addButton( "Switch Asset", QtWidgets.QMessageBox.ActionRole ) switch_btn.clicked.connect(lambda: self._show_switch_dialog(items)) dialog.addButton(QtWidgets.QMessageBox.Cancel) msg = ( "Version update to '{}' failed as representation doesn't exist." "\n\nPlease update to version with a valid representation" " OR \n use 'Switch Asset' button to change asset." ).format(version_str) dialog.setText(msg) dialog.exec_()