BigRoy/usd_list_layer_edits_editor.py

## usd_list_layer_edits_editor.py
from PySide2 import QtCore, QtWidgets, QtGui
from pxr import Usd, Tf, Sdf

# See: https://github.com/PixarAnimationStudios/OpenUSD/blob/release/pxr/usd/sdf/fileIO_Common.cpp#L879-L892
SPECIFIER_LABEL = {
    Sdf.SpecifierDef: "def",
    Sdf.SpecifierOver: "over",
    Sdf.SpecifierClass: "abstract"
}


def shorten(s, width, placeholder="..."):
    """Shorten string to `width`"""
    if len(s) <= width:
        return s
    return "{}{}".format(s[:width], placeholder)


def remove_spec(spec):
    """Remove Sdf.Spec authored opinion."""
    if spec.expired:
        return

    if isinstance(spec, Sdf.PrimSpec):
        # PrimSpec
        parent = spec.nameParent
        if parent:
            view = parent.nameChildren
        else:
            # Assume PrimSpec is root prim
            view = spec.layer.rootPrims
        del view[spec.name]

    elif isinstance(spec, Sdf.PropertySpec):
        # Relationship and Attribute specs
        del spec.owner.properties[spec.name]
    else:
        raise TypeError(f"Unsupported spec type: {spec}")


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=None):
        if parent is None or not parent.isValid():
            parent_item = self._root_item
        else:
            parent_item = parent.internalPointer()
        return parent_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

                self.dataChanged.emit(index, index)
                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=None):
        """Return index for row/column under parent"""

        if parent is None or 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 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)
        return -1

    def add_child(self, child):
        """Add a child to this item"""
        child._parent = self
        self._children.append(child)


class StageSdfModel(TreeModel):

    Columns = ["name", "specifier", "typeName", "default", "type"]

    Colors = {
        "Layer": QtGui.QColor("#008EC5"),
        "PseudoRootSpec": QtGui.QColor("#A2D2EF"),
        "PrimSpec": QtGui.QColor("#A2D2EF"),
        "RelationshipSpec": QtGui.QColor("#FCD057"),
        "AttributeSpec": QtGui.QColor("#FFC8DD"),
    }


    def __init__(self, stage, parent=None):
        super(StageSdfModel, self).__init__(parent)
        self._stage = stage

    def refresh(self):
        self.clear()

        for layer in stage.GetLayerStack():

            layer_item = Item({
                "name": layer.identifier,
                "identifier": layer.identifier,
                "specifier": None,
                "color": "red",
                "type": layer.__class__.__name__
            })
            self.add_child(layer_item)

            items_by_path = {}

            def _traverse(path):
                spec = layer.GetObjectAtPath(path)
                if not spec:
                    # ignore target list binding entries
                    items_by_path[path] = Item({
                        "name": path.elementString,
                        "path": path,
                        "type": path.__class__.__name__
                    })
                    return

                spec_item = Item({
                    "name": spec.name,
                    "spec": spec,
                    "path": path,
                    "type": spec.__class__.__name__
                })

                if isinstance(spec, Sdf.PrimSpec):
                    spec_item["specifier"] = SPECIFIER_LABEL.get(spec.specifier)
                    type_name = spec.typeName
                    spec_item["typeName"] = type_name

                elif isinstance(spec, Sdf.AttributeSpec):
                    spec_item["default"] = shorten(str(spec.default), 60)

                items_by_path[path] = spec_item

            layer.Traverse("/", _traverse)

            # Build hierarchy of item of specs
            for path, item in sorted(items_by_path.items()):
                parent = path.GetParentPath()
                parent_item = items_by_path.get(parent, layer_item)
                parent_item.add_child(item)

    def data(self, index, role):

        if role == QtCore.Qt.ForegroundRole:
            item = index.data(TreeModel.ItemRole)
            class_type_name = item.get("type")
            color = self.Colors.get(class_type_name)
            if color:
                return color


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


class SpecEditsWidget(QtWidgets.QDialog):
    def __init__(self, stage, parent=None):
        super(SpecEditsWidget, self).__init__(parent=parent)

        layout = QtWidgets.QVBoxLayout(self)

        self.setWindowTitle("USD Layer Spec Editor")
        model = StageSdfModel(stage)
        view = QtWidgets.QTreeView()
        view.setModel(model)
        view.setIndentation(10)
        view.setStyleSheet("QTreeView::item { padding: 3px }")
        view.setSelectionMode(QtWidgets.QAbstractItemView.ExtendedSelection)

        auto_refresh = QtWidgets.QCheckBox("Auto Refresh on Stage Changes")
        auto_refresh.setChecked(True)
        refresh = QtWidgets.QPushButton("Refresh")
        delete = QtWidgets.QPushButton("Delete")

        layout.addWidget(view)
        layout.addWidget(auto_refresh)
        layout.addWidget(refresh)
        layout.addWidget(delete)

        self.auto_refresh = auto_refresh
        self.model = model
        self.view = view

        auto_refresh.stateChanged.connect(self.set_refresh_on_changes)
        refresh.clicked.connect(self.on_refresh)
        delete.clicked.connect(self.on_delete)

        self._listeners = []

        self.set_refresh_on_changes(True)
        self.on_refresh()

    def set_refresh_on_changes(self, state):
        if state:
            if self._listeners:
                return
            print("Adding listener")
            sender = self.model._stage
            listener = Tf.Notice.Register(Usd.Notice.StageContentsChanged,
                                          self.on_stage_changed_notice,
                                          sender)
            self._listeners.append(listener)
        else:
            if not self._listeners:
                return
            print("Removing listener")
            for listener in self._listeners:
                listener.Revoke()
            self._listeners.clear()

    def on_stage_changed_notice(self, notice, sender):
        self.on_refresh()

    def showEvent(self, event):
        state = self.auto_refresh.checkState() == QtCore.Qt.Checked
        self.set_refresh_on_changes(state)

    def hideEvent(self, event):
        # Remove any callbacks if they exist
        self.set_refresh_on_changes(False)

    def on_refresh(self):
        self.model.refresh()
        self.view.resizeColumnToContents(0)
        self.view.expandAll()
        self.view.resizeColumnToContents(1)
        self.view.resizeColumnToContents(2)
        self.view.resizeColumnToContents(3)
        self.view.resizeColumnToContents(4)

    def on_delete(self):
        selection_model = self.view.selectionModel()
        rows = selection_model.selectedRows()
        specs = []
        for row in rows:
            item = row.data(TreeModel.ItemRole)
            spec = item.get("spec")
            if spec:
                specs.append(spec)

        if not specs:
            return

        with Sdf.ChangeBlock():
            for spec in specs:
                print(f"Removing spec: {spec.path}")
                remove_spec(spec)

        self.on_refresh()
	from PySide2 import QtCore, QtWidgets, QtGui
	from pxr import Usd, Tf, Sdf

	# See: https://github.com/PixarAnimationStudios/OpenUSD/blob/release/pxr/usd/sdf/fileIO_Common.cpp#L879-L892
	SPECIFIER_LABEL = {
	Sdf.SpecifierDef: "def",
	Sdf.SpecifierOver: "over",
	Sdf.SpecifierClass: "abstract"
	}


	def shorten(s, width, placeholder="..."):
	"""Shorten string to `width`"""
	if len(s) <= width:
	return s
	return "{}{}".format(s[:width], placeholder)


	def remove_spec(spec):
	"""Remove Sdf.Spec authored opinion."""
	if spec.expired:
	return

	if isinstance(spec, Sdf.PrimSpec):
	# PrimSpec
	parent = spec.nameParent
	if parent:
	view = parent.nameChildren
	else:
	# Assume PrimSpec is root prim
	view = spec.layer.rootPrims
	del view[spec.name]

	elif isinstance(spec, Sdf.PropertySpec):
	# Relationship and Attribute specs
	del spec.owner.properties[spec.name]
	else:
	raise TypeError(f"Unsupported spec type: {spec}")



	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=None):
	if parent is None or not parent.isValid():
	parent_item = self._root_item
	else:
	parent_item = parent.internalPointer()
	return parent_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

	self.dataChanged.emit(index, index)
	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=None):
	"""Return index for row/column under parent"""

	if parent is None or 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 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)
	return -1

	def add_child(self, child):
	"""Add a child to this item"""
	child._parent = self
	self._children.append(child)



	class StageSdfModel(TreeModel):

	Columns = ["name", "specifier", "typeName", "default", "type"]

	Colors = {
	"Layer": QtGui.QColor("#008EC5"),
	"PseudoRootSpec": QtGui.QColor("#A2D2EF"),
	"PrimSpec": QtGui.QColor("#A2D2EF"),
	"RelationshipSpec": QtGui.QColor("#FCD057"),
	"AttributeSpec": QtGui.QColor("#FFC8DD"),
	}


	def __init__(self, stage, parent=None):
	super(StageSdfModel, self).__init__(parent)
	self._stage = stage

	def refresh(self):
	self.clear()

	for layer in stage.GetLayerStack():

	layer_item = Item({
	"name": layer.identifier,
	"identifier": layer.identifier,
	"specifier": None,
	"color": "red",
	"type": layer.__class__.__name__
	})
	self.add_child(layer_item)

	items_by_path = {}

	def _traverse(path):
	spec = layer.GetObjectAtPath(path)
	if not spec:
	# ignore target list binding entries
	items_by_path[path] = Item({
	"name": path.elementString,
	"path": path,
	"type": path.__class__.__name__
	})
	return

	spec_item = Item({
	"name": spec.name,
	"spec": spec,
	"path": path,
	"type": spec.__class__.__name__
	})

	if isinstance(spec, Sdf.PrimSpec):
	spec_item["specifier"] = SPECIFIER_LABEL.get(spec.specifier)
	type_name = spec.typeName
	spec_item["typeName"] = type_name

	elif isinstance(spec, Sdf.AttributeSpec):
	spec_item["default"] = shorten(str(spec.default), 60)

	items_by_path[path] = spec_item

	layer.Traverse("/", _traverse)

	# Build hierarchy of item of specs
	for path, item in sorted(items_by_path.items()):
	parent = path.GetParentPath()
	parent_item = items_by_path.get(parent, layer_item)
	parent_item.add_child(item)

	def data(self, index, role):

	if role == QtCore.Qt.ForegroundRole:
	item = index.data(TreeModel.ItemRole)
	class_type_name = item.get("type")
	color = self.Colors.get(class_type_name)
	if color:
	return color


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


	class SpecEditsWidget(QtWidgets.QDialog):
	def __init__(self, stage, parent=None):
	super(SpecEditsWidget, self).__init__(parent=parent)

	layout = QtWidgets.QVBoxLayout(self)

	self.setWindowTitle("USD Layer Spec Editor")
	model = StageSdfModel(stage)
	view = QtWidgets.QTreeView()
	view.setModel(model)
	view.setIndentation(10)
	view.setStyleSheet("QTreeView::item { padding: 3px }")
	view.setSelectionMode(QtWidgets.QAbstractItemView.ExtendedSelection)

	auto_refresh = QtWidgets.QCheckBox("Auto Refresh on Stage Changes")
	auto_refresh.setChecked(True)
	refresh = QtWidgets.QPushButton("Refresh")
	delete = QtWidgets.QPushButton("Delete")

	layout.addWidget(view)
	layout.addWidget(auto_refresh)
	layout.addWidget(refresh)
	layout.addWidget(delete)

	self.auto_refresh = auto_refresh
	self.model = model
	self.view = view

	auto_refresh.stateChanged.connect(self.set_refresh_on_changes)
	refresh.clicked.connect(self.on_refresh)
	delete.clicked.connect(self.on_delete)

	self._listeners = []

	self.set_refresh_on_changes(True)
	self.on_refresh()

	def set_refresh_on_changes(self, state):
	if state:
	if self._listeners:
	return
	print("Adding listener")
	sender = self.model._stage
	listener = Tf.Notice.Register(Usd.Notice.StageContentsChanged,
	self.on_stage_changed_notice,
	sender)
	self._listeners.append(listener)
	else:
	if not self._listeners:
	return
	print("Removing listener")
	for listener in self._listeners:
	listener.Revoke()
	self._listeners.clear()

	def on_stage_changed_notice(self, notice, sender):
	self.on_refresh()

	def showEvent(self, event):
	state = self.auto_refresh.checkState() == QtCore.Qt.Checked
	self.set_refresh_on_changes(state)

	def hideEvent(self, event):
	# Remove any callbacks if they exist
	self.set_refresh_on_changes(False)

	def on_refresh(self):
	self.model.refresh()
	self.view.resizeColumnToContents(0)
	self.view.expandAll()
	self.view.resizeColumnToContents(1)
	self.view.resizeColumnToContents(2)
	self.view.resizeColumnToContents(3)
	self.view.resizeColumnToContents(4)

	def on_delete(self):
	selection_model = self.view.selectionModel()
	rows = selection_model.selectedRows()
	specs = []
	for row in rows:
	item = row.data(TreeModel.ItemRole)
	spec = item.get("spec")
	if spec:
	specs.append(spec)

	if not specs:
	return

	with Sdf.ChangeBlock():
	for spec in specs:
	print(f"Removing spec: {spec.path}")
	remove_spec(spec)

	self.on_refresh()