From 624c04c4aa44edabddcf30dd0bb393dc47472676 Mon Sep 17 00:00:00 2001
From: Krystian Ligenza <krystian.ligenza@autodesk.com>
Date: Tue, 25 Jan 2022 05:45:04 -0500
Subject: [PATCH] Improve edit as maya context menu

---
 lib/mayaUsd/fileio/primUpdaterManager.cpp |  2 +-
 plugin/adsk/scripts/mayaUsdMenu.mel       | 52 ++++++++++++++++++-----
 2 files changed, 42 insertions(+), 12 deletions(-)

diff --git a/lib/mayaUsd/fileio/primUpdaterManager.cpp b/lib/mayaUsd/fileio/primUpdaterManager.cpp
index 063f1bc5db..7e4e1d42eb 100644
--- a/lib/mayaUsd/fileio/primUpdaterManager.cpp
+++ b/lib/mayaUsd/fileio/primUpdaterManager.cpp
@@ -81,7 +81,7 @@ namespace {
 const std::string kPullParentPathKey("Maya:Pull:ParentPath");
 
 // Set name that will be used to hold all pulled objects
-const MString kPullSetName("pullStateSet");
+const MString kPullSetName("usdEditAsMaya");
 
 // Metadata key used to store pull information on a prim
 const TfToken kPullPrimMetadataKey("Maya:Pull:DagPath");
diff --git a/plugin/adsk/scripts/mayaUsdMenu.mel b/plugin/adsk/scripts/mayaUsdMenu.mel
index 3225ffd88d..26ff7362df 100644
--- a/plugin/adsk/scripts/mayaUsdMenu.mel
+++ b/plugin/adsk/scripts/mayaUsdMenu.mel
@@ -559,16 +559,45 @@ proc termAEShowMenu()
 // initPullMenu
 // Marking menu with USD options for DAG objects
 
-proc int isPulledUsdObject(string $dagObject) {
-    string $pullAttribute = ( $dagObject + ".Pull_UfePath" );
-    if (`attributeExists "Pull_UfePath" $dagObject`) {
-        string $ufePath = `getAttr $pullAttribute`;
-        if ($ufePath != "") {
-            return 1;
+proc string getPulledUsdObject(string $dagObject) {
+    string $root = "";
+    string $pulledDagObject = "";
+    
+    if (`objExists $dagObject`) {
+        // For some reason, outliner will provide a set instead of real
+        // selection when populating the menu under set. We go back to read the
+        // real selection
+        if (size(`ls -set $dagObject`) > 0) {
+            string $selection[] = `ls -sl`;
+            $dagObject = $selection[0];
+        }
+        if (size(`ls -dagObjects $dagObject`) > 0) {
+            for (;;) {
+                string $parents[] = `listRelatives -parent -fullPath $dagObject`;
+                if (size($parents) == 0) {
+                    $root = longNameOf($dagObject);
+                    break;
+                }
+                
+                // Not root, then maybe it's pulled object. Look for "metadata" attribute
+                if (`attributeExists "Pull_UfePath" $dagObject`) {
+                    string $pullAttribute = ( $dagObject + ".Pull_UfePath" );
+                    string $ufePath = `getAttr $pullAttribute`;
+                    if ($ufePath != "") {
+                        $pulledDagObject = $dagObject;
+                    }
+                }
+                
+                $dagObject = $parents[0];
+            }
         }
     }
     
-    return 0;
+    if ($root == "|__mayaUsd__" && $pulledDagObject != "") {
+        return $pulledDagObject;
+    }
+    
+    return "";
 }
 
 proc int isObjectSelected(string $obj)
@@ -618,10 +647,11 @@ global proc mayaUsdMenu_markingMenuCallback( string $obj )
 {
     if (!hasPrimUpdater())
         return;
-
-    if (isPulledUsdObject($obj)) {
-        string $pushback = `menuItem -label "Merge Maya Edits to USD" -insertAfter "" -image "merge_to_USD.png" -command ("mayaUsdMenu_pushBackToUSD " + $obj)`;
-        string $pushclear = `menuItem -label "Discard Maya Edits" -insertAfter $pushback -image "discard_edits.png" -command ("mayaUsdDiscardEdits " + $obj)`;
+    
+    string $pulledObject = getPulledUsdObject($obj);
+    if ($pulledObject != "") {
+        string $pushback = `menuItem -label "Merge Maya Edits to USD" -insertAfter "" -image "merge_to_USD.png" -command ("mayaUsdMenu_pushBackToUSD " + $pulledObject)`;
+        string $pushclear = `menuItem -label "Discard Maya Edits" -insertAfter $pushback -image "discard_edits.png" -command ("mayaUsdDiscardEdits " + $pulledObject)`;
         menuItem -divider true -insertAfter $pushclear;
     }
     else {