Update getUnitDefs() API

autodesk-forks · Oct 9, 2024 · 6251ac8 · 6251ac8
1 parent 3dd6999
commit 6251ac8
Show file tree

Hide file tree

Showing 2 changed files with 15 additions and 5 deletions.
diff --git a/source/MaterialXCore/Definition.cpp b/source/MaterialXCore/Definition.cpp
@@ -185,16 +185,26 @@ StringVec TargetDef::getMatchingTargets() const
 
 vector<UnitDefPtr> UnitTypeDef::getUnitDefs() const
 {
-    const auto datalibrary = getDocument()->hasDataLibrary() ? getDocument()->getRegisteredDataLibrary() : getDocument();
-
     vector<UnitDefPtr> unitDefs;
-    for (UnitDefPtr unitDef : datalibrary->getChildrenOfType<UnitDef>())
+    for (UnitDefPtr unitDef : getDocument()->getChildrenOfType<UnitDef>())
     {
         if (unitDef->getUnitType() == _name)
         {
             unitDefs.push_back(unitDef);
         }
     }
+    // Gather unitdefs from Data library
+    if (getDocument()->hasDataLibrary())
+    {
+        for (UnitDefPtr unitDef : getDocument()->getRegisteredDataLibrary()->getChildrenOfType<UnitDef>())
+        {
+            if (unitDef->getUnitType() == _name)
+            {
+                unitDefs.push_back(unitDef);
+            }
+        }
+
+    }
     return unitDefs;
 }
 

diff --git a/source/MaterialXCore/Document.h b/source/MaterialXCore/Document.h
@@ -554,7 +554,7 @@ class MX_CORE_API Document : public GraphElement
     /// Return a vector of all Member elements in the TypeDef.
     vector<UnitDefPtr> getUnitDefs() const
     {
-        return getChildrenOfType<UnitDef>();
+        return hasDataLibrary() ? getChildrenOfType<UnitDef>(getRegisteredDataLibrary()) : getChildrenOfType<UnitDef>();
     }
 
     /// Remove the UnitDef, if any, with the given name.
@@ -585,7 +585,7 @@ class MX_CORE_API Document : public GraphElement
     /// Return a vector of all UnitTypeDef elements in the document.
     vector<UnitTypeDefPtr> getUnitTypeDefs() const
     {
-        return getChildrenOfType<UnitTypeDef>();
+        return hasDataLibrary() ? getChildrenOfType<UnitTypeDef>(getRegisteredDataLibrary()) : getChildrenOfType<UnitTypeDef>();
     }
 
     /// Remove the UnitTypeDef, if any, with the given name.