From 6251ac8f83c8ba35105ee9faf64ba65ad97874ef Mon Sep 17 00:00:00 2001
From: Ashwin Bhat <ashwin.bhat@autodesk.com>
Date: Wed, 9 Oct 2024 11:59:04 -0700
Subject: [PATCH] Update getUnitDefs() API

---
 source/MaterialXCore/Definition.cpp | 16 +++++++++++++---
 source/MaterialXCore/Document.h     |  4 ++--
 2 files changed, 15 insertions(+), 5 deletions(-)

diff --git a/source/MaterialXCore/Definition.cpp b/source/MaterialXCore/Definition.cpp
index 6b45baf195..cf89f93de6 100644
--- 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
index 6fc9523cca..6c0fbb6149 100644
--- 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.