From 54fd84dc4798c1bb1a8d2f472cf25c6496ac28b9 Mon Sep 17 00:00:00 2001
From: Samuel Liu <samuel.liu@autodesk.com>
Date: Mon, 4 Dec 2023 17:04:20 -0500
Subject: [PATCH] Add bulk editing support for unload and load with descendants

---
 lib/usdUfe/ufe/UsdContextOps.cpp | 35 ++++++++++++++++++++++++
 test/lib/ufe/testContextOps.py   | 47 ++++++++++++++++++++++++++++++--
 2 files changed, 80 insertions(+), 2 deletions(-)

diff --git a/lib/usdUfe/ufe/UsdContextOps.cpp b/lib/usdUfe/ufe/UsdContextOps.cpp
index dc58352b8d..a5344e5b5c 100644
--- a/lib/usdUfe/ufe/UsdContextOps.cpp
+++ b/lib/usdUfe/ufe/UsdContextOps.cpp
@@ -490,6 +490,8 @@ void UsdContextOps::addBulkEditHeader(Ufe::ContextOps::Items& items) const
  *
  *      "{countOfPrimsSelected} {PrimType} Prims Selected" - disbled item has no action
  *      -----------------
+ *      Unload
+ *      Load with Descendants
  *      Make Visible
  *      Make Invisible
  *      Activate Prim
@@ -504,6 +506,12 @@ Ufe::ContextOps::Items UsdContextOps::getBulkItems(const ItemPath& itemPath) con
     if (itemPath.empty()) {
         addBulkEditHeader(items);
 
+        // Unload
+        items.emplace_back(kUSDUnloadItem, kUSDUnloadLabel);
+
+        // Load With Descendants
+        items.emplace_back(kUSDLoadWithDescendantsItem, kUSDLoadWithDescendantsLabel);
+
         // Visibility:
         items.emplace_back(kUSDMakeVisibleItem, kUSDMakeVisibleLabel);
         items.emplace_back(kUSDMakeInvisibleItem, kUSDMakeInvisibleLabel);
@@ -616,6 +624,33 @@ Ufe::UndoableCommand::Ptr UsdContextOps::doBulkOpCmd(const ItemPath& itemPath)
                                 : nullptr;
     };
 
+    // Unload:
+    if (itemPath[0u] == kUSDUnloadItem) {
+        for (auto& selItem : _bulkItems) {
+            UsdSceneItem::Ptr usdItem = std::dynamic_pointer_cast<UsdSceneItem>(selItem);
+            if (usdItem) {
+                auto cmd = std::make_shared<UsdUndoUnloadPayloadCommand>(usdItem->prim());
+                cmdList.emplace_back(cmd);
+            }
+        }
+        return compositeCmdReturn(_bulkItems);
+    }
+
+    // Load With Descendants:
+    if (itemPath[0u] == kUSDLoadWithDescendantsItem) {
+        for (auto& selItem : _bulkItems) {
+            UsdSceneItem::Ptr   usdItem = std::dynamic_pointer_cast<UsdSceneItem>(selItem);
+            const UsdLoadPolicy policy = (itemPath[0u] == kUSDLoadWithDescendantsItem)
+                ? UsdLoadWithDescendants
+                : UsdLoadWithoutDescendants;
+            if (usdItem) {
+                auto cmd = std::make_shared<UsdUndoLoadPayloadCommand>(usdItem->prim(), policy);
+                cmdList.emplace_back(cmd);
+            }
+        }
+        return compositeCmdReturn(_bulkItems);
+    }
+
     // Prim Visibility:
     const bool makeVisible = itemPath[0u] == kUSDMakeVisibleItem;
     const bool makeInvisible = itemPath[0u] == kUSDMakeInvisibleItem;
diff --git a/test/lib/ufe/testContextOps.py b/test/lib/ufe/testContextOps.py
index 3256184bfa..758a36b705 100644
--- a/test/lib/ufe/testContextOps.py
+++ b/test/lib/ufe/testContextOps.py
@@ -37,6 +37,7 @@
 import maya.api.OpenMaya as om
 
 import ufe
+import usdUfe
 
 import os
 import unittest
@@ -188,8 +189,8 @@ def testGetBulkItems(self):
 
         # Not supported in bulk (from UsdContextOps).
         self.assertNotIn('Load', contextItemStrings)
-        self.assertNotIn('Load with Descendants', contextItemStrings)
-        self.assertNotIn('Unload', contextItemStrings)
+        self.assertIn('Load with Descendants', contextItemStrings)
+        self.assertIn('Unload', contextItemStrings)
         self.assertNotIn('Variant Sets', contextItemStrings)
         self.assertNotIn('Add New Prim', contextItemStrings)
 
@@ -500,6 +501,48 @@ def testDoBulkOp(self):
         # Since the context item is in the selection, we get a bulk context.
         self.contextOps = ufe.ContextOps.contextOps(self.ball35Item)
 
+        # Test Bulk Unload and Load with Descendants
+        payloadFile = testUtils.getTestScene('twoSpheres', 'sphere.usda')
+
+        def addPayLoads(payloadFile, ballPrims):
+            for ballPrim in ballPrims.values():
+                cmd = usdUfe.AddPayloadCommand(ballPrim, payloadFile, True)
+                self.assertIsNotNone(cmd)
+
+                # Verify state after add payload
+                cmd.execute()
+                self.assertTrue(ballPrim.HasPayload())
+                self.assertTrue(ballPrim.IsLoaded())
+    
+        def verifyBulkPrimPayload(ballPrims, ifLoaded):
+            for ballPrim in ballPrims.values():
+                self.assertEqual(ballPrim.IsLoaded(), ifLoaded)
+
+        addPayLoads(payloadFile, ballPrims)
+
+        # Unload
+        cmd = self.contextOps.doOpCmd(['Unload'])
+        self.assertIsNotNone(cmd)
+        self.assertIsInstance(cmd, ufe.CompositeUndoableCommand)
+
+        ufeCmd.execute(cmd)
+        verifyBulkPrimPayload(ballPrims, False)
+        cmds.undo()
+        verifyBulkPrimPayload(ballPrims, True)
+
+        # Load with Descendants
+        # Unload the payloads first, using the unload command
+        ufeCmd.execute(cmd)
+        cmd = self.contextOps.doOpCmd(['Load with Descendants'])
+        self.assertIsNotNone(cmd)
+        self.assertIsInstance(cmd, ufe.CompositeUndoableCommand)
+
+        ufeCmd.execute(cmd)
+        verifyBulkPrimPayload(ballPrims, True)
+        cmds.undo()
+        verifyBulkPrimPayload(ballPrims, False)
+
+
         # Change visility, undo/redo.
         cmd = self.contextOps.doOpCmd(['Make Invisible'])
         self.assertIsNotNone(cmd)