added widgets for asset selection

pype/tools/standalonepublish/widgets/__init__.py

@@ -9,3 +9,13 @@ PluginRole = QtCore.Qt.UserRole + 5
 
 from ..resources import get_resource
 from .button_from_svgs import SvgResizable, SvgButton
+
+from .model_node import Node
+from .model_tree import TreeModel
+from .model_asset import AssetModel
+from .model_filter_proxy_exact_match import ExactMatchesFilterProxyModel
+from .model_filter_proxy_recursive_sort import RecursiveSortFilterProxyModel
+from .model_tree_view_deselectable import DeselectableTreeView
+
+from .widget_asset_view import AssetView
+from .widget_asset import AssetWidget

pype/tools/standalonepublish/widgets/model_asset.py

@@ -0,0 +1,158 @@
+import logging
+from . import QtCore, QtGui
+from . import TreeModel, Node
+from . import style, awesome
+
+
+log = logging.getLogger(__name__)
+
+
+def _iter_model_rows(model,
+                     column,
+                     include_root=False):
+    """Iterate over all row indices in a model"""
+    indices = [QtCore.QModelIndex()]  # start iteration at root
+
+    for index in indices:
+
+        # Add children to the iterations
+        child_rows = model.rowCount(index)
+        for child_row in range(child_rows):
+            child_index = model.index(child_row, column, index)
+            indices.append(child_index)
+
+        if not include_root and not index.isValid():
+            continue
+
+        yield index
+
+
+class AssetModel(TreeModel):
+    """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.
+
+    """
+
+    COLUMNS = ["label"]
+    Name = 0
+    Deprecated = 2
+    ObjectId = 3
+
+    DocumentRole = QtCore.Qt.UserRole + 2
+    ObjectIdRole = QtCore.Qt.UserRole + 3
+
+    def __init__(self, parent):
+        super(AssetModel, self).__init__(parent=parent)
+        self.parent_widget = parent
+        self.refresh()
+
+    @property
+    def db(self):
+        return self.parent_widget.db
+
+    def _add_hierarchy(self, parent=None):
+
+        # Find the assets under the parent
+        find_data = {
+            "type": "asset"
+        }
+        if parent is None:
+            find_data['$or'] = [
+                {'data.visualParent': {'$exists': False}},
+                {'data.visualParent': None}
+            ]
+        else:
+            find_data["data.visualParent"] = parent['_id']
+
+        assets = self.db.find(find_data).sort('name', 1)
+        for asset in assets:
+            # get label from data, otherwise use name
+            data = asset.get("data", {})
+            label = data.get("label", asset['name'])
+            tags = data.get("tags", [])
+
+            # store for the asset for optimization
+            deprecated = "deprecated" in tags
+
+            node = Node({
+                "_id": asset['_id'],
+                "name": asset["name"],
+                "label": label,
+                "type": asset['type'],
+                "tags": ", ".join(tags),
+                "deprecated": deprecated,
+                "_document": asset
+            })
+            self.add_child(node, parent=parent)
+
+            # Add asset's children recursively
+            self._add_hierarchy(node)
+
+    def refresh(self):
+        """Refresh the data for the model."""
+
+        self.clear()
+        if (
+            self.db.active_project() is None or
+            self.db.active_project() == ''
+        ):
+            return
+        self.beginResetModel()
+        self._add_hierarchy(parent=None)
+        self.endResetModel()
+
+    def flags(self, index):
+        return QtCore.Qt.ItemIsEnabled | QtCore.Qt.ItemIsSelectable
+
+    def data(self, index, role):
+
+        if not index.isValid():
+            return
+
+        node = index.internalPointer()
+        if role == QtCore.Qt.DecorationRole:        # icon
+
+            column = index.column()
+            if column == self.Name:
+
+                # Allow a custom icon and custom icon color to be defined
+                data = node["_document"]["data"]
+                icon = data.get("icon", None)
+                color = data.get("color", style.colors.default)
+
+                if icon is None:
+                    # Use default icons if no custom one is specified.
+                    # If it has children show a full folder, otherwise
+                    # show an open folder
+                    has_children = self.rowCount(index) > 0
+                    icon = "folder" if has_children else "folder-o"
+
+                # Make the color darker when the asset is deprecated
+                if node.get("deprecated", False):
+                    color = QtGui.QColor(color).darker(250)
+
+                try:
+                    key = "fa.{0}".format(icon)  # font-awesome key
+                    icon = awesome.icon(key, color=color)
+                    return icon
+                except Exception as exception:
+                    # Log an error message instead of erroring out completely
+                    # when the icon couldn't be created (e.g. invalid name)
+                    log.error(exception)
+
+                return
+
+        if role == QtCore.Qt.ForegroundRole:        # font color
+            if "deprecated" in node.get("tags", []):
+                return QtGui.QColor(style.colors.light).darker(250)
+
+        if role == self.ObjectIdRole:
+            return node.get("_id", None)
+
+        if role == self.DocumentRole:
+            return node.get("_document", None)
+
+        return super(AssetModel, self).data(index, role)

pype/tools/standalonepublish/widgets/model_filter_proxy_exact_match.py

@@ -0,0 +1,28 @@
+from . import QtCore
+
+
+class ExactMatchesFilterProxyModel(QtCore.QSortFilterProxyModel):
+    """Filter model to where key column's value is in the filtered tags"""
+
+    def __init__(self, *args, **kwargs):
+        super(ExactMatchesFilterProxyModel, self).__init__(*args, **kwargs)
+        self._filters = set()
+
+    def setFilters(self, filters):
+        self._filters = set(filters)
+
+    def filterAcceptsRow(self, source_row, source_parent):
+
+        # No filter
+        if not self._filters:
+            return True
+
+        else:
+            model = self.sourceModel()
+            column = self.filterKeyColumn()
+            idx = model.index(source_row, column, source_parent)
+            data = model.data(idx, self.filterRole())
+            if data in self._filters:
+                return True
+            else:
+                return False

pype/tools/standalonepublish/widgets/model_filter_proxy_recursive_sort.py

@@ -0,0 +1,30 @@
+from . import QtCore
+
+
+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)

pype/tools/standalonepublish/widgets/model_node.py

@@ -0,0 +1,56 @@
+import logging
+
+
+log = logging.getLogger(__name__)
+
+
+class Node(dict):
+    """A node that can be represented in a tree view.
+
+    The node can store data just like a dictionary.
+
+    >>> data = {"name": "John", "score": 10}
+    >>> node = Node(data)
+    >>> assert node["name"] == "John"
+
+    """
+
+    def __init__(self, data=None):
+        super(Node, 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 node 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 node"""
+        child._parent = self
+        self._children.append(child)

pype/tools/standalonepublish/widgets/model_tree.py

@@ -0,0 +1,122 @@
+from . import QtCore
+from . import Node
+
+
+class TreeModel(QtCore.QAbstractItemModel):
+
+    COLUMNS = list()
+    NodeRole = QtCore.Qt.UserRole + 1
+
+    def __init__(self, parent=None):
+        super(TreeModel, self).__init__(parent)
+        self._root_node = Node()
+
+    def rowCount(self, parent):
+        if parent.isValid():
+            node = parent.internalPointer()
+        else:
+            node = self._root_node
+
+        return node.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:
+
+            node = index.internalPointer()
+            column = index.column()
+
+            key = self.COLUMNS[column]
+            return node.get(key, None)
+
+        if role == self.NodeRole:
+            return index.internalPointer()
+
+    def setData(self, index, value, role=QtCore.Qt.EditRole):
+        """Change the data on the nodes.
+
+        Returns:
+            bool: Whether the edit was successful
+        """
+
+        if index.isValid():
+            if role == QtCore.Qt.EditRole:
+
+                node = index.internalPointer()
+                column = index.column()
+                key = self.COLUMNS[column]
+                node[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):
+        return (
+            QtCore.Qt.ItemIsEnabled |
+            QtCore.Qt.ItemIsSelectable
+        )
+
+    def parent(self, index):
+
+        node = index.internalPointer()
+        parent_node = node.parent()
+
+        # If it has no parents we return invalid
+        if parent_node == self._root_node or not parent_node:
+            return QtCore.QModelIndex()
+
+        return self.createIndex(parent_node.row(), 0, parent_node)
+
+    def index(self, row, column, parent):
+        """Return index for row/column under parent"""
+
+        if not parent.isValid():
+            parentNode = self._root_node
+        else:
+            parentNode = parent.internalPointer()
+
+        childItem = parentNode.child(row)
+        if childItem:
+            return self.createIndex(row, column, childItem)
+        else:
+            return QtCore.QModelIndex()
+
+    def add_child(self, node, parent=None):
+        if parent is None:
+            parent = self._root_node
+
+        parent.add_child(node)
+
+    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_node = Node()
+        self.endResetModel()

pype/tools/standalonepublish/widgets/model_tree_view_deselectable.py

@@ -0,0 +1,16 @@
+from . import QtWidgets, QtCore
+
+
+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)

pype/tools/standalonepublish/widgets/widget_asset.py

@@ -0,0 +1,255 @@
+import contextlib
+from . import QtWidgets, QtCore
+from . import RecursiveSortFilterProxyModel, AssetModel, AssetView
+from . import awesome, style
+
+@contextlib.contextmanager
+def preserve_expanded_rows(tree_view,
+                           column=0,
+                           role=QtCore.Qt.DisplayRole):
+    """Preserves expanded row in QTreeView by column's data role.
+
+    This function is created to maintain the expand vs collapse status of
+    the model items. When refresh is triggered the items which are expanded
+    will stay expanded and vise versa.
+
+    Arguments:
+        tree_view (QWidgets.QTreeView): the tree view which is
+            nested in the application
+        column (int): the column to retrieve the data from
+        role (int): the role which dictates what will be returned
+
+    Returns:
+        None
+
+    """
+
+    model = tree_view.model()
+
+    expanded = set()
+
+    for index in _iter_model_rows(model,
+                                  column=column,
+                                  include_root=False):
+        if tree_view.isExpanded(index):
+            value = index.data(role)
+            expanded.add(value)
+
+    try:
+        yield
+    finally:
+        if not expanded:
+            return
+
+        for index in _iter_model_rows(model,
+                                      column=column,
+                                      include_root=False):
+            value = index.data(role)
+            state = value in expanded
+            if state:
+                tree_view.expand(index)
+            else:
+                tree_view.collapse(index)
+
+
+@contextlib.contextmanager
+def preserve_selection(tree_view,
+                       column=0,
+                       role=QtCore.Qt.DisplayRole,
+                       current_index=True):
+    """Preserves row selection in QTreeView by column's data role.
+
+    This function is created to maintain the selection status of
+    the model items. When refresh is triggered the items which are expanded
+    will stay expanded and vise versa.
+
+        tree_view (QWidgets.QTreeView): the tree view nested in the application
+        column (int): the column to retrieve the data from
+        role (int): the role which dictates what will be returned
+
+    Returns:
+        None
+
+    """
+
+    model = tree_view.model()
+    selection_model = tree_view.selectionModel()
+    flags = selection_model.Select | selection_model.Rows
+
+    if current_index:
+        current_index_value = tree_view.currentIndex().data(role)
+    else:
+        current_index_value = None
+
+    selected_rows = selection_model.selectedRows()
+    if not selected_rows:
+        yield
+        return
+
+    selected = set(row.data(role) for row in selected_rows)
+    try:
+        yield
+    finally:
+        if not selected:
+            return
+
+        # Go through all indices, select the ones with similar data
+        for index in _iter_model_rows(model,
+                                      column=column,
+                                      include_root=False):
+
+            value = index.data(role)
+            state = value in selected
+            if state:
+                tree_view.scrollTo(index)  # Ensure item is visible
+                selection_model.select(index, flags)
+
+            if current_index_value and value == current_index_value:
+                tree_view.setCurrentIndex(index)
+
+
+class AssetWidget(QtWidgets.QWidget):
+    """A Widget to display a tree of assets with filter
+
+    To list the assets of the active project:
+        >>> # widget = AssetWidget()
+        >>> # widget.refresh()
+        >>> # widget.show()
+
+    """
+
+    assets_refreshed = QtCore.Signal()   # on model refresh
+    selection_changed = QtCore.Signal()  # on view selection change
+    current_changed = QtCore.Signal()    # on view current index change
+
+    def __init__(self, parent):
+        super(AssetWidget, self).__init__(parent=parent)
+        self.setContentsMargins(0, 0, 0, 0)
+
+        self.parent_widget = parent
+
+        layout = QtWidgets.QVBoxLayout(self)
+        layout.setContentsMargins(0, 0, 0, 0)
+        layout.setSpacing(4)
+
+        # Project
+        self.combo_projects = QtWidgets.QComboBox()
+        self._set_projects()
+        self.combo_projects.currentTextChanged.connect(self.on_project_change)
+        # Tree View
+        model = AssetModel(self)
+        proxy = RecursiveSortFilterProxyModel()
+        proxy.setSourceModel(model)
+        proxy.setFilterCaseSensitivity(QtCore.Qt.CaseInsensitive)
+        view = AssetView()
+        view.setModel(proxy)
+
+        # Header
+        header = QtWidgets.QHBoxLayout()
+
+        icon = awesome.icon("fa.refresh", color=style.colors.light)
+        refresh = QtWidgets.QPushButton(icon, "")
+        refresh.setToolTip("Refresh items")
+
+        filter = QtWidgets.QLineEdit()
+        filter.textChanged.connect(proxy.setFilterFixedString)
+        filter.setPlaceholderText("Filter assets..")
+
+        header.addWidget(filter)
+        header.addWidget(refresh)
+
+        # Layout
+        layout.addWidget(self.combo_projects)
+        layout.addLayout(header)
+        layout.addWidget(view)
+
+        # Signals/Slots
+        selection = view.selectionModel()
+        selection.selectionChanged.connect(self.selection_changed)
+        selection.currentChanged.connect(self.current_changed)
+        refresh.clicked.connect(self.refresh)
+
+        self.refreshButton = refresh
+        self.model = model
+        self.proxy = proxy
+        self.view = view
+
+    @property
+    def db(self):
+        return self.parent_widget.db
+
+    def _set_projects(self):
+        projects = list()
+        for project in self.db.projects():
+            projects.append(project['name'])
+
+        self.combo_projects.clear()
+        if len(projects) > 0:
+            self.combo_projects.addItems(projects)
+            self.db.activate_project(projects[0])
+
+    def on_project_change(self):
+        projects = list()
+        for project in self.db.projects():
+            projects.append(project['name'])
+        project_name = self.combo_projects.currentText()
+        if project_name in projects:
+            self.db.activate_project(project_name)
+        self.refresh()
+
+    def _refresh_model(self):
+        self.model.refresh()
+        self.assets_refreshed.emit()
+
+    def refresh(self):
+        self._refresh_model()
+
+    def get_active_asset(self):
+        """Return the asset id the current asset."""
+        current = self.view.currentIndex()
+        return current.data(self.model.ObjectIdRole)
+
+    def get_active_index(self):
+        return self.view.currentIndex()
+
+    def get_selected_assets(self):
+        """Return the assets' ids that are selected."""
+        selection = self.view.selectionModel()
+        rows = selection.selectedRows()
+        return [row.data(self.model.ObjectIdRole) for row in rows]
+
+    def select_assets(self, assets, expand=True):
+        """Select assets by name.
+
+        Args:
+            assets (list): List of asset names
+            expand (bool): Whether to also expand to the asset in the view
+
+        Returns:
+            None
+
+        """
+        # TODO: Instead of individual selection optimize for many assets
+
+        assert isinstance(assets,
+                          (tuple, list)), "Assets must be list or tuple"
+
+        # Clear selection
+        selection_model = self.view.selectionModel()
+        selection_model.clearSelection()
+
+        # Select
+        mode = selection_model.Select | selection_model.Rows
+        for index in _iter_model_rows(self.proxy,
+                                      column=0,
+                                      include_root=False):
+            data = index.data(self.model.NodeRole)
+            name = data['name']
+            if name in assets:
+                selection_model.select(index, mode)
+
+                if expand:
+                    self.view.expand(index)
+
+                # Set the currently active index
+                self.view.setCurrentIndex(index)

pype/tools/standalonepublish/widgets/widget_asset_view.py

@@ -0,0 +1,16 @@
+from . import QtCore
+from . import DeselectableTreeView
+
+
+class AssetView(DeselectableTreeView):
+    """Item view.
+
+    This implements a context menu.
+
+    """
+
+    def __init__(self):
+        super(AssetView, self).__init__()
+        self.setIndentation(15)
+        self.setContextMenuPolicy(QtCore.Qt.CustomContextMenu)
+        self.setHeaderHidden(True)