From 3fb772da6a507716cef7f47e8611f23ea67492c5 Mon Sep 17 00:00:00 2001
From: Nick Wu <nickw@slid.al.com.au>
Date: Fri, 8 Apr 2022 14:51:48 +1000
Subject: [PATCH] allowed edit target layer to be written a new prim spec

---
 lib/mayaUsd/ufe/Utils.cpp          | 19 +++++--
 test/lib/ufe/testAttributeBlock.py | 81 +++++++++++++++++++++++++++++-
 2 files changed, 94 insertions(+), 6 deletions(-)

diff --git a/lib/mayaUsd/ufe/Utils.cpp b/lib/mayaUsd/ufe/Utils.cpp
index 03826023bd..b77a9e6592 100644
--- a/lib/mayaUsd/ufe/Utils.cpp
+++ b/lib/mayaUsd/ufe/Utils.cpp
@@ -78,7 +78,7 @@ bool stringBeginsWithDigit(const std::string& inputString)
 
 // This function calculates the position index for a given layer across all
 // the site's local LayerStacks
-uint32_t findLayerIndex(const UsdPrim& prim, const SdfLayerHandle& layer)
+uint32_t findLayerIndex(const UsdPrim& prim, const SdfLayerHandle& layer, bool& foundLayer)
 {
     uint32_t position { 0 };
 
@@ -96,6 +96,7 @@ uint32_t findLayerIndex(const UsdPrim& prim, const SdfLayerHandle& layer)
         // to find the layer
         for (SdfLayerRefPtr const& l : layerStack->GetLayers()) {
             if (l == layer) {
+                foundLayer = true;
                 return position;
             }
             ++position;
@@ -451,8 +452,15 @@ bool isAttributeEditAllowed(const PXR_NS::UsdAttribute& attr, std::string* errMs
         return false;
     }
 
+    bool foundLayer = false;
+
     // get the index to edit target layer
-    const auto targetLayerIndex = findLayerIndex(prim, editTarget.GetLayer());
+    const auto targetLayerIndex = findLayerIndex(prim, editTarget.GetLayer(), foundLayer);
+
+    // If the edit target layer not found in the prim's layer stack, we return true to allow
+    // new future edits go to the edit target layer.
+    if (!foundLayer)
+        return true;
 
     // HS March 22th,2021
     // TODO: "Value Clips" are UsdStage-level feature, unknown to Pcp.So if the attribute in
@@ -470,14 +478,15 @@ bool isAttributeEditAllowed(const PXR_NS::UsdAttribute& attr, std::string* errMs
     if (!propertyStack.empty()) {
         // get the strongest layer that has the attr.
         auto strongestLayer = attr.GetPropertyStack().front()->GetLayer();
+        foundLayer = false;
 
         // compare the calculated index between the "attr" and "edit target" layers.
-        if (findLayerIndex(prim, strongestLayer) < targetLayerIndex) {
+        if (findLayerIndex(prim, strongestLayer, foundLayer) < targetLayerIndex) {
             if (errMsg) {
                 *errMsg = TfStringPrintf(
                     "Cannot edit [%s] attribute because there is a stronger opinion in [%s].",
-                    attr.GetBaseName().GetText(),
-                    strongestLayer->GetDisplayName().c_str());
+                    attr.GetName().GetText(),
+                    strongestLayer->GetIdentifier().c_str());
             }
 
             return false;
diff --git a/test/lib/ufe/testAttributeBlock.py b/test/lib/ufe/testAttributeBlock.py
index 41cd679d52..46cde68ac9 100644
--- a/test/lib/ufe/testAttributeBlock.py
+++ b/test/lib/ufe/testAttributeBlock.py
@@ -21,7 +21,7 @@
 import ufeUtils
 import usdUtils
 
-from pxr import UsdGeom, Vt, Gf
+from pxr import UsdGeom, Vt, Gf, Usd, Sdf
 
 from maya import cmds
 from maya import standalone
@@ -32,6 +32,8 @@
 import ufe
 import os
 import unittest
+import textwrap
+
 
 class AttributeBlockTestCase(unittest.TestCase):
 
@@ -347,6 +349,83 @@ def testAttributeBlocking3dMatrixOps(self):
 
         # authoring new transformation edit is allowed.
         self.assertTrue(mayaUsdUfe.isAttributeEditAllowed(transformAttr))
+    
+    def testIsAttributeEditAllowedFunction(self):
+        rootLayerStr = """\
+        #sdf 1.4.32
+        ()
+        def "root"
+        {
+        }
+        """
+        subLayerAStr = """\
+        #sdf 1.4.32
+        over "root"
+        {
+            def "a" 
+            {
+                float foo = 0.0
+            }
+
+            def "b" 
+            {
+                float foo = 0.0
+            }
+        }
+        """
+        subLayerBStr = """\
+        #sdf 1.4.32
+        over "root"
+        {
+        }
+        """
+        subLayerCStr = """\
+        #sdf 1.4.32
+        over "root"
+        {
+            over "a" 
+            {
+                float foo = 10.0
+            }
+        }
+        """
+
+        rootLayer = Sdf.Layer.CreateAnonymous()
+        rootLayer.ImportFromString(textwrap.dedent(rootLayerStr))
+
+        subLayerA = Sdf.Layer.CreateAnonymous()
+        subLayerA.ImportFromString(textwrap.dedent(subLayerAStr))
+
+        subLayerB = Sdf.Layer.CreateAnonymous()
+        subLayerB.ImportFromString(textwrap.dedent(subLayerBStr))
+
+        subLayerC = Sdf.Layer.CreateAnonymous()
+        subLayerC.ImportFromString(textwrap.dedent(subLayerCStr))
+
+        rootLayer.subLayerPaths = [
+            subLayerC.identifier,
+            subLayerB.identifier,
+            subLayerA.identifier,
+        ]
+        stage = Usd.Stage.Open(rootLayer)
+
+        primB = stage.GetPrimAtPath("/root/b")
+        self.assertTrue(primB.IsValid())
+
+        fooattrB = primB.GetAttribute("foo")
+        self.assertTrue(fooattrB.IsValid())
+
+        stage.SetEditTarget(Usd.EditTarget(subLayerB))
+        self.assertTrue(mayaUsdUfe.isAttributeEditAllowed(fooattrB))
+
+        primA = stage.GetPrimAtPath("/root/a")
+        self.assertTrue(primA.IsValid())
+
+        fooattrA = primA.GetAttribute("foo")
+        self.assertTrue(fooattrA.IsValid())
+
+        self.assertFalse(mayaUsdUfe.isAttributeEditAllowed(fooattrA))
+
 
 if __name__ == '__main__':
     unittest.main(verbosity=2)