import copy import re import math import time from uuid import uuid4 from qtpy import QtCore, QtGui import qtawesome from openpype.client import ( get_assets, get_subsets, get_last_versions, get_versions, get_hero_versions, get_version_by_name, get_representations ) from openpype.pipeline import ( registered_host, HeroVersionType, schema, ) from openpype.style import get_default_entity_icon_color from openpype.tools.utils.models import TreeModel, Item from openpype.tools.utils import lib from openpype.host import ILoadHost from openpype.modules import ModulesManager from openpype.tools.utils.constants import ( LOCAL_PROVIDER_ROLE, REMOTE_PROVIDER_ROLE, LOCAL_AVAILABILITY_ROLE, REMOTE_AVAILABILITY_ROLE ) ITEM_ID_ROLE = QtCore.Qt.UserRole + 90 def is_filtering_recursible(): """Does Qt binding support recursive filtering for QSortFilterProxyModel? (NOTE) Recursive filtering was introduced in Qt 5.10. """ return hasattr(QtCore.QSortFilterProxyModel, "setRecursiveFilteringEnabled") class BaseRepresentationModel(object): """Methods for SyncServer useful in multiple models""" # Cheap & hackish way how to avoid refreshing of whole sync server module # on each selection change _last_project = None _modules_manager = None _last_project_cache = 0 _last_manager_cache = 0 _max_project_cache_time = 30 _max_manager_cache_time = 60 def reset_sync_server(self, project_name=None): """Sets/Resets sync server vars after every change (refresh.)""" repre_icons = {} sync_server = None active_site = active_provider = None remote_site = remote_provider = None if not project_name: project_name = self.dbcon.active_project() else: self.dbcon.Session["AVALON_PROJECT"] = project_name if not project_name: self.repre_icons = repre_icons self.sync_server = sync_server self.active_site = active_site self.active_provider = active_provider self.remote_site = remote_site self.remote_provider = remote_provider return now_time = time.time() project_cache_diff = now_time - self._last_project_cache if project_cache_diff > self._max_project_cache_time: self._last_project = None if project_name == self._last_project: return self._last_project = project_name self._last_project_cache = now_time manager_cache_diff = now_time - self._last_manager_cache if manager_cache_diff > self._max_manager_cache_time: self._modules_manager = None if self._modules_manager is None: self._modules_manager = ModulesManager() self._last_manager_cache = now_time sync_server = self._modules_manager.modules_by_name["sync_server"] if sync_server.is_project_enabled(project_name, single=True): active_site = sync_server.get_active_site(project_name) active_provider = sync_server.get_provider_for_site( project_name, active_site) if active_site == 'studio': # for studio use explicit icon active_provider = 'studio' remote_site = sync_server.get_remote_site(project_name) remote_provider = sync_server.get_provider_for_site( project_name, remote_site) if remote_site == 'studio': # for studio use explicit icon remote_provider = 'studio' repre_icons = lib.get_repre_icons() self.repre_icons = repre_icons self.sync_server = sync_server self.active_site = active_site self.active_provider = active_provider self.remote_site = remote_site self.remote_provider = remote_provider class SubsetsModel(TreeModel, BaseRepresentationModel): doc_fetched = QtCore.Signal() refreshed = QtCore.Signal(bool) Columns = [ "subset", "asset", "family", "version", "time", "author", "frames", "duration", "handles", "step", "loaded_in_scene", "repre_info" ] column_labels_mapping = { "subset": "Subset", "asset": "Asset", "family": "Family", "version": "Version", "time": "Time", "author": "Author", "frames": "Frames", "duration": "Duration", "handles": "Handles", "step": "Step", "loaded_in_scene": "In scene", "repre_info": "Availability" } SortAscendingRole = QtCore.Qt.UserRole + 2 SortDescendingRole = QtCore.Qt.UserRole + 3 merged_subset_colors = [ (55, 161, 222), # Light Blue (231, 176, 0), # Yellow (154, 13, 255), # Purple (130, 184, 30), # Light Green (211, 79, 63), # Light Red (179, 181, 182), # Grey (194, 57, 179), # Pink (0, 120, 215), # Dark Blue (0, 204, 106), # Dark Green (247, 99, 12), # Orange ] not_last_hero_brush = QtGui.QBrush(QtGui.QColor(254, 121, 121)) # Should be minimum of required asset document keys asset_doc_projection = { "name": 1, "label": 1 } # Should be minimum of required subset document keys subset_doc_projection = { "name": 1, "parent": 1, "schema": 1, "data.families": 1, "data.subsetGroup": 1 } def __init__( self, dbcon, groups_config, family_config_cache, grouping=True, parent=None, asset_doc_projection=None, subset_doc_projection=None ): super(SubsetsModel, self).__init__(parent=parent) self.dbcon = dbcon # Projections for Mongo queries # - let ability to modify them if used in tools that require more than # defaults if asset_doc_projection: self.asset_doc_projection = asset_doc_projection if subset_doc_projection: self.subset_doc_projection = subset_doc_projection self.repre_icons = {} self.sync_server = None self.active_site = self.active_provider = None self.columns_index = dict( (key, idx) for idx, key in enumerate(self.Columns) ) self._asset_ids = None self.groups_config = groups_config self.family_config_cache = family_config_cache self._sorter = None self._grouping = grouping self._icons = { "subset": qtawesome.icon( "fa.file-o", color=get_default_entity_icon_color() ) } self._items_by_id = {} self._doc_fetching_thread = None self._doc_fetching_stop = False self._doc_payload = {} self._host = registered_host() self._loaded_representation_ids = set() # Refresh loaded scene containers only every 3 seconds at most self._host_loaded_refresh_timeout = 3 self._host_loaded_refresh_time = 0 self.doc_fetched.connect(self._on_doc_fetched) self.refresh() def get_item_by_id(self, item_id): return self._items_by_id.get(item_id) def add_child(self, new_item, *args, **kwargs): item_id = str(uuid4()) new_item["id"] = item_id self._items_by_id[item_id] = new_item super(SubsetsModel, self).add_child(new_item, *args, **kwargs) def set_assets(self, asset_ids): self._asset_ids = asset_ids self.refresh() def set_grouping(self, state): self._grouping = state self._on_doc_fetched() def get_subsets_families(self): return self._doc_payload.get("subset_families") or set() def setData(self, index, value, role=QtCore.Qt.EditRole): # Trigger additional edit when `version` column changed # because it also updates the information in other columns if index.column() == self.columns_index["version"]: item = index.internalPointer() subset_id = item["_id"] if isinstance(value, HeroVersionType): version_doc = self._get_hero_version(subset_id) else: project_name = self.dbcon.active_project() version_doc = get_version_by_name( project_name, value, subset_id ) # update availability on active site when version changes if self.sync_server.enabled and version_doc: repres_info = list( self.sync_server.get_repre_info_for_versions( project_name, [version_doc["_id"]], self.active_site, self.remote_site ) ) if repres_info: version_doc["data"].update( self._get_repre_dict(repres_info[0])) self.set_version(index, version_doc) return super(SubsetsModel, self).setData(index, value, role) def _get_hero_version(self, subset_id): project_name = self.dbcon.active_project() version_docs = get_versions( project_name, subset_ids=[subset_id], hero=True ) standard_versions = [] hero_version_doc = None for version_doc in version_docs: if version_doc["type"] == "hero_version": hero_version_doc = version_doc continue standard_versions.append(version_doc) src_version_id = hero_version_doc["version_id"] src_version = None is_from_latest = True for version_doc in reversed(sorted( standard_versions, key=lambda item: item["name"] )): if version_doc["_id"] == src_version_id: src_version = version_doc break is_from_latest = False hero_version_doc["data"] = src_version["data"] hero_version_doc["name"] = src_version["name"] hero_version_doc["is_from_latest"] = is_from_latest return hero_version_doc def set_version(self, index, version): """Update the version data of the given index. Arguments: index (QtCore.QModelIndex): The model index. version (dict) Version document in the database. """ assert isinstance(index, QtCore.QModelIndex) if not index.isValid(): return item = index.internalPointer() assert version["parent"] == item["_id"], ( "Version does not belong to subset" ) # Get the data from the version version_data = version.get("data", dict()) # Compute frame ranges (if data is present) frame_start = version_data.get( "frameStart", # backwards compatibility version_data.get("startFrame", None) ) frame_end = version_data.get( "frameEnd", # backwards compatibility version_data.get("endFrame", None) ) handle_start = version_data.get("handleStart", None) handle_end = version_data.get("handleEnd", None) if handle_start is not None and handle_end is not None: handles = "{}-{}".format(str(handle_start), str(handle_end)) if frame_start is not None and frame_end is not None: # Remove superfluous zeros from numbers (3.0 -> 3) to improve # readability for most frame ranges start_clean = ("%f" % frame_start).rstrip("0").rstrip(".") end_clean = ("%f" % frame_end).rstrip("0").rstrip(".") frames = "{0}-{1}".format(start_clean, end_clean) duration = frame_end - frame_start + 1 else: frames = None duration = None schema_maj_version, _ = schema.get_schema_version(item["schema"]) if schema_maj_version < 3: families = version_data.get("families", [None]) else: families = item["data"]["families"] family = None if families: family = families[0] family_config = self.family_config_cache.family_config(family) item.update({ "version": version["name"], "version_document": version, "author": version_data.get("author", None), "time": version_data.get("time", None), "family": family, "familyLabel": family_config.get("label", family), "familyIcon": family_config.get("icon", None), "families": set(families), "frameStart": frame_start, "frameEnd": frame_end, "duration": duration, "handles": handles, "frames": frames, "step": version_data.get("step", None), }) repre_info = version_data.get("repre_info") if repre_info: item["repre_info"] = repre_info def _fetch(self): project_name = self.dbcon.active_project() asset_docs = get_assets( project_name, asset_ids=self._asset_ids, fields=self.asset_doc_projection.keys() ) asset_docs_by_id = { asset_doc["_id"]: asset_doc for asset_doc in asset_docs } subset_docs_by_id = {} subset_docs = get_subsets( project_name, asset_ids=self._asset_ids, fields=self.subset_doc_projection.keys() ) subset_families = set() for subset_doc in subset_docs: if self._doc_fetching_stop: return families = subset_doc.get("data", {}).get("families") if families: subset_families.add(families[0]) subset_docs_by_id[subset_doc["_id"]] = subset_doc subset_ids = list(subset_docs_by_id.keys()) last_versions_by_subset_id = get_last_versions( project_name, subset_ids, fields=["_id", "parent", "name", "type", "data", "schema"] ) hero_versions = get_hero_versions(project_name, subset_ids=subset_ids) missing_versions = [] for hero_version in hero_versions: version_id = hero_version["version_id"] if version_id not in last_versions_by_subset_id: missing_versions.append(version_id) missing_versions_by_id = {} if missing_versions: missing_version_docs = get_versions( project_name, version_ids=missing_versions ) missing_versions_by_id = { missing_version_doc["_id"]: missing_version_doc for missing_version_doc in missing_version_docs } for hero_version in hero_versions: version_id = hero_version["version_id"] subset_id = hero_version["parent"] version_doc = last_versions_by_subset_id.get(subset_id) if version_doc is None: version_doc = missing_versions_by_id.get(version_id) if version_doc is None: continue hero_version["data"] = version_doc["data"] hero_version["name"] = HeroVersionType(version_doc["name"]) # Add information if hero version is from latest version hero_version["is_from_latest"] = version_id == version_doc["_id"] last_versions_by_subset_id[subset_id] = hero_version # Check loaded subsets loaded_subset_ids = set() ids = self._loaded_representation_ids if ids: if self._doc_fetching_stop: return # Get subset ids from loaded representations in workfile # todo: optimize with aggregation query to distinct subset id representations = get_representations(project_name, representation_ids=ids, fields=["parent"]) version_ids = set(repre["parent"] for repre in representations) versions = get_versions(project_name, version_ids=version_ids, fields=["parent"]) loaded_subset_ids = set(version["parent"] for version in versions) if self._doc_fetching_stop: return repre_info_by_version_id = {} if self.sync_server.enabled: versions_by_id = {} for _subset_id, doc in last_versions_by_subset_id.items(): versions_by_id[doc["_id"]] = doc repres_info = self.sync_server.get_repre_info_for_versions( project_name, list(versions_by_id.keys()), self.active_site, self.remote_site ) for repre_info in repres_info: if self._doc_fetching_stop: return version_id = repre_info["_id"] doc = versions_by_id[version_id] doc["active_provider"] = self.active_provider doc["remote_provider"] = self.remote_provider repre_info_by_version_id[version_id] = repre_info self._doc_payload = { "asset_docs_by_id": asset_docs_by_id, "subset_docs_by_id": subset_docs_by_id, "subset_families": subset_families, "last_versions_by_subset_id": last_versions_by_subset_id, "repre_info_by_version_id": repre_info_by_version_id, "subsets_loaded_by_id": loaded_subset_ids } self.doc_fetched.emit() def fetch_subset_and_version(self): """Query all subsets and latest versions from aggregation (NOTE) The returned version documents are NOT the real version document, it's generated from the MongoDB's aggregation so some of the first level field may not be presented. """ self._doc_payload = {} self._doc_fetching_stop = False self._doc_fetching_thread = lib.create_qthread(self._fetch) self._doc_fetching_thread.start() def stop_fetch_thread(self): if self._doc_fetching_thread is not None: self._doc_fetching_stop = True while self._doc_fetching_thread.isRunning(): pass def refresh(self): self.stop_fetch_thread() self.clear() self._items_by_id = {} self.reset_sync_server() if not self._asset_ids: self.doc_fetched.emit() return # Collect scene container representations to compare loaded state # This runs in the main thread because it involves the host DCC if self._host: time_since_refresh = time.time() - self._host_loaded_refresh_time if time_since_refresh > self._host_loaded_refresh_timeout: if isinstance(self._host, ILoadHost): containers = self._host.get_containers() else: containers = self._host.ls() repre_ids = {con.get("representation") for con in containers} self._loaded_representation_ids = repre_ids self._host_loaded_refresh_time = time.time() self.fetch_subset_and_version() def _on_doc_fetched(self): self.clear() self._items_by_id = {} self.beginResetModel() asset_docs_by_id = self._doc_payload.get( "asset_docs_by_id" ) subset_docs_by_id = self._doc_payload.get( "subset_docs_by_id" ) last_versions_by_subset_id = self._doc_payload.get( "last_versions_by_subset_id" ) repre_info_by_version_id = self._doc_payload.get( "repre_info_by_version_id" ) subsets_loaded_by_id = self._doc_payload.get( "subsets_loaded_by_id" ) if ( asset_docs_by_id is None or subset_docs_by_id is None or last_versions_by_subset_id is None or len(self._asset_ids) == 0 ): self.endResetModel() self.refreshed.emit(False) return self._fill_subset_items( asset_docs_by_id, subset_docs_by_id, last_versions_by_subset_id, repre_info_by_version_id, subsets_loaded_by_id ) self.endResetModel() self.refreshed.emit(True) def create_multiasset_group( self, subset_name, asset_ids, subset_counter, parent_item=None ): subset_color = self.merged_subset_colors[ subset_counter % len(self.merged_subset_colors) ] merge_group = Item() merge_group.update({ "subset": "{} ({})".format(subset_name, len(asset_ids)), "isMerged": True, "subsetColor": subset_color, "assetIds": list(asset_ids), "icon": qtawesome.icon( "fa.circle", color="#{0:02x}{1:02x}{2:02x}".format(*subset_color) ) }) self.add_child(merge_group, parent_item) return merge_group def _fill_subset_items( self, asset_docs_by_id, subset_docs_by_id, last_versions_by_subset_id, repre_info_by_version_id, subsets_loaded_by_id ): _groups_tuple = self.groups_config.split_subsets_for_groups( subset_docs_by_id.values(), self._grouping ) groups, subset_docs_without_group, subset_docs_by_group = _groups_tuple group_item_by_name = {} for group_data in groups: group_name = group_data["name"] group_item = Item() group_item.update({ "subset": group_name, "isGroup": True }) group_item.update(group_data) self.add_child(group_item) group_item_by_name[group_name] = { "item": group_item, "index": self.index(group_item.row(), 0) } def _add_subset_item(subset_doc, parent_item, parent_index): last_version = last_versions_by_subset_id.get( subset_doc["_id"] ) # do not show subset without version if not last_version: return data = copy.deepcopy(subset_doc) data["subset"] = subset_doc["name"] asset_id = subset_doc["parent"] data["asset"] = asset_docs_by_id[asset_id]["name"] data["last_version"] = last_version data["loaded_in_scene"] = subset_doc["_id"] in subsets_loaded_by_id # Sync server data data.update( self._get_last_repre_info(repre_info_by_version_id, last_version["_id"])) item = Item() item.update(data) self.add_child(item, parent_item) index = self.index(item.row(), 0, parent_index) self.set_version(index, last_version) subset_counter = 0 for group_name, subset_docs_by_name in subset_docs_by_group.items(): parent_item = group_item_by_name[group_name]["item"] parent_index = group_item_by_name[group_name]["index"] for subset_name in sorted(subset_docs_by_name.keys()): subset_docs = subset_docs_by_name[subset_name] asset_ids = [ subset_doc["parent"] for subset_doc in subset_docs ] if len(subset_docs) > 1: _parent_item = self.create_multiasset_group( subset_name, asset_ids, subset_counter, parent_item ) _parent_index = self.index( _parent_item.row(), 0, parent_index ) subset_counter += 1 else: _parent_item = parent_item _parent_index = parent_index for subset_doc in subset_docs: _add_subset_item(subset_doc, parent_item=_parent_item, parent_index=_parent_index) for subset_name in sorted(subset_docs_without_group.keys()): subset_docs = subset_docs_without_group[subset_name] asset_ids = [subset_doc["parent"] for subset_doc in subset_docs] parent_item = None parent_index = None if len(subset_docs) > 1: parent_item = self.create_multiasset_group( subset_name, asset_ids, subset_counter ) parent_index = self.index(parent_item.row(), 0) subset_counter += 1 for subset_doc in subset_docs: _add_subset_item(subset_doc, parent_item=parent_item, parent_index=parent_index) def data(self, index, role): if not index.isValid(): return item = index.internalPointer() if role == ITEM_ID_ROLE: return item["id"] if role == self.SortDescendingRole: if item.get("isGroup"): # Ensure groups be on top when sorting by descending order prefix = "2" order = item["order"] else: if item.get("isMerged"): prefix = "1" else: prefix = "0" order = str(super(SubsetsModel, self).data( index, QtCore.Qt.DisplayRole )) return prefix + order if role == self.SortAscendingRole: if item.get("isGroup"): # Ensure groups be on top when sorting by ascending order prefix = "0" order = item["order"] else: if item.get("isMerged"): prefix = "1" else: prefix = "2" order = str(super(SubsetsModel, self).data( index, QtCore.Qt.DisplayRole )) return prefix + order if role == QtCore.Qt.DisplayRole: if index.column() == self.columns_index["family"]: # Show familyLabel instead of family return item.get("familyLabel", None) elif role == QtCore.Qt.DecorationRole: # Add icon to subset column if index.column() == self.columns_index["subset"]: if item.get("isGroup") or item.get("isMerged"): return item["icon"] else: return self._icons["subset"] # Add icon to family column if index.column() == self.columns_index["family"]: return item.get("familyIcon", None) elif role == QtCore.Qt.ForegroundRole: version_doc = item.get("version_document") if version_doc and version_doc.get("type") == "hero_version": if not version_doc["is_from_latest"]: return self.not_last_hero_brush elif role == LOCAL_AVAILABILITY_ROLE: if not item.get("isGroup"): return item.get("repre_info_local") else: return None elif role == REMOTE_AVAILABILITY_ROLE: if not item.get("isGroup"): return item.get("repre_info_remote") else: return None elif role == LOCAL_PROVIDER_ROLE: return self.active_provider elif role == REMOTE_PROVIDER_ROLE: return self.remote_provider return super(SubsetsModel, self).data(index, role) def flags(self, index): flags = QtCore.Qt.ItemIsEnabled | QtCore.Qt.ItemIsSelectable # Make the version column editable if index.column() == self.columns_index["version"]: flags |= QtCore.Qt.ItemIsEditable return flags def headerData(self, section, orientation, role): """Remap column names to labels""" if role == QtCore.Qt.DisplayRole: if section < len(self.Columns): key = self.Columns[section] return self.column_labels_mapping.get(key) or key super(TreeModel, self).headerData(section, orientation, role) def _get_last_repre_info(self, repre_info_by_version_id, last_version_id): data = {} if repre_info_by_version_id: repre_info = repre_info_by_version_id.get(last_version_id) return self._get_repre_dict(repre_info) return data def _get_repre_dict(self, repre_info): """Returns str representation of availability""" data = {} if repre_info: repres_str = "{}/{}".format( int(math.floor(float(repre_info['avail_repre_local']))), int(math.floor(float(repre_info['repre_count'])))) data["repre_info_local"] = repres_str repres_str = "{}/{}".format( int(math.floor(float(repre_info['avail_repre_remote']))), int(math.floor(float(repre_info['repre_count'])))) data["repre_info_remote"] = repres_str return data class GroupMemberFilterProxyModel(QtCore.QSortFilterProxyModel): """Provide the feature of filtering group by the acceptance of members The subset group nodes will not be filtered directly, the group node's acceptance depends on it's child subsets' acceptance. """ if is_filtering_recursible(): def _is_group_acceptable(self, index, node): # (NOTE) With the help of `RecursiveFiltering` feature from # Qt 5.10, group always not be accepted by default. return False filter_accepts_group = _is_group_acceptable else: # Patch future function setRecursiveFilteringEnabled = (lambda *args: None) def _is_group_acceptable(self, index, model): # (NOTE) This is not recursive. for child_row in range(model.rowCount(index)): if self.filterAcceptsRow(child_row, index): return True return False filter_accepts_group = _is_group_acceptable def __init__(self, *args, **kwargs): super(GroupMemberFilterProxyModel, self).__init__(*args, **kwargs) self.setRecursiveFilteringEnabled(True) class SubsetFilterProxyModel(GroupMemberFilterProxyModel): def filterAcceptsRow(self, row, parent): model = self.sourceModel() index = model.index(row, self.filterKeyColumn(), parent) item = index.internalPointer() if item.get("isGroup"): return self.filter_accepts_group(index, model) return super( SubsetFilterProxyModel, self ).filterAcceptsRow(row, parent) class FamiliesFilterProxyModel(GroupMemberFilterProxyModel): """Filters to specified families""" def __init__(self, *args, **kwargs): super(FamiliesFilterProxyModel, self).__init__(*args, **kwargs) self._families = set() def familyFilter(self): return self._families def setFamiliesFilter(self, values): """Set the families to include""" assert isinstance(values, (tuple, list, set)) self._families = set(values) self.invalidateFilter() def filterAcceptsRow(self, row=0, parent=None): if not self._families: return False model = self.sourceModel() index = model.index(row, 0, parent=parent or QtCore.QModelIndex()) # Ensure index is valid if not index.isValid() or index is None: return True # Get the item data and validate item = model.data(index, TreeModel.ItemRole) if item.get("isGroup"): return self.filter_accepts_group(index, model) family = item.get("family") if not family: return True # We want to keep the families which are not in the list return family in self._families def sort(self, column, order): proxy = self.sourceModel() model = proxy.sourceModel() # We need to know the sorting direction for pinning groups on top if order == QtCore.Qt.AscendingOrder: self.setSortRole(model.SortAscendingRole) else: self.setSortRole(model.SortDescendingRole) super(FamiliesFilterProxyModel, self).sort(column, order) class RepresentationSortProxyModel(GroupMemberFilterProxyModel): """To properly sort progress string""" def lessThan(self, left, right): source_model = self.sourceModel() progress_indexes = [source_model.Columns.index("active_site"), source_model.Columns.index("remote_site")] if left.column() in progress_indexes: left_data = self.sourceModel().data(left, QtCore.Qt.DisplayRole) right_data = self.sourceModel().data(right, QtCore.Qt.DisplayRole) left_val = re.sub("[^0-9]", '', left_data) right_val = re.sub("[^0-9]", '', right_data) return int(left_val) < int(right_val) return super(RepresentationSortProxyModel, self).lessThan(left, right) class RepresentationModel(TreeModel, BaseRepresentationModel): doc_fetched = QtCore.Signal() refreshed = QtCore.Signal(bool) SiteNameRole = QtCore.Qt.UserRole + 2 ProgressRole = QtCore.Qt.UserRole + 3 SiteSideRole = QtCore.Qt.UserRole + 4 IdRole = QtCore.Qt.UserRole + 5 ContextRole = QtCore.Qt.UserRole + 6 Columns = [ "name", "subset", "asset", "active_site", "remote_site" ] column_labels_mapping = { "name": "Name", "subset": "Subset", "asset": "Asset", "active_site": "Active", "remote_site": "Remote" } repre_projection = { "_id": 1, "name": 1, "context.subset": 1, "context.asset": 1, "context.version": 1, "context.representation": 1, 'files.sites': 1 } def __init__(self, dbcon, header): super(RepresentationModel, self).__init__() self.dbcon = dbcon self._data = [] self._header = header self._version_ids = [] manager = ModulesManager() sync_server = active_site = remote_site = None active_provider = remote_provider = None project_name = dbcon.current_project() if project_name: sync_server = manager.modules_by_name["sync_server"] active_site = sync_server.get_active_site(project_name) remote_site = sync_server.get_remote_site(project_name) # TODO refactor active_provider = sync_server.get_provider_for_site( project_name, active_site ) if active_site == 'studio': active_provider = 'studio' remote_provider = sync_server.get_provider_for_site( project_name, remote_site ) if remote_site == 'studio': remote_provider = 'studio' self.sync_server = sync_server self.active_site = active_site self.active_provider = active_provider self.remote_site = remote_site self.remote_provider = remote_provider self.doc_fetched.connect(self._on_doc_fetched) self._docs = {} self._icons = lib.get_repre_icons() self._icons["repre"] = qtawesome.icon( "fa.file-o", color=get_default_entity_icon_color() ) self._items_by_id = {} def set_version_ids(self, version_ids): self._version_ids = version_ids self.refresh() def data(self, index, role): item = index.internalPointer() if role == ITEM_ID_ROLE: return item["id"] if role == self.IdRole: return item.get("_id") if role == QtCore.Qt.DecorationRole: # Add icon to subset column if index.column() == self.Columns.index("name"): if item.get("isMerged"): return item["icon"] return self._icons["repre"] active_index = self.Columns.index("active_site") remote_index = self.Columns.index("remote_site") if role == QtCore.Qt.DisplayRole: progress = None label = '' if index.column() == active_index: progress = item.get("active_site_progress", 0) elif index.column() == remote_index: progress = item.get("remote_site_progress", 0) if progress is not None: # site added, sync in progress progress_str = "not avail." if progress >= 0: if progress == 0 and item.get("isMerged"): progress_str = "not avail." else: progress_str = "{}% {}".format( int(progress * 100), label ) return progress_str if role == QtCore.Qt.DecorationRole: if index.column() == active_index: return item.get("active_site_icon", None) if index.column() == remote_index: return item.get("remote_site_icon", None) if role == self.SiteNameRole: if index.column() == active_index: return item.get("active_site_name", None) if index.column() == remote_index: return item.get("remote_site_name", None) if role == self.SiteSideRole: if index.column() == active_index: return "active" if index.column() == remote_index: return "remote" if role == self.ProgressRole: if index.column() == active_index: return item.get("active_site_progress", 0) if index.column() == remote_index: return item.get("remote_site_progress", 0) return super(RepresentationModel, self).data(index, role) def _on_doc_fetched(self): self.clear() self.beginResetModel() subsets = set() assets = set() repre_groups = {} repre_groups_items = {} group = None self._items_by_id = {} for doc in self._docs: if len(self._version_ids) > 1: group = repre_groups.get(doc["name"]) if not group: group_item = Item() item_id = str(uuid4()) group_item.update({ "id": item_id, "_id": doc["_id"], "name": doc["name"], "isMerged": True, "active_site_name": self.active_site, "remote_site_name": self.remote_site, "icon": qtawesome.icon( "fa.folder", color=get_default_entity_icon_color() ) }) self._items_by_id[item_id] = group_item self.add_child(group_item, None) repre_groups[doc["name"]] = group_item repre_groups_items[doc["name"]] = 0 group = group_item progress = lib.get_progress_for_repre( doc, self.active_site, self.remote_site ) active_site_icon = self._icons.get(self.active_provider) remote_site_icon = self._icons.get(self.remote_provider) item_id = str(uuid4()) data = { "id": item_id, "_id": doc["_id"], "name": doc["name"], "subset": doc["context"]["subset"], "asset": doc["context"]["asset"], "isMerged": False, "active_site_icon": active_site_icon, "remote_site_icon": remote_site_icon, "active_site_name": self.active_site, "remote_site_name": self.remote_site, "active_site_progress": progress[self.active_site], "remote_site_progress": progress[self.remote_site] } subsets.add(doc["context"]["subset"]) assets.add(doc["context"]["subset"]) item = Item() item.update(data) self._items_by_id[item_id] = item current_progress = { 'active_site_progress': progress[self.active_site], 'remote_site_progress': progress[self.remote_site] } if group: group = self._sum_group_progress( doc["name"], group, current_progress, repre_groups_items ) self.add_child(item, group) # finalize group average progress for group_name, group in repre_groups.items(): items_cnt = repre_groups_items[group_name] active_progress = group.get("active_site_progress", 0) group["active_site_progress"] = active_progress / items_cnt remote_progress = group.get("remote_site_progress", 0) group["remote_site_progress"] = remote_progress / items_cnt self.endResetModel() self.refreshed.emit(False) def get_item_by_id(self, item_id): return self._items_by_id.get(item_id) def refresh(self): project_name = self.dbcon.current_project() if not project_name: return repre_docs = [] if self._version_ids: # Simple find here for now, expected to receive lower number of # representations and logic could be in Python repre_docs = list(get_representations( project_name, version_ids=self._version_ids, fields=self.repre_projection.keys() )) self._docs = repre_docs self.doc_fetched.emit() def _sum_group_progress( self, repre_name, group, current_item_progress, repre_groups_items ): """Update final group progress Called after every item in group is added Args: repre_name(string) group(dict): info about group of selected items current_item_progress(dict): {'active_site_progress': XX, 'remote_site_progress': YY} repre_groups_items(dict) Returns: (dict): updated group info """ repre_groups_items[repre_name] += 1 for key, progress in current_item_progress.items(): group[key] = (group.get(key, 0) + max(progress, 0)) return group