Fix for MAYA-122381 Target Layer moves and stays on Session Layer whe…

lib/mayaUsd/fileio/jobs/jobArgs.cpp

@@ -88,6 +88,19 @@ bool _Boolean(const VtDictionary& userArgs, const TfToken& key)
     return VtDictionaryGet<bool>(userArgs, key);
 }
 
+/// Extracts a pointer at \p key from \p userArgs, or nullptr if it can't extract.
+UsdStageRefPtr _UsdStageRefPtr(const VtDictionary& userArgs, const TfToken& key)
+{
+    if (!VtDictionaryIsHolding<UsdStageRefPtr>(userArgs, key)) {
+        TF_CODING_ERROR(
+            "Dictionary is missing required key '%s' or key is "
+            "not pointer type",
+            key.GetText());
+        return nullptr;
+    }
+    return VtDictionaryGet<UsdStageRefPtr>(userArgs, key);
+}
+
 /// Extracts a double at \p key from \p userArgs, or defaultValue if it can't extract.
 double _Double(const VtDictionary& userArgs, const TfToken& key, double defaultValue)
 {
@@ -1071,6 +1084,7 @@ UsdMayaJobImportArgs::UsdMayaJobImportArgs(
     , importInstances(_Boolean(userArgs, UsdMayaJobImportArgsTokens->importInstances))
     , useAsAnimationCache(_Boolean(userArgs, UsdMayaJobImportArgsTokens->useAsAnimationCache))
     , importWithProxyShapes(importWithProxyShapes)
+    , pullImportStage(_UsdStageRefPtr(userArgs, UsdMayaJobImportArgsTokens->pullImportStage))
     , timeInterval(timeInterval)
     , chaserNames(_Vector<std::string>(userArgs, UsdMayaJobImportArgsTokens->chaser))
     , allChaserArgs(_ChaserArgs(userArgs, UsdMayaJobImportArgsTokens->chaserArgs))
@@ -1122,6 +1136,7 @@ const VtDictionary& UsdMayaJobImportArgs::GetDefaultDictionary()
         d[UsdMayaJobImportArgsTokens->importInstances] = true;
         d[UsdMayaJobImportArgsTokens->importUSDZTextures] = false;
         d[UsdMayaJobImportArgsTokens->importUSDZTexturesFilePath] = "";
+        d[UsdMayaJobImportArgsTokens->pullImportStage] = UsdStageRefPtr();
         d[UsdMayaJobImportArgsTokens->useAsAnimationCache] = false;
         d[UsdMayaJobExportArgsTokens->chaser] = std::vector<VtValue>();
         d[UsdMayaJobExportArgsTokens->chaserArgs] = std::vector<VtValue>();
@@ -1179,6 +1194,7 @@ const VtDictionary& UsdMayaJobImportArgs::GetGuideDictionary()
     std::call_once(once, []() {
         // Common types:
         const auto _boolean = VtValue(false);
+        const auto _usdStageRefPtr = VtValue(nullptr);
         const auto _string = VtValue(std::string());
         const auto _stringVector = VtValue(std::vector<VtValue>({ _string }));
         const auto _stringTuplet = VtValue(std::vector<VtValue>({ _string, _string, _string }));
@@ -1197,6 +1213,7 @@ const VtDictionary& UsdMayaJobImportArgs::GetGuideDictionary()
         d[UsdMayaJobImportArgsTokens->importInstances] = _boolean;
         d[UsdMayaJobImportArgsTokens->importUSDZTextures] = _boolean;
         d[UsdMayaJobImportArgsTokens->importUSDZTexturesFilePath] = _string;
+        d[UsdMayaJobImportArgsTokens->pullImportStage] = _usdStageRefPtr;
         d[UsdMayaJobImportArgsTokens->useAsAnimationCache] = _boolean;
         d[UsdMayaJobExportArgsTokens->chaser] = _stringVector;
         d[UsdMayaJobExportArgsTokens->chaserArgs] = _stringTripletVector;
@@ -1282,6 +1299,7 @@ std::ostream& operator<<(std::ostream& out, const UsdMayaJobImportArgs& importAr
         << "importInstances: " << TfStringify(importArgs.importInstances) << std::endl
         << "importUSDZTextures: " << TfStringify(importArgs.importUSDZTextures) << std::endl
         << "importUSDZTexturesFilePath: " << TfStringify(importArgs.importUSDZTexturesFilePath)
+        << "pullImportStage: " << TfStringify(importArgs.pullImportStage) << std::endl
         << std::endl
         << "timeInterval: " << importArgs.timeInterval << std::endl
         << "useAsAnimationCache: " << TfStringify(importArgs.useAsAnimationCache) << std::endl

lib/mayaUsd/fileio/jobs/jobArgs.h

@@ -142,6 +142,7 @@ TF_DECLARE_PUBLIC_TOKENS(
     (importInstances) \
     (importUSDZTextures) \
     (importUSDZTexturesFilePath) \
+    (pullImportStage) \
     /* assemblyRep values */ \
     (Collapsed) \
     (Full) \
@@ -310,13 +311,14 @@ struct UsdMayaJobImportArgs
         TfToken materialConversion;
     };
     using ShadingModes = std::vector<ShadingMode>;
-    ShadingModes      shadingModes; // XXX can we make this const?
-    const TfToken     preferredMaterial;
-    const std::string importUSDZTexturesFilePath;
-    const bool        importUSDZTextures;
-    const bool        importInstances;
-    const bool        useAsAnimationCache;
-    const bool        importWithProxyShapes;
+    ShadingModes         shadingModes; // XXX can we make this const?
+    const TfToken        preferredMaterial;
+    const std::string    importUSDZTexturesFilePath;
+    const bool           importUSDZTextures;
+    const bool           importInstances;
+    const bool           useAsAnimationCache;
+    const bool           importWithProxyShapes;
+    const UsdStageRefPtr pullImportStage;
     /// The interval over which to import animated data.
     /// An empty interval (<tt>GfInterval::IsEmpty()</tt>) means that no
     /// animated (time-sampled) data should be imported.

lib/mayaUsd/fileio/jobs/readJob.cpp

@@ -31,6 +31,7 @@
 #include <pxr/usd/sdf/fileFormat.h>
 #include <pxr/usd/sdf/layer.h>
 #include <pxr/usd/sdf/path.h>
+#include <pxr/usd/usd/editContext.h>
 #include <pxr/usd/usd/prim.h>
 #include <pxr/usd/usd/primFlags.h>
 #include <pxr/usd/usd/primRange.h>
@@ -141,12 +142,16 @@ bool UsdMaya_ReadJob::Read(std::vector<MDagPath>* addedDagPaths)
     } else {
         UsdStageCacheContext stageCacheContext(UsdMayaStageCache::Get(
             mImportData.stageInitialLoadSet() == UsdStage::InitialLoadSet::LoadAll));
-        stage = UsdStage::Open(rootLayer, sessionLayer, mImportData.stageInitialLoadSet());
+        if (mArgs.pullImportStage)
+            stage = mArgs.pullImportStage;
+        else
+            stage = UsdStage::Open(rootLayer, sessionLayer, mImportData.stageInitialLoadSet());
     }
     if (!stage) {
         return false;
     }
 
+    UsdEditContext editContext(stage, stage->GetSessionLayer());
     stage->SetEditTarget(stage->GetSessionLayer());
     _setTimeSampleMultiplierFrom(stage->GetTimeCodesPerSecond());
 

lib/mayaUsd/fileio/primUpdaterManager.cpp

@@ -272,7 +272,8 @@ PullImportPaths pullImport(
         return PullImportPaths(addedDagPaths, pulledUfePaths);
     }
 
-    const VtDictionary& userArgs = context.GetUserArgs();
+    VtDictionary userArgs(context.GetUserArgs());
+    userArgs[UsdMayaJobImportArgsTokens->pullImportStage] = PXR_NS::VtValue(context.GetUsdStage());
 
     UsdMayaJobImportArgs jobArgs = UsdMayaJobImportArgs::CreateFromDictionary(
         userArgs,

test/lib/mayaUsd/fileio/testEditAsMaya.py

@@ -26,7 +26,7 @@
 import mayaUtils
 import mayaUsd.ufe
 
-from pxr import Usd, UsdGeom, Gf
+from pxr import Usd, UsdGeom, Gf, Sdf
 
 from maya import cmds
 from maya import standalone
@@ -255,5 +255,44 @@ def testSessionLayer(self):
 
         self.assertEqual(prim.GetCustomDataByKey(kPullPrimMetadataKey), None)
 
+
+    @unittest.skipUnless(ufeFeatureSetVersion() >= 3, 'Test only available in UFE v3 or greater.')
+    def testTargetLayer(self):
+        '''Verify that the target layer is not moved after Edit As Maya.'''
+
+        import mayaUsd_createStageWithNewLayer
+
+        proxyShapePathStr = mayaUsd_createStageWithNewLayer.createStageWithNewLayer()
+        stage = mayaUsd.lib.GetPrim(proxyShapePathStr).GetStage()
+        rootLayer = stage.GetRootLayer()
+
+        currentLayer = stage.GetEditTarget().GetLayer()
+        self.assertEqual(currentLayer, rootLayer) # Current layer should be the Anonymous Root Layer
+
+        sessionLayer = stage.GetSessionLayer()
+        prim = stage.DefinePrim('/A', 'Xform')
+
+        primPathStr = proxyShapePathStr + ',/A'
+
+        self.assertTrue(stage.GetSessionLayer().empty)
+
+        otherLayerId = cmds.mayaUsdLayerEditor(rootLayer.identifier, edit=True, addAnonymous="otherLayer")[0]
+        otherLayer = Sdf.Layer.Find(otherLayerId)
+
+        stage.SetEditTarget(otherLayer)
+
+        currentLayer = stage.GetEditTarget().GetLayer()
+        self.assertEqual(currentLayer, otherLayer) # Current layer should be the Other Layer
+
+        with mayaUsd.lib.OpUndoItemList():
+            self.assertTrue(mayaUsd.lib.PrimUpdaterManager.canEditAsMaya(primPathStr))
+            self.assertTrue(mayaUsd.lib.PrimUpdaterManager.editAsMaya(primPathStr))
+
+        self.assertFalse(stage.GetSessionLayer().empty)
+
+        currentLayer = stage.GetEditTarget().GetLayer()
+        self.assertEqual(currentLayer, otherLayer) # Current layer should still be the Other Layer
+
+
 if __name__ == '__main__':
     unittest.main(verbosity=2)