Enhancement: Add shortcuts to duplicate, rename, delete or group selected prims in prim hierarchy

usd_qtpy/lib/usd.py

@@ -205,6 +205,74 @@ def rename_prim(prim: Usd.Prim, new_name: str) -> bool:
     return True
 
 
+def unique_name(stage: Usd.Stage, prim_path: Sdf.Path) -> Sdf.Path:
+    """Return Sdf.Path that is unique under the current composed stage.
+
+    Note that this technically does not ensure that the Sdf.Path does not
+    exist in any of the layers, e.g. it could be defined within a currently
+    unselected variant or a muted layer.
+
+    """
+    src = prim_path.pathString.rstrip("123456789")
+    i = 1
+    while stage.GetPrimAtPath(prim_path):
+        prim_path = Sdf.Path(f"{src}{i}")
+        i += 1
+    return prim_path
+
+
+def parent_prims(prims: list[Usd.Prim],
+                 new_parent: Sdf.Path,
+                 layers: list[Sdf.Layer] = None) -> bool:
+    """Move Prims to a new parent in given layers.
+
+    Note:
+        This will only reparent prims to the new parent if the new parent
+        exists in the layer.
+
+    Arguments:
+        prims (list[Usd.Prim]): The prims to move the new parent
+        new_parent (Sdf.Path): Parent path to be moved to.
+        layers (list[Sdf.Layer]): The layers to apply the reparenting
+            in. If None are provided the stage's full layer stack will be used.
+
+    """
+    if not prims:
+        return False
+
+    # Only consider prims not already parented to the new parent
+    prims = [
+        prim for prim in prims if prim.GetPath().GetParentPath() != new_parent
+    ]
+    if not prims:
+        return False
+
+    if layers is None:
+        stage = prims[0].GetStage()
+        layers = stage.GetLayerStack()
+
+    edit_batch = Sdf.BatchNamespaceEdit()
+    for prim in prims:
+        edit = Sdf.NamespaceEdit.Reparent(
+            prim.GetPath(),
+            new_parent,
+            -1
+        )
+        edit_batch.Add(edit)
+
+    any_edits_made = False
+    with Sdf.ChangeBlock():
+        for layer in layers:
+            applied = layer.Apply(edit_batch)
+            if applied:
+                any_edits_made = True
+                for edit in edit_batch.edits:
+                    repath_properties(layer,
+                                      edit.currentPath,
+                                      edit.newPath)
+    return any_edits_made
+
+
 def remove_spec(spec):
     """Remove Sdf.Spec authored opinion."""
     if spec.expired:

usd_qtpy/prim_hierarchy.py

@@ -1,14 +1,24 @@
+import logging
 from functools import partial
 
 from qtpy import QtWidgets, QtCore
 from pxr import Sdf
 
-from .lib.usd import get_prim_types_by_group
+from .lib.usd import (
+    get_prim_types_by_group,
+    parent_prims,
+    remove_spec,
+    unique_name,
+)
+from .lib.usd_merge_spec import copy_spec_merge
+from .lib.qt import iter_model_rows
 from .prim_delegate import DrawRectsDelegate
 from .prim_hierarchy_model import HierarchyModel
 from .references import ReferenceListWidget
 from .variants import CreateVariantSetDialog
 
+log = logging.getLogger(__name__)
+
 
 class View(QtWidgets.QTreeView):
     # TODO: Add shortcuts
@@ -42,13 +52,8 @@ def create_prim(action):
             type_name = action.text()
 
             # Ensure unique name
-            base_path = parent_path.AppendChild(type_name)
-            prim_path = base_path
-            i = 1
-            while stage.GetPrimAtPath(prim_path):
-                prim_path = Sdf.Path(f"{base_path.pathString}{i}")
-                i += 1
-
+            prim_path = parent_path.AppendChild(type_name)
+            prim_path = unique_name(stage, prim_path)
             if type_name == "Def":
                 # Typeless
                 type_name = ""
@@ -57,7 +62,8 @@ def create_prim(action):
             # TODO: Remove signaling once model listens to changes
             current_rows = model.rowCount(index)
             model.beginInsertRows(index, current_rows, current_rows+1)
-            stage.DefinePrim(prim_path, type_name)
+            new_prim = stage.DefinePrim(prim_path, type_name)
+            self.select_paths([new_prim.GetPath()])
             model.endInsertRows()
 
         # Create Prims
@@ -164,6 +170,219 @@ def on_prim_tag_clicked(self, event, index, block):
         elif text == "VAR":
             raise NotImplementedError("To be implemented")
 
+    def select_paths(self, paths: list[Sdf.Path]):
+        """Select prims in the hierarchy view that match the Sdf.Path
+
+        If an empty path list is provided or none matching paths are found
+        the selection is just cleared.
+
+        Arguments:
+            paths (list[Sdf.Path]): The paths to select.
+
+        """
+
+        model: HierarchyModel = self.model()
+        assert isinstance(model, HierarchyModel)
+        selection = QtCore.QItemSelection()
+
+        if not paths:
+            self.selectionModel().clear()
+            return
+
+        search = set(paths)
+        path_to_index = {}
+        for index in iter_model_rows(model, column=0):
+            # We iterate the model using its regular methods so we support both
+            # the model directly but also a proxy model. Also, this forces it
+            # to fetch the data if the model is lazy.
+            # TODO: This can be optimized by pruning the traversal if we
+            #   the current prim path is not a parent of the path we search for
+            if not search:
+                # Found all
+                break
+
+            prim = index.data(HierarchyModel.PrimRole)
+            if not prim:
+                continue
+
+            path = prim.GetPath()
+            path_to_index[path] = index
+            search.discard(path)
+
+        for path in paths:
+            index = path_to_index.get(path)
+            if not index:
+                # Path not found
+                continue
+
+            selection.select(index, index)
+
+        selection_model = self.selectionModel()
+        selection_model.select(selection,
+                               QtCore.QItemSelectionModel.ClearAndSelect |
+                               QtCore.QItemSelectionModel.Rows)
+
+    def keyPressEvent(self, event):
+        modifiers = event.modifiers()
+        ctrl_pressed = QtCore.Qt.ControlModifier & modifiers
+
+        # Group selected with Ctrl + G
+        if (
+            ctrl_pressed
+            and event.key() == QtCore.Qt.Key_G
+            and not event.isAutoRepeat()
+        ):
+            self._group_selected()
+            event.accept()
+            return
+
+        # Delete selected with delete key
+        if (
+            event.key() == QtCore.Qt.Key_Delete
+            and not event.isAutoRepeat()
+        ):
+            self._delete_selected()
+            event.accept()
+            return
+
+        # Duplicate selected with Ctrl + D
+        if (
+            ctrl_pressed
+            and event.key() == QtCore.Qt.Key_D
+            and not event.isAutoRepeat()
+        ):
+            self._duplicate_selected()
+            event.accept()
+            return
+
+        # Enter rename mode on current index when enter is pressed
+        if (
+            self.state() != QtWidgets.QAbstractItemView.EditingState
+            and event.key() in [QtCore.Qt.Key_Return, QtCore.Qt.Key_Enter]
+        ):
+            self.edit(self.currentIndex())
+            event.accept()
+            return
+
+        return super(View, self).keyPressEvent(event)
+
+    def _group_selected(self):
+        """Group selected prims under a new Xform"""
+        selected = self.selectionModel().selectedIndexes()
+        prims = [index.data(HierarchyModel.PrimRole) for index in selected]
+
+        # Exclude root prims
+        prims = [prim for prim in prims if not prim.IsPseudoRoot()]
+        if not prims:
+            return
+
+        stage = prims[0].GetStage()
+
+        # Consider only prims that have opinions in the stage's layer stack
+        # disregard opinions inside payloads/references
+        stage_layers = set(stage.GetLayerStack())
+
+        # Exclude prims not defined in the stage's layer stack
+        prims = [
+            prim for prim in prims
+            if any(spec.layer in stage_layers for spec in prim.GetPrimStack())
+        ]
+        if not prims:
+            log.warning("Skipped all prims because they are not defined in "
+                        "the stage's layer stack but likely originate from a "
+                        "reference or payload.")
+            return
+
+        parent_path = prims[0].GetPath().GetParentPath()
+
+        group_path = parent_path.AppendChild("group")
+        group_path = unique_name(stage, group_path)
+
+        # Define a group
+        stage.DefinePrim(group_path, "Xform")
+
+        # We want to group across all prim specs to ensure whatever we're
+        # moving gets put into the group, so we define the prim across all
+        # layers of the layer stack if it contains any of the objects
+        for layer in stage.GetLayerStack():
+            # If the layer has opinions on any of the source prims we ensure
+            # the new parent also exists, to ensure the movement of the input
+            # prims
+            if (
+                    any(layer.GetObjectAtPath(prim.GetPath())
+                        for prim in prims)
+                    and not layer.GetPrimAtPath(group_path)
+            ):
+                Sdf.CreatePrimInLayer(layer, group_path)
+
+        # Now we want to move all selected prims into this
+        parent_prims(prims, group_path)
+
+        # If the original group was renamed but there's now no conflict
+        # anymore, e.g. we grouped `group` itself from the parent path
+        # then now we can safely rename it to `group` without conflicts
+        # TODO: Ensure grouping `group` doesn't make a `group1`
+        self.select_paths([group_path])
+
+    def _delete_selected(self):
+        """Delete prims across all layers in the layer stack"""
+        selected = self.selectionModel().selectedIndexes()
+        prims = [index.data(HierarchyModel.PrimRole) for index in selected]
+
+        # Exclude root prims
+        prims = [prim for prim in prims if not prim.IsPseudoRoot()]
+        if not prims:
+            return
+
+        stage = prims[0].GetStage()
+        stage_layers = stage.GetLayerStack()
+
+        # We first collect the prim specs before removing because the Usd.Prim
+        # will become invalid as we start removing specs
+        specs = []
+        for prim in prims:
+            # We only allow deletions from layers in the current layer stack
+            # and exclude those that are from loaded references/payloads to
+            # avoid editing specs inside references/layers
+            for spec in prim.GetPrimStack():
+                if spec.layer in stage_layers:
+                    specs.append(spec)
+                else:
+                    logging.warning("Skipping prim spec not in "
+                                    "stage's layer stack: %s", spec)
+
+        with Sdf.ChangeBlock():
+            for spec in specs:
+                if spec.expired:
+                    continue
+
+                # Warning: This would also remove it from layers from
+                #   references/payloads!
+                # TODO: Filter specs for which their `.getLayer()` is a layer
+                #   from the Stage's layer stack?
+                remove_spec(spec)
+
+    def _duplicate_selected(self):
+        """Duplicate prim specs across all layers in the layer stack"""
+        selected = self.selectionModel().selectedIndexes()
+        prims = [index.data(HierarchyModel.PrimRole) for index in selected]
+
+        # Exclude root prims
+        prims = [prim for prim in prims if not prim.IsPseudoRoot()]
+        if not prims:
+            return []
+
+        new_paths = []
+        for prim in prims:
+            path = prim.GetPath()
+            stage = prim.GetStage()
+            new_path = unique_name(stage, path)
+            for spec in prim.GetPrimStack():
+                layer = spec.layer
+                copy_spec_merge(layer, path, layer, new_path)
+            new_paths.append(new_path)
+        self.select_paths(new_paths)
+
 
 class HierarchyWidget(QtWidgets.QDialog):
     def __init__(self, stage, parent=None):
@@ -175,6 +394,7 @@ def __init__(self, stage, parent=None):
 
         model = HierarchyModel(stage=stage)
         view = View()
+        view.setSelectionMode(QtWidgets.QTreeView.ExtendedSelection)
         view.setModel(model)
 
         self.model = model