diff --git a/src/interpreter.cpp b/src/interpreter.cpp
index c7f6644..22c1824 100644
--- a/src/interpreter.cpp
+++ b/src/interpreter.cpp
@@ -100,12 +100,12 @@
         auto leafType = m_datastore.schema()->leafType(path);
         auto typedefName = m_datastore.schema()->leafTypeName(path);
         std::string baseTypeStr;
-        if (std::holds_alternative<yang::LeafRef>(leafType)) {
+        if (std::holds_alternative<yang::LeafRef>(leafType.m_type)) {
             ss << "-> ";
             ss << m_datastore.schema()->leafrefPath(path) << " ";
-            baseTypeStr = leafDataTypeToString(*std::get<yang::LeafRef>(leafType).m_targetType);
+            baseTypeStr = leafDataTypeToString(std::get<yang::LeafRef>(leafType.m_type).m_targetType->m_type);
         } else {
-            baseTypeStr = leafDataTypeToString(leafType);
+            baseTypeStr = leafDataTypeToString(leafType.m_type);
         }
 
         if (typedefName) {
@@ -114,8 +114,8 @@
             ss << baseTypeStr;
         }
 
-        if (auto units = m_datastore.schema()->units(path)) {
-            ss << " [" + *units + "]";
+        if (leafType.m_units) {
+            ss << " [" + *leafType.m_units + "]";
         }
 
         if (m_datastore.schema()->leafIsKey(path)) {
diff --git a/src/leaf_data.hpp b/src/leaf_data.hpp
index 1c4176e..68705d0 100644
--- a/src/leaf_data.hpp
+++ b/src/leaf_data.hpp
@@ -171,12 +171,12 @@
     }
     bool operator()(const yang::LeafRef& leafRef) const
     {
-        return std::visit(*this, *leafRef.m_targetType);
+        return std::visit(*this, leafRef.m_targetType->m_type);
     }
     bool operator()(const yang::Union& unionInfo) const
     {
         return std::any_of(unionInfo.m_unionTypes.begin(), unionInfo.m_unionTypes.end(), [this](const auto& type) {
-            return std::visit(*this, type);
+            return std::visit(*this, type.m_type);
         });
     }
 };
@@ -190,14 +190,13 @@
     {
         ParserContext& parserContext = x3::get<parser_context_tag>(ctx);
         const Schema& schema = parserContext.m_schema;
-        auto type = schema.leafType(parserContext.m_tmpListKeyLeafPath.m_location, parserContext.m_tmpListKeyLeafPath.m_node);
+        auto type = schema.leafType(parserContext.m_tmpListKeyLeafPath.m_location, parserContext.m_tmpListKeyLeafPath.m_node).m_type;
 
         auto pass = std::visit(impl_LeafData<It, Ctx, RCtx, Attr>{first, last, ctx, rctx, attr, parserContext}, type);
 
         if (!pass) {
             if (parserContext.m_errorMsg.empty()) {
-                parserContext.m_errorMsg = "leaf data type mismatch: Expected " +
-                    leafDataTypeToString(schema.leafType(parserContext.m_tmpListKeyLeafPath.m_location, parserContext.m_tmpListKeyLeafPath.m_node)) + " here:";
+                parserContext.m_errorMsg = "leaf data type mismatch: Expected " + leafDataTypeToString(type) + " here:";
             }
         }
         return pass;
diff --git a/src/leaf_data_type.cpp b/src/leaf_data_type.cpp
index 8048efa..29fb355 100644
--- a/src/leaf_data_type.cpp
+++ b/src/leaf_data_type.cpp
@@ -8,6 +8,15 @@
 #include "leaf_data_type.hpp"
 
 namespace yang {
+bool TypeInfo::operator==(const TypeInfo& other) const
+{
+    return this->m_type == other.m_type && this->m_units == other.m_units;
+}
+TypeInfo::TypeInfo(const yang::LeafDataType& type, const std::optional<std::string> units)
+    : m_type(type)
+    , m_units(units)
+{
+}
 Enum::Enum(std::set<enum_>&& values)
     : m_allowedValues(std::move(values))
 {
@@ -27,10 +36,10 @@
 // Copy constructor needed, because unique_ptr is not copy-constructible
 LeafRef::LeafRef(const LeafRef& src)
     : m_targetXPath(src.m_targetXPath)
-    , m_targetType(std::make_unique<LeafDataType>(*src.m_targetType))
+    , m_targetType(std::make_unique<TypeInfo>(*src.m_targetType))
 {
 }
-LeafRef::LeafRef(const std::string& xpath, std::unique_ptr<LeafDataType>&& type)
+LeafRef::LeafRef(const std::string& xpath, std::unique_ptr<TypeInfo>&& type)
     : m_targetXPath(xpath)
     , m_targetType(std::move(type))
 {
diff --git a/src/leaf_data_type.hpp b/src/leaf_data_type.hpp
index a86691b..470c435 100644
--- a/src/leaf_data_type.hpp
+++ b/src/leaf_data_type.hpp
@@ -82,16 +82,23 @@
     yang::LeafRef,
     yang::Union
 >;
+struct TypeInfo;
 struct LeafRef {
     LeafRef(const LeafRef& src);
-    LeafRef(const std::string& xpath, std::unique_ptr<LeafDataType>&& type);
+    LeafRef(const std::string& xpath, std::unique_ptr<TypeInfo>&& type);
     bool operator==(const LeafRef& other) const;
     std::string m_targetXPath;
-    std::unique_ptr<LeafDataType> m_targetType;
+    std::unique_ptr<TypeInfo> m_targetType;
 };
 
 struct Union {
     bool operator==(const Union& other) const;
-    std::vector<LeafDataType> m_unionTypes;
+    std::vector<TypeInfo> m_unionTypes;
+};
+struct TypeInfo {
+    TypeInfo(const yang::LeafDataType& type, const std::optional<std::string> units = std::nullopt);
+    bool operator==(const TypeInfo& other) const;
+    yang::LeafDataType m_type;
+    std::optional<std::string> m_units;
 };
 }
diff --git a/src/schema.hpp b/src/schema.hpp
index 9a86045..8cdffa1 100644
--- a/src/schema.hpp
+++ b/src/schema.hpp
@@ -55,12 +55,11 @@
     virtual bool listHasKey(const schemaPath_& location, const ModuleNodePair& node, const std::string& key) const = 0;
     virtual bool leafIsKey(const std::string& leafPath) const = 0;
     virtual const std::set<std::string> listKeys(const schemaPath_& location, const ModuleNodePair& node) const = 0;
-    virtual yang::LeafDataType leafType(const schemaPath_& location, const ModuleNodePair& node) const = 0;
-    virtual yang::LeafDataType leafType(const std::string& path) const = 0;
+    virtual yang::TypeInfo leafType(const schemaPath_& location, const ModuleNodePair& node) const = 0;
+    virtual yang::TypeInfo leafType(const std::string& path) const = 0;
     virtual std::optional<std::string> leafTypeName(const std::string& path) const = 0;
     virtual std::string leafrefPath(const std::string& leafrefPath) const = 0;
     virtual std::optional<std::string> description(const std::string& location) const = 0;
-    virtual std::optional<std::string> units(const std::string& location) const = 0;
 
     virtual std::set<std::string> childNodes(const schemaPath_& path, const Recursion recursion) const = 0;
     virtual std::set<std::string> moduleNodes(const module_& module, const Recursion recursion) const = 0;
diff --git a/src/static_schema.cpp b/src/static_schema.cpp
index 3eb42f6..6815015 100644
--- a/src/static_schema.cpp
+++ b/src/static_schema.cpp
@@ -86,7 +86,7 @@
 
 void StaticSchema::addLeaf(const std::string& location, const std::string& name, const yang::LeafDataType& type)
 {
-    m_nodes.at(location).emplace(name, yang::leaf{type});
+    m_nodes.at(location).emplace(name, yang::leaf{yang::TypeInfo{type, std::nullopt}});
     std::string key = joinPaths(location, name);
     m_nodes.emplace(key, std::unordered_map<std::string, NodeType>());
 }
@@ -122,13 +122,13 @@
     return res;
 }
 
-yang::LeafDataType StaticSchema::leafType(const schemaPath_& location, const ModuleNodePair& node) const
+yang::TypeInfo StaticSchema::leafType(const schemaPath_& location, const ModuleNodePair& node) const
 {
     std::string locationString = pathToSchemaString(location, Prefixes::Always);
     return boost::get<yang::leaf>(children(locationString).at(fullNodeName(location, node))).m_type;
 }
 
-yang::LeafDataType StaticSchema::leafType(const std::string& path) const
+yang::TypeInfo StaticSchema::leafType(const std::string& path) const
 {
     auto locationString = stripLastNodeFromPath(path);
     auto node = lastNodeOfSchemaPath(path);
@@ -197,11 +197,6 @@
     throw std::runtime_error{"StaticSchema::description not implemented"};
 }
 
-std::optional<std::string> StaticSchema::units([[maybe_unused]] const std::string& path) const
-{
-    throw std::runtime_error{"StaticSchema::units not implemented"};
-}
-
 yang::NodeTypes StaticSchema::nodeType([[maybe_unused]] const std::string& path) const
 {
     throw std::runtime_error{"Internal error: StaticSchema::nodeType(std::string) not implemented. The tests should not have called this overload."};
diff --git a/src/static_schema.hpp b/src/static_schema.hpp
index a94c78d..8b160b0 100644
--- a/src/static_schema.hpp
+++ b/src/static_schema.hpp
@@ -27,7 +27,7 @@
 };
 
 struct leaf {
-    yang::LeafDataType m_type;
+    yang::TypeInfo m_type;
 };
 
 struct module {
@@ -51,14 +51,13 @@
     bool listHasKey(const schemaPath_& location, const ModuleNodePair& node, const std::string& key) const override;
     bool leafIsKey(const std::string& leafPath) const override;
     const std::set<std::string> listKeys(const schemaPath_& location, const ModuleNodePair& node) const override;
-    yang::LeafDataType leafType(const schemaPath_& location, const ModuleNodePair& node) const override;
-    yang::LeafDataType leafType(const std::string& path) const override;
+    yang::TypeInfo leafType(const schemaPath_& location, const ModuleNodePair& node) const override;
+    yang::TypeInfo leafType(const std::string& path) const override;
     std::optional<std::string> leafTypeName(const std::string& path) const override;
     std::string leafrefPath(const std::string& leafrefPath) const override;
     std::set<std::string> childNodes(const schemaPath_& path, const Recursion) const override;
     std::set<std::string> moduleNodes(const module_& module, const Recursion recursion) const override;
     std::optional<std::string> description(const std::string& path) const override;
-    std::optional<std::string> units(const std::string& path) const override;
 
     /** A helper for making tests a little bit easier. It returns all
      * identities which are based on the argument passed and which can then be
diff --git a/src/utils.cpp b/src/utils.cpp
index 31854b2..9496db3 100644
--- a/src/utils.cpp
+++ b/src/utils.cpp
@@ -118,7 +118,7 @@
     {
         std::ostringstream ss;
         std::transform(type.m_unionTypes.begin(), type.m_unionTypes.end(), std::experimental::make_ostream_joiner(ss, ", "), [this](const auto& unionType) {
-            return std::visit(*this, unionType);
+            return std::visit(*this, unionType.m_type);
         });
         return ss.str();
     }
diff --git a/src/yang_schema.cpp b/src/yang_schema.cpp
index a453b50..5a34e4f 100644
--- a/src/yang_schema.cpp
+++ b/src/yang_schema.cpp
@@ -201,65 +201,99 @@
 }
 }
 
-yang::LeafDataType YangSchema::impl_leafType(const libyang::S_Schema_Node& node) const
+yang::TypeInfo YangSchema::impl_leafType(const libyang::S_Schema_Node& node) const
 {
     using namespace std::string_literals;
     auto leaf = std::make_shared<libyang::Schema_Node_Leaf>(node);
-    std::function<yang::LeafDataType(std::shared_ptr<libyang::Type>)> resolveType;
-    resolveType = [this, &resolveType, leaf] (auto type) -> yang::LeafDataType {
+    auto leafUnits = leaf->units();
+    std::function<yang::TypeInfo(std::shared_ptr<libyang::Type>)> resolveType;
+    resolveType = [this, &resolveType, leaf, leafUnits] (std::shared_ptr<libyang::Type> type) -> yang::TypeInfo {
+        yang::LeafDataType resType;
         switch (type->base()) {
         case LY_TYPE_STRING:
-            return yang::String{};
+            resType.emplace<yang::String>();
+            break;
         case LY_TYPE_DEC64:
-            return yang::Decimal{};
+            resType.emplace<yang::Decimal>();
+            break;
         case LY_TYPE_BOOL:
-            return yang::Bool{};
+            resType.emplace<yang::Bool>();
+            break;
         case LY_TYPE_INT8:
-            return yang::Int8{};
+            resType.emplace<yang::Int8>();
+            break;
         case LY_TYPE_INT16:
-            return yang::Int16{};
+            resType.emplace<yang::Int16>();
+            break;
         case LY_TYPE_INT32:
-            return yang::Int32{};
+            resType.emplace<yang::Int32>();
+            break;
         case LY_TYPE_INT64:
-            return yang::Int64{};
+            resType.emplace<yang::Int64>();
+            break;
         case LY_TYPE_UINT8:
-            return yang::Uint8{};
+            resType.emplace<yang::Uint8>();
+            break;
         case LY_TYPE_UINT16:
-            return yang::Uint16{};
+            resType.emplace<yang::Uint16>();
+            break;
         case LY_TYPE_UINT32:
-            return yang::Uint32{};
+            resType.emplace<yang::Uint32>();
+            break;
         case LY_TYPE_UINT64:
-            return yang::Uint64{};
+            resType.emplace<yang::Uint64>();
+            break;
         case LY_TYPE_BINARY:
-            return yang::Binary{};
+            resType.emplace<yang::Binary>();
+            break;
         case LY_TYPE_ENUM:
-            return yang::Enum{enumValues(type)};
+            resType.emplace<yang::Enum>(enumValues(type));
+            break;
         case LY_TYPE_IDENT:
-            return yang::IdentityRef{validIdentities(type)};
+            resType.emplace<yang::IdentityRef>(validIdentities(type));
+            break;
         case LY_TYPE_LEAFREF:
-            return yang::LeafRef{::leafrefPath(type), std::make_unique<yang::LeafDataType>(leafType(::leafrefPath(type)))};
+            resType.emplace<yang::LeafRef>(::leafrefPath(type), std::make_unique<yang::TypeInfo>(leafType(::leafrefPath(type))));
+            break;
         case LY_TYPE_UNION:
             {
-                auto res = yang::Union{};
+                auto resUnion = yang::Union{};
                 for (auto unionType : type->info()->uni()->types()) {
-                    res.m_unionTypes.push_back(resolveType(unionType));
+                    resUnion.m_unionTypes.push_back(resolveType(unionType));
                 }
-                return res;
+                resType.emplace<yang::Union>(std::move(resUnion));
+                break;
             }
         default:
             using namespace std::string_literals;
             throw UnsupportedYangTypeException("the type of "s + leaf->name() + " is not supported: " + std::to_string(leaf->type()->base()));
         }
-        };
+
+        std::optional<std::string> resUnits;
+
+        if (leafUnits) {
+            resUnits = leafUnits;
+        } else {
+            for (auto parentTypedef = type->der(); parentTypedef; parentTypedef = parentTypedef->type()->der()) {
+                auto units = parentTypedef->units();
+                if (units) {
+                    resUnits = units;
+                    break;
+                }
+            }
+        }
+
+        return yang::TypeInfo(resType, resUnits);
+    };
     return resolveType(leaf->type());
 }
 
-yang::LeafDataType YangSchema::leafType(const schemaPath_& location, const ModuleNodePair& node) const
+yang::TypeInfo YangSchema::leafType(const schemaPath_& location, const ModuleNodePair& node) const
 {
     return impl_leafType(getSchemaNode(location, node));
 }
 
-yang::LeafDataType YangSchema::leafType(const std::string& path) const
+yang::TypeInfo YangSchema::leafType(const std::string& path) const
 {
     return impl_leafType(getSchemaNode(path));
 }
@@ -421,28 +455,3 @@
     auto node = getSchemaNode(path.c_str());
     return node->dsc() ? std::optional{node->dsc()} : std::nullopt;
 }
-
-std::optional<std::string> YangSchema::units(const std::string& path) const
-{
-    auto node = getSchemaNode(path.c_str());
-    if (node->nodetype() != LYS_LEAF) {
-        return std::nullopt;
-    }
-    libyang::Schema_Node_Leaf leaf{node};
-    auto units = leaf.units();
-
-    // A leaf can specify units as part of its definition.
-    if (units) {
-        return units;
-    }
-
-    // A typedef (or its parent typedefs) can specify units too. We'll use the first `units` we find.
-    for (auto parentTypedef = leaf.type()->der(); parentTypedef; parentTypedef = parentTypedef->type()->der()) {
-        units = parentTypedef->units();
-        if (units) {
-            return units;
-        }
-    }
-
-    return std::nullopt;
-}
diff --git a/src/yang_schema.hpp b/src/yang_schema.hpp
index e3a9025..a2981ed 100644
--- a/src/yang_schema.hpp
+++ b/src/yang_schema.hpp
@@ -37,14 +37,13 @@
     bool listHasKey(const schemaPath_& location, const ModuleNodePair& node, const std::string& key) const override;
     bool leafIsKey(const std::string& leafPath) const override;
     const std::set<std::string> listKeys(const schemaPath_& location, const ModuleNodePair& node) const override;
-    yang::LeafDataType leafType(const schemaPath_& location, const ModuleNodePair& node) const override;
-    yang::LeafDataType leafType(const std::string& path) const override;
+    yang::TypeInfo leafType(const schemaPath_& location, const ModuleNodePair& node) const override;
+    yang::TypeInfo leafType(const std::string& path) const override;
     std::optional<std::string> leafTypeName(const std::string& path) const override;
     std::string leafrefPath(const std::string& leafrefPath) const override;
     std::set<std::string> childNodes(const schemaPath_& path, const Recursion recursion) const override;
     std::set<std::string> moduleNodes(const module_& module, const Recursion recursion) const override;
     std::optional<std::string> description(const std::string& path) const override;
-    std::optional<std::string> units(const std::string& path) const override;
 
     void registerModuleCallback(const std::function<std::string(const char*, const char*, const char*, const char*)>& clb);
 
@@ -68,7 +67,7 @@
     std::shared_ptr<libyang::Module> getYangModule(const std::string& name);
 
 private:
-    yang::LeafDataType impl_leafType(const std::shared_ptr<libyang::Schema_Node>& node) const;
+    yang::TypeInfo impl_leafType(const std::shared_ptr<libyang::Schema_Node>& node) const;
     std::set<std::string> modules() const;
 
     /** @short Returns a set of nodes, that match the location and name criteria. */
diff --git a/tests/leaf_editing.cpp b/tests/leaf_editing.cpp
index ad88a99..64abbf2 100644
--- a/tests/leaf_editing.cpp
+++ b/tests/leaf_editing.cpp
@@ -50,26 +50,26 @@
     schema->addList("/", "mod:list", {"number"});
     schema->addLeaf("/mod:list", "mod:number", yang::Int32{});
     schema->addLeaf("/mod:list", "mod:leafInList", yang::String{});
-    schema->addLeaf("/", "mod:refToString", yang::LeafRef{"/mod:leafString", std::make_unique<yang::LeafDataType>(schema->leafType("/mod:leafString"))});
-    schema->addLeaf("/", "mod:refToInt8", yang::LeafRef{"/mod:leafInt8", std::make_unique<yang::LeafDataType>(schema->leafType("/mod:leafInt8"))});
-    schema->addLeaf("/", "mod:refToLeafInCont", yang::LeafRef{"/mod:contA/identInCont", std::make_unique<yang::LeafDataType>(schema->leafType("/mod:contA/mod:identInCont"))});
-    schema->addLeaf("/", "mod:intOrString", yang::Union{{yang::Int32{}, yang::String{}}});
-    schema->addLeaf("/", "mod:twoInts", yang::Union{{yang::Uint8{}, yang::Int16{}}});
+    schema->addLeaf("/", "mod:refToString", yang::LeafRef{"/mod:leafString", std::make_unique<yang::TypeInfo>(schema->leafType("/mod:leafString"))});
+    schema->addLeaf("/", "mod:refToInt8", yang::LeafRef{"/mod:leafInt8", std::make_unique<yang::TypeInfo>(schema->leafType("/mod:leafInt8"))});
+    schema->addLeaf("/", "mod:refToLeafInCont", yang::LeafRef{"/mod:contA/identInCont", std::make_unique<yang::TypeInfo>(schema->leafType("/mod:contA/mod:identInCont"))});
+    schema->addLeaf("/", "mod:intOrString", yang::Union{{yang::TypeInfo{yang::Int32{}}, yang::TypeInfo{yang::String{}}}});
+    schema->addLeaf("/", "mod:twoInts", yang::Union{{yang::TypeInfo{yang::Uint8{}}, yang::TypeInfo{yang::Int16{}}}});
     schema->addLeaf("/", "mod:unionStringEnumLeafref", yang::Union{{
-        yang::String{},
-        createEnum({"foo", "bar"}),
-        yang::LeafRef{"/mod:leafEnum", std::make_unique<yang::LeafDataType>(schema->leafType("/mod:leafEnum"))}
+            yang::LeafDataType{yang::String{}},
+        yang::LeafDataType{createEnum({"foo", "bar"})},
+        yang::LeafDataType{yang::LeafRef{"/mod:leafEnum", std::make_unique<yang::TypeInfo>(schema->leafType("/mod:leafEnum"))}}
     }});
 
     schema->addList("/", "mod:portSettings", {"port"});
     schema->addLeaf("/mod:portSettings", "mod:port", createEnum({"eth0", "eth1", "eth2"}));
     schema->addList("/", "mod:portMapping", {"port"});
     schema->addLeaf("/mod:portMapping", "mod:port", createEnum({"utf1", "utf2", "utf3"}));
-    schema->addLeaf("/", "mod:activeMappedPort", yang::LeafRef{"/mod:portMapping/mod:port", std::make_unique<yang::LeafDataType>(schema->leafType("/mod:portMapping/mod:port"))});
+    schema->addLeaf("/", "mod:activeMappedPort", yang::LeafRef{"/mod:portMapping/mod:port", std::make_unique<yang::TypeInfo>(schema->leafType("/mod:portMapping/mod:port"))});
     schema->addLeaf("/", "mod:activePort", yang::Union{{
-        createEnum({"wlan0", "wlan1"}),
-        yang::LeafRef{"/mod:portSettings/mod:port", std::make_unique<yang::LeafDataType>(schema->leafType("/mod:portSettings/mod:port"))},
-        yang::LeafRef{"/mod:activeMappedPort", std::make_unique<yang::LeafDataType>(schema->leafType("/mod:activeMappedPort"))},
+        yang::TypeInfo{createEnum({"wlan0", "wlan1"})},
+        yang::TypeInfo{yang::LeafRef{"/mod:portSettings/mod:port", std::make_unique<yang::TypeInfo>(schema->leafType("/mod:portSettings/mod:port"))}},
+        yang::TypeInfo{yang::LeafRef{"/mod:activeMappedPort", std::make_unique<yang::TypeInfo>(schema->leafType("/mod:activeMappedPort"))}},
     }});
 
     Parser parser(schema);
diff --git a/tests/pretty_printers.hpp b/tests/pretty_printers.hpp
index e5b6659..98638ea 100644
--- a/tests/pretty_printers.hpp
+++ b/tests/pretty_printers.hpp
@@ -64,14 +64,22 @@
         s << "}";
     }
     if (std::holds_alternative<yang::LeafRef>(type)) {
-        s << "{" << std::get<yang::LeafRef>(type).m_targetXPath << "," << *std::get<yang::LeafRef>(type).m_targetType  << "}";
+        s << "{" << std::get<yang::LeafRef>(type).m_targetXPath << "," << std::get<yang::LeafRef>(type).m_targetType->m_type  << "}";
     }
     if (std::holds_alternative<yang::Union>(type)) {
         s << "{" << std::endl;
         auto types = std::get<yang::Union>(type).m_unionTypes;
-        std::copy(types.begin(), types.end(), std::experimental::make_ostream_joiner(s, ",\n"));
+        std::transform(types.begin(), types.end(), std::experimental::make_ostream_joiner(s, ",\n"), [] (const auto& type) {
+            return type.m_type;
+        });
     }
     s << std::endl;
     return s;
 }
+
+std::ostream& operator<<(std::ostream& s, const yang::TypeInfo& type)
+{
+    s << type.m_type << (type.m_units ? " units: " + *type.m_units : "");
+    return s;
+}
 }
diff --git a/tests/utils.cpp b/tests/utils.cpp
index cea698b..509c2fe 100644
--- a/tests/utils.cpp
+++ b/tests/utils.cpp
@@ -89,10 +89,10 @@
         SECTION("union")
         {
             type = yang::Union{{
-                yang::String{},
-                createEnum({"foo", "bar"}),
-                yang::Int8{},
-                yang::Int64{},
+                yang::TypeInfo{yang::String{}},
+                yang::TypeInfo{createEnum({"foo", "bar"})},
+                yang::TypeInfo{yang::Int8{}},
+                yang::TypeInfo{yang::Int64{}},
             }};
             expected = "a string, an enum, an 8-bit integer, a 64-bit integer";
         }
diff --git a/tests/yang.cpp b/tests/yang.cpp
index a9fb65e..941586d 100644
--- a/tests/yang.cpp
+++ b/tests/yang.cpp
@@ -695,7 +695,7 @@
                 node.second = "activeNumber";
                 type.emplace<yang::LeafRef>(
                     "/example-schema:_list/number",
-                    std::make_unique<yang::LeafDataType>(ys.leafType("/example-schema:_list/number"))
+                    std::make_unique<yang::TypeInfo>(ys.leafType("/example-schema:_list/number"))
                 );
             }
 
@@ -718,17 +718,18 @@
                 }();
 
                 type = yang::Union{{
-                    createEnum({"wlan0", "wlan1"}),
-                    yang::LeafRef{
+                    yang::TypeInfo{createEnum({"wlan0", "wlan1"})},
+                    yang::TypeInfo{yang::LeafRef{
                         "/example-schema:portSettings/port",
-                        std::make_unique<yang::LeafDataType>(createEnum({"eth0", "eth1", "eth2"}))
-                    },
-                    yang::LeafRef{
+                        std::make_unique<yang::TypeInfo>(createEnum({"eth0", "eth1", "eth2"}))
+                    }},
+                    yang::TypeInfo{yang::LeafRef{
                         "/example-schema:activeMappedPort",
-                        std::make_unique<yang::LeafDataType>(yang::LeafRef{
+                        std::make_unique<yang::TypeInfo>(yang::LeafRef{
                                 "/example-schema:portMapping/port",
-                                std::make_unique<yang::LeafDataType>(enums)
-                        })},
+                                std::make_unique<yang::TypeInfo>(enums)
+                        })
+                    }},
                 }};
             }
 
@@ -841,37 +842,42 @@
 
         SECTION("units")
         {
-            std::optional<std::string> expected;
+            yang::LeafDataType expectedType;
+            std::optional<std::string> expectedUnits;
             SECTION("length")
             {
                 path.m_nodes.push_back(schemaNode_(module_{"example-schema"}, leaf_("length")));
-                expected = "m";
+                expectedType.emplace<yang::Int32>();
+                expectedUnits = "m";
             }
 
             SECTION("wavelength")
             {
                 path.m_nodes.push_back(schemaNode_(module_{"example-schema"}, leaf_("wavelength")));
-                expected = "nm";
+                expectedType.emplace<yang::Decimal>();
+                expectedUnits = "nm";
             }
 
             SECTION("leafInt32")
             {
                 path.m_nodes.push_back(schemaNode_(module_{"example-schema"}, leaf_("leafInt32")));
+                expectedType.emplace<yang::Int32>();
             }
 
             SECTION("duration")
             {
                 path.m_nodes.push_back(schemaNode_(module_{"example-schema"}, leaf_("duration")));
-                expected = "s";
+                expectedType.emplace<yang::Int32>();
+                expectedUnits = "s";
             }
 
             SECTION("another-duration")
             {
                 path.m_nodes.push_back(schemaNode_(module_{"example-schema"}, leaf_("another-duration")));
-                expected = "vt";
+                expectedType.emplace<yang::Int32>();
+                expectedUnits = "vt";
             }
-
-            REQUIRE(ys.units(pathToSchemaString(path, Prefixes::WhenNeeded)) == expected);
+            REQUIRE(ys.leafType(pathToSchemaString(path, Prefixes::WhenNeeded)) == yang::TypeInfo{expectedType, expectedUnits});
         }
 
         SECTION("nodeType")