mirror of
https://github.com/ynput/ayon-core.git
synced 2026-01-01 16:34:53 +01:00
311 lines
8.5 KiB
Python
311 lines
8.5 KiB
Python
import re
|
|
import logging
|
|
import collections
|
|
|
|
from avalon.vendor.Qt import QtCore, QtWidgets
|
|
from avalon.vendor import qtawesome
|
|
from avalon import io
|
|
from avalon import style
|
|
|
|
log = logging.getLogger(__name__)
|
|
|
|
|
|
class Item(dict):
|
|
"""An item that can be represented in a tree view using `TreeModel`.
|
|
|
|
The item can store data just like a regular dictionary.
|
|
|
|
>>> data = {"name": "John", "score": 10}
|
|
>>> item = Item(data)
|
|
>>> assert item["name"] == "John"
|
|
|
|
"""
|
|
|
|
def __init__(self, data=None):
|
|
super(Item, self).__init__()
|
|
|
|
self._children = list()
|
|
self._parent = None
|
|
|
|
if data is not None:
|
|
assert isinstance(data, dict)
|
|
self.update(data)
|
|
|
|
def childCount(self):
|
|
return len(self._children)
|
|
|
|
def child(self, row):
|
|
|
|
if row >= len(self._children):
|
|
log.warning("Invalid row as child: {0}".format(row))
|
|
return
|
|
|
|
return self._children[row]
|
|
|
|
def children(self):
|
|
return self._children
|
|
|
|
def parent(self):
|
|
return self._parent
|
|
|
|
def row(self):
|
|
"""
|
|
Returns:
|
|
int: Index of this item under parent"""
|
|
if self._parent is not None:
|
|
siblings = self.parent().children()
|
|
return siblings.index(self)
|
|
|
|
def add_child(self, child):
|
|
"""Add a child to this item"""
|
|
child._parent = self
|
|
self._children.append(child)
|
|
|
|
|
|
class TreeModel(QtCore.QAbstractItemModel):
|
|
|
|
Columns = list()
|
|
ItemRole = QtCore.Qt.UserRole + 1
|
|
|
|
def __init__(self, parent=None):
|
|
super(TreeModel, self).__init__(parent)
|
|
self._root_item = Item()
|
|
|
|
def rowCount(self, parent):
|
|
if parent.isValid():
|
|
item = parent.internalPointer()
|
|
else:
|
|
item = self._root_item
|
|
|
|
return item.childCount()
|
|
|
|
def columnCount(self, parent):
|
|
return len(self.Columns)
|
|
|
|
def data(self, index, role):
|
|
|
|
if not index.isValid():
|
|
return None
|
|
|
|
if role == QtCore.Qt.DisplayRole or role == QtCore.Qt.EditRole:
|
|
|
|
item = index.internalPointer()
|
|
column = index.column()
|
|
|
|
key = self.Columns[column]
|
|
return item.get(key, None)
|
|
|
|
if role == self.ItemRole:
|
|
return index.internalPointer()
|
|
|
|
def setData(self, index, value, role=QtCore.Qt.EditRole):
|
|
"""Change the data on the items.
|
|
|
|
Returns:
|
|
bool: Whether the edit was successful
|
|
"""
|
|
|
|
if index.isValid():
|
|
if role == QtCore.Qt.EditRole:
|
|
|
|
item = index.internalPointer()
|
|
column = index.column()
|
|
key = self.Columns[column]
|
|
item[key] = value
|
|
|
|
# passing `list()` for PyQt5 (see PYSIDE-462)
|
|
self.dataChanged.emit(index, index, list())
|
|
|
|
# must return true if successful
|
|
return True
|
|
|
|
return False
|
|
|
|
def setColumns(self, keys):
|
|
assert isinstance(keys, (list, tuple))
|
|
self.Columns = keys
|
|
|
|
def headerData(self, section, orientation, role):
|
|
|
|
if role == QtCore.Qt.DisplayRole:
|
|
if section < len(self.Columns):
|
|
return self.Columns[section]
|
|
|
|
super(TreeModel, self).headerData(section, orientation, role)
|
|
|
|
def flags(self, index):
|
|
flags = QtCore.Qt.ItemIsEnabled
|
|
|
|
item = index.internalPointer()
|
|
if item.get("enabled", True):
|
|
flags |= QtCore.Qt.ItemIsSelectable
|
|
|
|
return flags
|
|
|
|
def parent(self, index):
|
|
|
|
item = index.internalPointer()
|
|
parent_item = item.parent()
|
|
|
|
# If it has no parents we return invalid
|
|
if parent_item == self._root_item or not parent_item:
|
|
return QtCore.QModelIndex()
|
|
|
|
return self.createIndex(parent_item.row(), 0, parent_item)
|
|
|
|
def index(self, row, column, parent):
|
|
"""Return index for row/column under parent"""
|
|
|
|
if not parent.isValid():
|
|
parent_item = self._root_item
|
|
else:
|
|
parent_item = parent.internalPointer()
|
|
|
|
child_item = parent_item.child(row)
|
|
if child_item:
|
|
return self.createIndex(row, column, child_item)
|
|
else:
|
|
return QtCore.QModelIndex()
|
|
|
|
def add_child(self, item, parent=None):
|
|
if parent is None:
|
|
parent = self._root_item
|
|
|
|
parent.add_child(item)
|
|
|
|
def column_name(self, column):
|
|
"""Return column key by index"""
|
|
|
|
if column < len(self.Columns):
|
|
return self.Columns[column]
|
|
|
|
def clear(self):
|
|
self.beginResetModel()
|
|
self._root_item = Item()
|
|
self.endResetModel()
|
|
|
|
|
|
class TasksModel(TreeModel):
|
|
"""A model listing the tasks combined for a list of assets"""
|
|
|
|
Columns = ["Tasks"]
|
|
|
|
def __init__(self):
|
|
super(TasksModel, self).__init__()
|
|
self._num_assets = 0
|
|
self._icons = {
|
|
"__default__": qtawesome.icon("fa.male",
|
|
color=style.colors.default),
|
|
"__no_task__": qtawesome.icon("fa.exclamation-circle",
|
|
color=style.colors.mid)
|
|
}
|
|
|
|
self._get_task_icons()
|
|
|
|
def _get_task_icons(self):
|
|
# Get the project configured icons from database
|
|
project = io.find_one({"type": "project"})
|
|
tasks = project["config"].get("tasks", [])
|
|
for task in tasks:
|
|
icon_name = task.get("icon", None)
|
|
if icon_name:
|
|
icon = qtawesome.icon("fa.{}".format(icon_name),
|
|
color=style.colors.default)
|
|
self._icons[task["name"]] = icon
|
|
|
|
def set_tasks(self, tasks):
|
|
"""Set assets to track by their database id
|
|
|
|
Arguments:
|
|
asset_ids (list): List of asset ids.
|
|
|
|
"""
|
|
|
|
self.clear()
|
|
|
|
# let cleared task view if no tasks are available
|
|
if len(tasks) == 0:
|
|
return
|
|
|
|
self.beginResetModel()
|
|
|
|
icon = self._icons["__default__"]
|
|
for task in tasks:
|
|
item = Item({
|
|
"Tasks": task,
|
|
"icon": icon
|
|
})
|
|
|
|
self.add_child(item)
|
|
|
|
self.endResetModel()
|
|
|
|
def flags(self, index):
|
|
return QtCore.Qt.ItemIsEnabled | QtCore.Qt.ItemIsSelectable
|
|
|
|
def headerData(self, section, orientation, role):
|
|
|
|
# Override header for count column to show amount of assets
|
|
# it is listing the tasks for
|
|
if role == QtCore.Qt.DisplayRole:
|
|
if orientation == QtCore.Qt.Horizontal:
|
|
if section == 1: # count column
|
|
return "count ({0})".format(self._num_assets)
|
|
|
|
return super(TasksModel, self).headerData(section, orientation, role)
|
|
|
|
def data(self, index, role):
|
|
|
|
if not index.isValid():
|
|
return
|
|
|
|
# Add icon to the first column
|
|
if role == QtCore.Qt.DecorationRole:
|
|
if index.column() == 0:
|
|
return index.internalPointer()["icon"]
|
|
|
|
return super(TasksModel, self).data(index, role)
|
|
|
|
|
|
class DeselectableTreeView(QtWidgets.QTreeView):
|
|
"""A tree view that deselects on clicking on an empty area in the view"""
|
|
|
|
def mousePressEvent(self, event):
|
|
|
|
index = self.indexAt(event.pos())
|
|
if not index.isValid():
|
|
# clear the selection
|
|
self.clearSelection()
|
|
# clear the current index
|
|
self.setCurrentIndex(QtCore.QModelIndex())
|
|
|
|
QtWidgets.QTreeView.mousePressEvent(self, event)
|
|
|
|
|
|
class RecursiveSortFilterProxyModel(QtCore.QSortFilterProxyModel):
|
|
"""Filters to the regex if any of the children matches allow parent"""
|
|
def filterAcceptsRow(self, row, parent):
|
|
|
|
regex = self.filterRegExp()
|
|
if not regex.isEmpty():
|
|
pattern = regex.pattern()
|
|
model = self.sourceModel()
|
|
source_index = model.index(row, self.filterKeyColumn(), parent)
|
|
if source_index.isValid():
|
|
|
|
# Check current index itself
|
|
key = model.data(source_index, self.filterRole())
|
|
if re.search(pattern, key, re.IGNORECASE):
|
|
return True
|
|
|
|
# Check children
|
|
rows = model.rowCount(source_index)
|
|
for i in range(rows):
|
|
if self.filterAcceptsRow(i, source_index):
|
|
return True
|
|
|
|
# Otherwise filter it
|
|
return False
|
|
|
|
return super(RecursiveSortFilterProxyModel,
|
|
self).filterAcceptsRow(row, parent)
|