EMSUSD-895 - Payload command undo to restore payload rules

lib/usdUfe/ufe/UsdUndoPayloadCommand.cpp

@@ -47,23 +47,24 @@ UsdUndoLoadUnloadBaseCommand::UsdUndoLoadUnloadBaseCommand(const PXR_NS::UsdPrim
         : PXR_NS::UsdLoadPolicy::UsdLoadWithoutDescendants;
 }
 
-void UsdUndoLoadUnloadBaseCommand::doLoad() const
+void UsdUndoLoadUnloadBaseCommand::doCommand(UsdUndoLoadUnloadBaseCommand::commandFn fn, bool undo)
 {
+    // handle the redo/undo commands on stage load/unload rules
+    // The load rules are saved before applying new rules and are
+    // applied back when the commands are undo
     if (!_stage)
         return;
-
-    _stage->Load(_primPath, _policy);
+    if (!undo)
+        _undoRules = _stage->GetLoadRules();
+    fn(this);
+    if (undo)
+        _stage->SetLoadRules(_undoRules);
     saveModifiedLoadRules();
 }
 
-void UsdUndoLoadUnloadBaseCommand::doUnload() const
-{
-    if (!_stage)
-        return;
+void UsdUndoLoadUnloadBaseCommand::doLoad() const { _stage->Load(_primPath, _policy); }
 
-    _stage->Unload(_primPath);
-    saveModifiedLoadRules();
-}
+void UsdUndoLoadUnloadBaseCommand::doUnload() const { _stage->Unload(_primPath); }
 
 void UsdUndoLoadUnloadBaseCommand::saveModifiedLoadRules() const
 {
@@ -81,8 +82,11 @@ UsdUndoLoadPayloadCommand::UsdUndoLoadPayloadCommand(
 {
 }
 
-void UsdUndoLoadPayloadCommand::redo() { doLoad(); }
-void UsdUndoLoadPayloadCommand::undo() { doUnload(); }
+void UsdUndoLoadPayloadCommand::redo() { doCommand(&UsdUndoLoadPayloadCommand::doLoad); }
+void UsdUndoLoadPayloadCommand::undo()
+{
+    doCommand(&UsdUndoLoadPayloadCommand::doUnload, true /* undo */);
+}
 
 USDUFE_VERIFY_CLASS_SETUP(UsdUndoLoadUnloadBaseCommand, UsdUndoUnloadPayloadCommand);
 
@@ -91,7 +95,10 @@ UsdUndoUnloadPayloadCommand::UsdUndoUnloadPayloadCommand(const PXR_NS::UsdPrim&
 {
 }
 
-void UsdUndoUnloadPayloadCommand::redo() { doUnload(); }
-void UsdUndoUnloadPayloadCommand::undo() { doLoad(); }
+void UsdUndoUnloadPayloadCommand::redo() { doCommand(&UsdUndoUnloadPayloadCommand::doUnload); }
+void UsdUndoUnloadPayloadCommand::undo()
+{
+    doCommand(&UsdUndoUnloadPayloadCommand::doLoad, true /* undo */);
+}
 
 } // namespace USDUFE_NS_DEF

lib/usdUfe/ufe/UsdUndoPayloadCommand.h

@@ -36,6 +36,8 @@ class USDUFE_PUBLIC UsdUndoLoadUnloadBaseCommand : public Ufe::UndoableCommand
 
     USDUFE_DISALLOW_COPY_MOVE_AND_ASSIGNMENT(UsdUndoLoadUnloadBaseCommand);
 
+    using commandFn = std::function<void(const UsdUndoLoadUnloadBaseCommand*)>;
+    void doCommand(commandFn fn, bool undo = false);
     void doLoad() const;
     void doUnload() const;
 
@@ -45,6 +47,7 @@ class USDUFE_PUBLIC UsdUndoLoadUnloadBaseCommand : public Ufe::UndoableCommand
     const PXR_NS::UsdStageWeakPtr _stage;
     const PXR_NS::SdfPath         _primPath;
     PXR_NS::UsdLoadPolicy         _policy;
+    PXR_NS::UsdStageLoadRules     _undoRules;
 };
 
 //! \brief Undoable command for loading a USD prim.

test/lib/ufe/testPayloadCommands.py

@@ -144,6 +144,63 @@ def testLoadAndUnloadPayloadCommands(self):
         cmd.redo()
         self.assertTrue(prim.IsLoaded())
 
+    def testUnloadPayloadUndoStackCommands(self):
+        '''
+        Test unload payload undo stack commands.
+        '''
+        withPayloadFile = testUtils.getTestScene("payload", "multipots.usda")
+        # scene to be loaded
+        # \- FlowerPot_set
+        #  \- Variants_grp
+        #   \- FlowerPotA (payload)
+        #   \- FlowerPotB (payload)
+        withPayloadDagPath, withPayloadStage = mayaUtils.createProxyFromFile(withPayloadFile)
+        flowerPotAUfePathString = ','.join([withPayloadDagPath, "/FlowerPot_set/Variants_grp/FlowerPotA"])
+        primA = mayaUsd.ufe.ufePathToPrim(flowerPotAUfePathString)
+        flowerPotBUfePathString = ','.join([withPayloadDagPath, "/FlowerPot_set/Variants_grp/FlowerPotB"])
+        primB = mayaUsd.ufe.ufePathToPrim(flowerPotBUfePathString)
+
+        # Verify initial state
+        self.assertTrue(primA.HasPayload())
+        self.assertTrue(primB.HasPayload())
+
+        # Verify unload payload sequence
+        # Emulate a stack of unload commands in a tree
+        # - first 'unload' FlowerPotA
+        # - then 'unload' FlowerPot_set; will unload all
+        # - undo 'unload' FlowerPot_set; will reload and set back the payload rules as they were
+        #                                prior to the 'unload', where FlowerPotA is unloaded
+        # - redo 'unload' FlowerPot_set
+        # - undo 'unload' FlowerPot_set
+        # - undo 'unload' FlowerPotA; which brings back everything to the original payload state
+        primAUnloadCmd = usdUfe.UnloadPayloadCommand(primA)
+        self.assertIsNotNone(primAUnloadCmd)
+        primAUnloadCmd.execute()
+        self.assertFalse(primA.IsLoaded())
+        self.assertTrue(primB.IsLoaded())
+
+        flowerPotSetUfePathString = ','.join([withPayloadDagPath, "/FlowerPot_set"])
+        primSet = mayaUsd.ufe.ufePathToPrim(flowerPotSetUfePathString)
+        primSetUnloadCmd = usdUfe.UnloadPayloadCommand(primSet)
+        self.assertIsNotNone(primSetUnloadCmd)
+
+        primSetUnloadCmd.execute()
+        self.assertFalse(primA.IsLoaded())
+        self.assertFalse(primB.IsLoaded())
+
+        primSetUnloadCmd.undo()
+        self.assertFalse(primA.IsLoaded())
+        self.assertTrue(primB.IsLoaded())
+
+        primSetUnloadCmd.redo()
+        self.assertFalse(primA.IsLoaded())
+        self.assertFalse(primB.IsLoaded())
+
+        primSetUnloadCmd.undo()
+        primAUnloadCmd.undo()
+        self.assertTrue(primA.IsLoaded())
+        self.assertTrue(primB.IsLoaded())
+
     @unittest.skipUnless(mayaUtils.mayaMajorVersion() >= 2023, 'Delete restriction on delete requires Maya 2023 or greater.')
     def testDeletePrimContainingPayload(self):
         '''

test/testSamples/payload/multipots.usda

@@ -0,0 +1,36 @@
+#usda 1.0
+(
+    defaultPrim = "Item_set"
+    upAxis = "Z"
+)
+
+def Xform "FlowerPot_set" (
+    kind = "assembly"
+)
+{
+    def Xform "Variants_grp" (
+        kind = "group"
+    )
+    {
+		def "FlowerPotA" (
+			add references = @./FlowerPot.usda@
+			variants = {
+				string modelingVariant = "FlowerPotA"
+			}
+		)
+		{
+		}
+		def "FlowerPotB" (
+			add references = @./FlowerPot.usda@
+			variants = {
+				string modelingVariant = "FlowerPotB"
+			}
+		)
+		{
+			double3 xformOp:translate = (30.0, 0.0, 0.0)
+			uniform token[] xformOpOrder = ["xformOp:translate"]
+		}
+	}
+}
+
+