ayon-core/openpype/tools/loader/model.py
2023-03-29 23:57:35 +02:00

1269 lines
42 KiB
Python

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