Last active
November 7, 2024 15:42
-
-
Save BigRoy/4d2bf2eef6c6a83f4fda3c58db1489a5 to your computer and use it in GitHub Desktop.
Quick and dirty "List USD Layer Edits" to allow removal of Sdf.PrimSpec, Sdf.PropertySpec, Sdf.AttributeSpec, Sdf.RelationshipSpec through a Python Qt interface
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
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() |
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Slightly more elaborate example with:
Code