LOOKDEVX-2403 - Recognize NodeGraph EnumString attributes

lib/mayaUsd/ufe/UsdAttributeHolder.cpp

@@ -555,12 +555,69 @@ PXR_NS::SdfValueTypeName UsdAttributeHolder::usdAttributeType() const
 Ufe::AttributeEnumString::EnumValues UsdAttributeHolder::getEnumValues() const
 {
     Ufe::AttributeEnumString::EnumValues retVal;
+    if (_usdAttr.IsValid()) {
+        for (auto const& option : getEnums()) {
+            retVal.push_back(option.first);
+        }
+    }
+
+    return retVal;
+}
+
+UsdAttributeHolder::EnumOptions UsdAttributeHolder::getEnums() const
+{
+    UsdAttributeHolder::EnumOptions retVal;
     if (_usdAttr.IsValid()) {
         VtTokenArray allowedTokens;
         if (_usdAttr.GetPrim().GetPrimDefinition().GetPropertyMetadata(
                 _usdAttr.GetName(), SdfFieldKeys->AllowedTokens, &allowedTokens)) {
             for (auto const& token : allowedTokens) {
-                retVal.push_back(token.GetString());
+                retVal.emplace_back(token.GetString(), "");
+            }
+        }
+        // We might have a propagated enum copied into the created NodeGraph port, resulting from
+        // connecting a shader enum property.
+        PXR_NS::UsdShadeNodeGraph ngPrim(_usdAttr.GetPrim());
+        if (ngPrim && UsdShadeInput::IsInput(_usdAttr)) {
+            const auto shaderInput = UsdShadeInput { _usdAttr };
+            const auto enumLabels = shaderInput.GetSdrMetadataByKey(TfToken("enum"));
+            const auto enumValues = shaderInput.GetSdrMetadataByKey(TfToken("enumvalues"));
+            const std::vector<std::string> allLabels = splitString(enumLabels, ", ");
+            std::vector<std::string>       allValues = splitString(enumValues, ", ");
+
+            if (!allValues.empty() && allValues.size() != allLabels.size()) {
+                // An array of vector2 values will produce twice the expected number of
+                // elements. We can fix that by regrouping them.
+                if (allValues.size() > allLabels.size()
+                    && allValues.size() % allLabels.size() == 0) {
+
+                    size_t                   stride = allValues.size() / allLabels.size();
+                    std::vector<std::string> rebuiltValues;
+                    std::string              currentValue;
+                    for (size_t i = 0; i < allValues.size(); ++i) {
+                        if (i % stride != 0) {
+                            currentValue += ",";
+                        }
+                        currentValue += allValues[i];
+                        if ((i + 1) % stride == 0) {
+                            rebuiltValues.push_back(currentValue);
+                            currentValue = "";
+                        }
+                    }
+                    allValues.swap(rebuiltValues);
+                } else {
+                    // Can not reconcile the size difference:
+                    allValues.clear();
+                }
+            }
+
+            const bool hasValues = allLabels.size() == allValues.size();
+            for (size_t i = 0; i < allLabels.size(); ++i) {
+                if (hasValues) {
+                    retVal.emplace_back(allLabels[i], allValues[i]);
+                } else {
+                    retVal.emplace_back(allLabels[i], "");
+                }
             }
         }
     }

lib/mayaUsd/ufe/UsdAttributeHolder.h

@@ -29,10 +29,8 @@ namespace ufe {
 //! \brief Internal helper class holding a USD attributes for query:
 class UsdAttributeHolder
 {
-protected:
-    UsdAttributeHolder(const PXR_NS::UsdAttribute& usdAttr);
-
 public:
+    UsdAttributeHolder(const PXR_NS::UsdAttribute& usdAttr);
     typedef std::unique_ptr<UsdAttributeHolder> UPtr;
     static UPtr                                 create(const PXR_NS::UsdAttribute& usdAttr);
     virtual ~UsdAttributeHolder() = default;
@@ -62,6 +60,8 @@ class UsdAttributeHolder
     virtual PXR_NS::UsdAttribute                 usdAttribute() const { return _usdAttr; }
     virtual PXR_NS::SdfValueTypeName             usdAttributeType() const;
     virtual Ufe::AttributeEnumString::EnumValues getEnumValues() const;
+    using EnumOptions = std::vector<std::pair<std::string, std::string>>;
+    virtual EnumOptions getEnums() const;
 
 protected:
     PXR_NS::UsdAttribute _usdAttr;

lib/mayaUsd/ufe/UsdAttributes.cpp

@@ -747,14 +747,19 @@ Ufe::Attribute::Ptr UsdAttributes::doRenameAttribute(
 #ifdef UFE_ATTRIBUTES_GET_ENUMS
 UFE_ATTRIBUTES_BASE::Enums UsdAttributes::getEnums(const std::string& attrName) const
 {
-    UFE_ATTRIBUTES_BASE::Enums        result;
-    PXR_NS::SdrShaderPropertyConstPtr shaderProp = _GetSdrPropertyAndType(fItem, attrName).first;
-    if (shaderProp) {
-        for (const auto& option : shaderProp->GetOptions()) {
-            result.emplace_back(option.first.GetString(), option.second.GetString());
+    auto shaderPropAndType = _GetSdrPropertyAndType(fItem, attrName);
+    if (shaderPropAndType.first) {
+        const auto shaderPropertyHolder = UsdShaderAttributeHolder(
+            fItem->prim(), shaderPropAndType.first, shaderPropAndType.second);
+        return shaderPropertyHolder.getEnums();
+    } else {
+        PXR_NS::UsdAttribute usdAttr = _GetAttributeType(fItem->prim(), attrName);
+        if (usdAttr.IsValid()) {
+            const auto attrHolder = UsdAttributeHolder { usdAttr };
+            return attrHolder.getEnums();
         }
     }
-    return result;
+    return {};
 }
 #endif
 #endif

lib/mayaUsd/ufe/UsdShaderAttributeHolder.cpp

@@ -224,13 +224,22 @@ PXR_NS::SdfValueTypeName UsdShaderAttributeHolder::usdAttributeType() const
 
 Ufe::AttributeEnumString::EnumValues UsdShaderAttributeHolder::getEnumValues() const
 {
-    Ufe::AttributeEnumString::EnumValues retVal = _Base::getEnumValues();
-    for (auto const& option : _sdrProp->GetOptions()) {
+    Ufe::AttributeEnumString::EnumValues retVal;
+    for (auto const& option : getEnums()) {
         retVal.push_back(option.first);
     }
     return retVal;
 }
 
+UsdAttributeHolder::EnumOptions UsdShaderAttributeHolder::getEnums() const
+{
+    auto retVal = _Base::getEnums();
+    for (auto const& option : _sdrProp->GetOptions()) {
+        retVal.emplace_back(option.first, option.second);
+    }
+    return retVal;
+}
+
 void UsdShaderAttributeHolder::_CreateUsdAttribute()
 {
     PXR_NS::UsdShadeShader shader(usdPrim());

lib/mayaUsd/ufe/UsdShaderAttributeHolder.h

@@ -34,12 +34,12 @@ class UsdShaderAttributeHolder : public UsdAttributeHolder
 {
     typedef UsdAttributeHolder _Base;
 
+public:
     UsdShaderAttributeHolder(
         PXR_NS::UsdPrim                   usdPrim,
         PXR_NS::SdrShaderPropertyConstPtr sdrProp,
         PXR_NS::UsdShadeAttributeType     sdrType);
 
-public:
     static UPtr create(
         PXR_NS::UsdPrim                   usdPrim,
         PXR_NS::SdrShaderPropertyConstPtr sdrProp,
@@ -66,6 +66,7 @@ class UsdShaderAttributeHolder : public UsdAttributeHolder
     virtual PXR_NS::UsdPrim                      usdPrim() const { return _usdAttr.GetPrim(); }
     virtual PXR_NS::SdfValueTypeName             usdAttributeType() const;
     virtual Ufe::AttributeEnumString::EnumValues getEnumValues() const;
+    virtual EnumOptions                          getEnums() const;
 
 private:
     PXR_NS::SdrShaderPropertyConstPtr _sdrProp;

lib/mayaUsd/ufe/Utils.cpp

@@ -581,11 +581,20 @@ Ufe::Attribute::Type usdTypeToUfe(const PXR_NS::UsdAttribute& usdAttr)
                     usdAttr.GetName(), SdfFieldKeys->AllowedTokens, nullptr)) {
                 type = Ufe::Attribute::kEnumString;
             }
-            // TfToken is also used in UsdShade as a Generic placeholder for connecting struct I/O.
             UsdShadeNodeGraph asNodeGraph(usdAttr.GetPrim());
-            if (asNodeGraph && usdAttr.GetTypeName() == SdfValueTypeNames->Token) {
-                if (UsdShadeUtils::GetBaseNameAndType(usdAttr.GetName()).second
-                    != UsdShadeAttributeType::Invalid) {
+            if (asNodeGraph) {
+                // NodeGraph inputs can have enum metadata on them when they export an inner enum.
+                const auto portType = UsdShadeUtils::GetBaseNameAndType(usdAttr.GetName()).second;
+                if (portType == UsdShadeAttributeType::Input) {
+                    const auto input = UsdShadeInput(usdAttr);
+                    if (!input.GetSdrMetadataByKey(TfToken("enum")).empty()) {
+                        return Ufe::Attribute::kEnumString;
+                    }
+                }
+                // TfToken is also used in UsdShade as a Generic placeholder for connecting struct
+                // I/O.
+                if (usdAttr.GetTypeName() == SdfValueTypeNames->Token
+                    && portType != UsdShadeAttributeType::Invalid) {
                     type = Ufe::Attribute::kGeneric;
                 }
             }

test/lib/ufe/testConnections.py

@@ -951,35 +951,67 @@ def testCreateNodeGraphAttributes(self):
             testAttrs.addObserver(testItem, testObserver)
             testObserver.assertNotificationCount(self)
 
-            testAttrs.addAttribute("inputs:foo", ufe.Attribute.kFloat)
+            testAttrs.addAttribute("inputs:foo", ufe.Attribute.kFloat4)
             testObserver.assertNotificationCount(self, numAdded = 1)
-            testAttrs.addAttribute("outputs:bar", ufe.Attribute.kFloat4)
+            testAttrs.addAttribute("outputs:bar", ufe.Attribute.kFloat)
             testObserver.assertNotificationCount(self, numAdded = 2)
-
+            testAttrs.addAttribute("inputs:enumString", ufe.Attribute.kString)
+            testObserver.assertNotificationCount(self, numAdded = 3)
+            enumAttr = testAttrs.attribute("inputs:enumString")
+            # Not yet an enum:
+            self.assertEqual(enumAttr.type, ufe.Attribute.kString)
+            # But with proper metadata:
+            enumAttr.setMetadata("enum", "foo, bar, baz")
+            enumAttr = testAttrs.attribute("inputs:enumString")
+            # It is now an enum:
+            self.assertEqual(enumAttr.type, ufe.Attribute.kEnumString)
+            # Test the other enum API:
+            if hasattr(testAttrs, "getEnums"):
+                enums = testAttrs.getEnums("inputs:enumString")
+                self.assertEqual(len(enums), 3)
+                self.assertEqual(len(enums[0]), 2)
+                self.assertEqual(enums[0][0], "foo")
+                self.assertEqual(enums[0][1], "")
+                # Compare against the original API:
+                self.assertEqual([i[0] for i in enums], enumAttr.getEnumValues())
+            testObserver.assertNotificationCount(self, numAdded = 3, numValue=1, numMetadata=1)
+
+            # Enumify the Float4 attribute:
+            expectedEnums = [("X", "1,0,0"), ("Y", "0,1,0"), ("Z", "0,0,1")]
+            enumAttr = testAttrs.attribute("inputs:foo")
+            enumAttr.setMetadata("enum", ", ".join([i[0] for i in expectedEnums]))
+            enumAttr.setMetadata("enumvalues", ", ".join([i[1] for i in expectedEnums]))
+            if hasattr(testAttrs, "getEnums"):
+                enums = testAttrs.getEnums("inputs:foo")
+                self.assertEqual(enums, expectedEnums)
+            testObserver.assertNotificationCount(self, numAdded = 3, numValue=3, numMetadata=3)
+
             # Testing custom NodeGraph data types
             testAttrs.addAttribute("inputs:edf", "EDF")
             # The custom type is saved as metadata, which emits one value and one metadata changes
-            testObserver.assertNotificationCount(self, numAdded = 3, numValue=1, numMetadata=1)
+            testObserver.assertNotificationCount(self, numAdded = 4, numValue=4, numMetadata=4)
             customAttr = testAttrs.attribute("inputs:edf")
             self.assertEqual(customAttr.type, "Generic")
             # Make sure the custom shader type was remembered
             self.assertEqual(customAttr.nativeType(), "EDF")
 
             # Same thing, on the output side
             testAttrs.addAttribute("outputs:srf", "surfaceshader")
-            testObserver.assertNotificationCount(self, numAdded = 4, numValue=2, numMetadata=2)
+            testObserver.assertNotificationCount(self, numAdded = 5, numValue=5, numMetadata=5)
             customAttr = testAttrs.attribute("outputs:srf")
             self.assertEqual(customAttr.type, "Generic")
             self.assertEqual(customAttr.nativeType(), "surfaceshader")
 
             testAttrs.removeAttribute("inputs:foo")
-            testObserver.assertNotificationCount(self, numAdded = 4, numRemoved = 1, numValue=2, numMetadata=2)
+            testObserver.assertNotificationCount(self, numAdded = 5, numRemoved = 1, numValue=5, numMetadata=5)
             testAttrs.removeAttribute("outputs:bar")
-            testObserver.assertNotificationCount(self, numAdded = 4, numRemoved = 2, numValue=2, numMetadata=2)
+            testObserver.assertNotificationCount(self, numAdded = 5, numRemoved = 2, numValue=5, numMetadata=5)
             testAttrs.removeAttribute("inputs:edf")
-            testObserver.assertNotificationCount(self, numAdded = 4, numRemoved = 3, numValue=2, numMetadata=2)
+            testObserver.assertNotificationCount(self, numAdded = 5, numRemoved = 3, numValue=5, numMetadata=5)
             testAttrs.removeAttribute("outputs:srf")
-            testObserver.assertNotificationCount(self, numAdded = 4, numRemoved = 4, numValue=2, numMetadata=2)
+            testObserver.assertNotificationCount(self, numAdded = 5, numRemoved = 4, numValue=5, numMetadata=5)
+            testAttrs.removeAttribute("inputs:enumString")
+            testObserver.assertNotificationCount(self, numAdded = 5, numRemoved = 5, numValue=5, numMetadata=5)
 
 
     @unittest.skipUnless(ufeUtils.ufeFeatureSetVersion() >= 4, 'Test only available in UFE v4 or greater')