Quick 'n' dirty prototype for Input/Outputs tracking UI for Avalon

avalon_show_loaded_containers_inputs.py

from Qt import QtWidgets, QtCore

import avalon.api
import avalon.io
import avalon.style as style
import avalon.tools.representationinput


def get_main_window():
    """Return top level QMainWindow"""
    top_widgets = QtWidgets.QApplication.topLevelWidgets()
    main_window = next((widget for widget in top_widgets if
                       widget.inherits("QMainWindow")), None)
    return main_window
    
    
class ContainerInputDialog(QtWidgets.QDialog):
    
    ItemRole = QtCore.Qt.UserRole + 1
    
    def __init__(self, parent=None):
        super(ContainerInputDialog, self).__init__(parent=parent)
        
        representationinput = avalon.tools.representationinput.RepresentationWidget()
        
        refresh_containers = QtWidgets.QPushButton("Refresh Containers")
        containers = QtWidgets.QListWidget()
        containers.setSelectionMode(containers.ExtendedSelection)
        refresh_representations = QtWidgets.QPushButton("Refresh Inputs/Outputs")
        
        left = QtWidgets.QVBoxLayout()
        left.addWidget(refresh_containers)
        left.addWidget(containers)
        left.addWidget(refresh_representations)
        refresh_containers.setFixedWidth(250)
        containers.setFixedWidth(250)
        refresh_representations.setFixedWidth(250)
        
        hbox = QtWidgets.QHBoxLayout(self)
        hbox.addLayout(left)
        hbox.addWidget(representationinput)
        
        refresh_containers.clicked.connect(self.refresh_containers)
        refresh_representations.clicked.connect(self.refresh_representations)
        
        self.representationinput = representationinput
        self.containers = containers

    def refresh_containers(self):

        host = avalon.api.registered_host()
        if not host:
            raise RuntimeError("No registered avalon host. Is Avalon installed correctly?")
        
        containers = list(host.ls())
        self.containers.clear()
        for container in sorted(containers, key=lambda x: x["name"]):
            item = QtWidgets.QListWidgetItem()
            label = "{0[name]} ({0[loader]})".format(container)
            item.setData(QtCore.Qt.DisplayRole, label)
            item.setData(self.ItemRole, container)
            
            self.containers.addItem(item)
        
    def refresh_representations(self):
        
        containers = [item.data(self.ItemRole) for item in self.containers.selectedItems()]
        representations = [avalon.io.ObjectId(c["representation"]) for c in containers]
        
        self.representationinput.model.load(representations)
        self.representationinput.on_mode_changed()


parent = get_main_window()

widget = ContainerInputDialog(parent=parent)
widget.setWindowTitle("List inputs/outputs for scene containers")
widget.resize(1200, 600)
widget.show()
widget.setStyleSheet(style.load_stylesheet())
widget.refresh_containers()

representationinput.py

from collections import defaultdict

from .. import io, style, api
from ..vendor.Qt import QtCore, QtGui, QtWidgets
from ..vendor import qtawesome as qta

from .models import (
    TreeModel,
    Item,
    RecursiveSortFilterProxyModel
)
from . import lib
from .delegates import PrettyTimeDelegate, VersionDelegate


@io.auto_reconnect
def aggregate(*args, **kwargs):
    return io._database[api.Session["AVALON_PROJECT"]].aggregate(
        *args, **kwargs)


def get_representations_inputs_recursive(representations,
                                         max_depth=0):
    """Returns representations with extra 'inputs_recursive' key"""

    assert all(isinstance(x, io.ObjectId) for x in representations)

    graph_lookup = {
        "from": api.Session["AVALON_PROJECT"],
        "startWith": "$data.inputs",
        "connectFromField": "data.inputs",
        "connectToField": "_id",
        "as": "inputs_recursive",
        "depthField": "depth"
    }
    if max_depth:
        graph_lookup["maxDepth"] = max_depth

    pipeline_ = [
        # Match
        {"$match": {"_id": {"$in": representations},
                    "type": "representation"}},
        # Recursive graph lookup for inputs
        {"$graphLookup": graph_lookup}
    ]

    result = aggregate(pipeline_)

    return list(result)


def get_representations_outputs_recursive(representations,
                                          max_depth=0):
    """Returns representations with extra 'outputs_recursive' key.

    This will also embed an extra "data.outputs" value on all of the
    representations so we can easily reiterate the tree in that direction.

    """

    assert all(isinstance(x, io.ObjectId) for x in representations)

    graph_lookup = {
        "from": api.Session["AVALON_PROJECT"],
        "startWith": "$_id",
        "connectFromField": "_id",
        "connectToField": "data.inputs",
        "as": "outputs_recursive",
        "depthField": "depth"
    }
    if max_depth:
        graph_lookup["maxDepth"] = max_depth

    pipeline_ = [
        # Match
        {"$match": {"_id": {"$in": representations},
                    "type": "representation"}},
        # Recursive graph lookup for inputs
        {"$graphLookup": graph_lookup}
    ]

    result = aggregate(pipeline_)
    result = list(result)

    for representation in result:
        outputs_by_id = defaultdict(list)
        representation["data"]["outputs"] = []

        # Build a lookup for the outputs per _id
        for output in representation["outputs_recursive"]:
            for input_ in output["data"]["inputs"]:
                outputs_by_id[input_].append(output["_id"])

            if output["depth"] == 0:
                # Direct output for the representation, ignore for lookup
                representation["data"]["outputs"].append(output["_id"])
                continue

        # Add the relevant outputs to the recursively collected output
        # dictionaries
        for output in representation["outputs_recursive"]:
            output["data"]["outputs"] = outputs_by_id.get(output["_id"], [])

    return result


def get_families(version):
    """Return families from version document"""
    if version.get("schema") == "avalon-core:version-3.0":
        subset = io.find_one({"_id": version["parent"], "type": "subset"})
        return subset["data"].get("families", None)
    else:
        # Older version schema
        return version["data"].get("families", None)


class RepresentationInputOutputModel(TreeModel):
    """A model listing the inputs/outputs for representations"""

    Columns = ["label", "family", "version", "time", "author"]

    # Different modes for listing
    INPUTS = 0
    OUTPUTS = 1

    def __init__(self):
        super(RepresentationInputOutputModel, self).__init__()

        # Default mode is listing inputs
        self.mode = self.INPUTS
        self._representations = []

        self._icons = {"subset": qta.icon("fa.file-o",
                                          color=style.colors.default)}

    def load(self, representation_ids):
        """Set representation to track by their database id.

        Arguments:
            representation_ids (list): List of representation ids.

        """
        assert isinstance(representation_ids, (list, tuple))
        self._representations = representation_ids
        self.refresh()

    def set_mode(self, mode):
        self.mode = mode

    def refresh(self):

        self.clear()
        self.beginResetModel()

        representation_ids = self._representations

        # The recursive aggregated queries have collected all dependencies
        # in full so we don't need to perform additional queries. We will
        # use that as a cached lookup for the actual representations
        cache = {}
        if self.mode == self.INPUTS:
            result = get_representations_inputs_recursive(representation_ids)
            for representation in result:
                cache[representation["_id"]] = representation
                cache.update({i["_id"]: i for i
                              in representation["inputs_recursive"]})

        elif self.mode == self.OUTPUTS:
            result = get_representations_outputs_recursive(representation_ids)
            for representation in result:
                cache[representation["_id"]] = representation
                cache.update({i["_id"]: i for i
                              in representation["outputs_recursive"]})

        for representation_id in representation_ids:
            # Populate the representations
            self._add_representation(representation_id,
                                     parent=None,
                                     cache=cache)

        self.endResetModel()

    def _add_representation(self, representation_id, parent, cache):

        representation = cache.get(representation_id)
        context = representation["context"]

        # Get the version since we want to display some information from it
        # like the specific time/date it was published.
        version = io.find_one({"_id": representation["parent"],
                               "type": "version"})
        version_data = version.get("data", dict())

        # Get family from version
        families = get_families(version)
        family = families[0] if families else "???"
        family_config = lib.get_family_cached_config(family)

        label = "{subset} .{representation} ({asset})".format(
            **context
        )

        node = Item({
            "label": label,
            "version": context["version"],

            # Version data
            "time": version_data.get("time", None),
            "author": version_data.get("author", None),
            "family": family_config.get("label", family),
            "familyIcon": family_config.get('icon', None),
            "families": families
        })

        # Collect its dependencies and recursively populate them too
        if self.mode == self.INPUTS:
            for input_id in representation["data"].get("inputs", []):
                self._add_representation(input_id, parent=node, cache=cache)
        elif self.mode == self.OUTPUTS:
            for output_id in representation["data"].get("outputs", []):
                self._add_representation(output_id, parent=node, cache=cache)
        else:
            raise ValueError("Invalid mode: %s" % self.mode)

        self.add_child(node, parent=parent)

    def data(self, index, role):

        if not index.isValid():
            return

        # Show icons
        if role == QtCore.Qt.DecorationRole:
            if index.column() == 0:
                return self._icons["subset"]

            if index.column() == 1:
                return index.internalPointer()['familyIcon']

        return super(RepresentationInputOutputModel, self).data(index, role)


class RecursiveFamilyFilter(RecursiveSortFilterProxyModel):
    """Filters to specified families"""

    def __init__(self, *args, **kwargs):
        super(RecursiveFamilyFilter, 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 is_item_family_allowed(self, row, parent):

        model = self.sourceModel()
        index = model.index(row, 0, parent=parent)

        # Ensure index is valid
        if not index.isValid() or index is None:
            return True

        # Get the node data and validate
        item = model.data(index, TreeModel.ItemRole)

        families = item.get("families", [])
        filterable_families = set()
        for name in families:
            family_config = lib.get_family_cached_config(name)
            if not family_config.get("hideFilter"):
                filterable_families.add(name)

        if not filterable_families:
            return True

        # Check children
        rows = model.rowCount(index)
        for i in range(rows):
            if self.is_item_family_allowed(i, index):
                return True

        return filterable_families.issubset(self._families)

    def filterAcceptsRow(self, row=0, parent=QtCore.QModelIndex()):
        if not self._families:
            return False

        return self.is_item_family_allowed(row, parent) and \
            super(RecursiveFamilyFilter, self).filterAcceptsRow(row, parent)


class RepresentationWidget(QtWidgets.QDialog):
    """A Widget showing the inputs/outputs for a Representation."""

    def __init__(self, parent=None):
        super(RepresentationWidget, self).__init__(parent=parent)

        mode_checkboxes = QtWidgets.QGroupBox()
        mode_checkboxes.setTitle("Inputs/Outputs")
        hbox = QtWidgets.QHBoxLayout(mode_checkboxes)
        hbox.setContentsMargins(0, 0, 0, 0)
        input = QtWidgets.QRadioButton("Inputs")
        input.setChecked(True)
        output = QtWidgets.QRadioButton("Outputs")
        hbox.addWidget(input)
        hbox.addWidget(output)
        hbox.addStretch()

        from .loader.widgets import FamilyListWidget

        search = QtWidgets.QLineEdit()
        search.setPlaceholderText("Search..")

        model = RepresentationInputOutputModel()

        proxy = RecursiveFamilyFilter()
        proxy.setDynamicSortFilter(True)
        proxy.setFilterCaseSensitivity(QtCore.Qt.CaseInsensitive)

        view = QtWidgets.QTreeView()
        view.setIndentation(15)
        view.setStyleSheet("""
           QTreeView::item{
               padding: 5px 1px;
               border: 0px;
           }
        """)
        view.setAllColumnsShowFocus(True)

        families_label = QtWidgets.QLabel("Filter by family")
        families_label.setStyleSheet("padding: 5px;")
        families = FamilyListWidget()
        families.setFixedHeight(200)

        # Set view delegates
        version_delegate = VersionDelegate()
        column = model.Columns.index("version")
        view.setItemDelegateForColumn(column, version_delegate)

        time_delegate = PrettyTimeDelegate()
        column = model.Columns.index("time")
        view.setItemDelegateForColumn(column, time_delegate)

        layout = QtWidgets.QVBoxLayout(self)
        layout.setContentsMargins(0, 0, 0, 0)
        layout.setSpacing(0)
        layout.addWidget(mode_checkboxes)
        layout.addWidget(search)
        layout.addWidget(view)
        layout.addWidget(families_label)
        layout.addWidget(families)

        view.setContextMenuPolicy(QtCore.Qt.CustomContextMenu)
        view.setSelectionMode(QtWidgets.QAbstractItemView.ExtendedSelection)
        view.setSortingEnabled(True)
        view.sortByColumn(1, QtCore.Qt.AscendingOrder)
        view.setAlternatingRowColors(True)

        self.data = {
            "delegates": {
                "version": version_delegate,
                "time": time_delegate
            }
        }

        self.mode_input = input
        self.mode_output = output

        self.model = model
        self.proxy = proxy
        self.view = view

        # settings and connections
        proxy.setSourceModel(model)
        view.setModel(proxy)

        input.clicked.connect(self.on_mode_changed)
        output.clicked.connect(self.on_mode_changed)
        families.active_changed.connect(proxy.setFamiliesFilter)
        search.textChanged.connect(proxy.setFilterRegExp)

        families.refresh()
        # By default force all checkboxes to be enabled
        families._set_checkstate_all(True)

    def on_mode_changed(self):

        if self.mode_input.isChecked():
            # assume it's input
            self.model.set_mode(self.model.INPUTS)
            self.model.refresh()
        else:
            self.model.set_mode(self.model.OUTPUTS)
            self.model.refresh()

        # For now auto expand root level  auto-resize columns for readability
        self.view.expandToDepth(0)
        for i in range(self.model.rowCount(parent=QtCore.QModelIndex())):
            self.view.resizeColumnToContents(i)

show_representation_inputs.py

from avalon.vendor.Qt import QtWidgets, QtCore
from avalon import api
from avalon.tools import representationinput


def get_main_window():
    """Return top level QMainWindow"""
    top_widgets = QtWidgets.QApplication.topLevelWidgets()
    main_window = next((widget for widget in top_widgets if
                       widget.inherits("QMainWindow")), None)
    return main_window


class ShowRepresentationInputs(api.Loader):
    """Show Inputs/Outputs"""

    families = ["*"]
    representations = ["*"]

    label = "Show Inputs/Outputs"
    order = 999
    icon = "play-circle"
    color = "#444444"

    def load(self, context, name, namespace, data):

        representation = context["representation"]["_id"]

        parent = get_main_window()
        widget = representationinput.RepresentationWidget(parent=parent)
        widget.model.load([representation])
        widget.on_mode_changed()

        widget.setWindowTitle("List inputs/outputs for %s" % name)
        widget.resize(850, 400)
        widget.show()

Above is code related to a prototype as described here: getavalon/core#395 (comment)

Files (how to)

• representationinput.py - contains the main logic/widget for listing the inputs/outputs. On our end saved as avalon.tools.representationinput hence the imports on the other files pointing there. You will need to put it there due to the relative imports and the other files importing it from there.
• avalon_show_loaded_containers_inputs.py - is the prototype tool that allows to get inputs/outputs for loaded containers. You can run that code snippet in e.g. Houdini (from shelf) or Maya (from script editor)
• show_representation_inputs.py is a Loader plug-in that you can store under your global loaders so that any entry in the Loader allows to right click and "show inputs/outputs" ui for that single entry.

Known limitations

There is a known bug with the show loaded containers inputs UI where it will raise an error on something 'families.get()' on a None type. It's due to families in Avalon needing some internal cache refreshed. Either openthe Loader tool once first OR run this:

import avalon.tools.lib
avalon.tools.lib.refresh_family_config_cache()

Technically that code could also just be added in that code snippet for avalon_show_loaded_containers_inputs.py and it'll be fine.

---

I believe the aggregate function embedded in the code can actually be dropped as latest avalon core has it implemented in avalon.io with this PR: getavalon/core#555