diff --git a/src/common/model/element.cc b/src/common/model/element.cc
index ec1b2549..00fbc612 100644
--- a/src/common/model/element.cc
+++ b/src/common/model/element.cc
@@ -66,4 +66,9 @@ const std::vector<relationship> &element::relationships() const
 }
 
 void element::append(const element &e) { decorated_element::append(e); }
+
+bool operator==(const element &l, const element &r)
+{
+    return l.full_name(false) == r.full_name(false);
+}
 }
diff --git a/src/common/model/element.h b/src/common/model/element.h
index 6704a124..6206893c 100644
--- a/src/common/model/element.h
+++ b/src/common/model/element.h
@@ -30,6 +30,8 @@ class element : public decorated_element {
 public:
     element(const std::vector<std::string> &using_namespaces);
 
+    virtual ~element() = default;
+
     std::string alias() const;
 
     void set_name(const std::string &name) { name_ = name; }
@@ -54,6 +56,8 @@ public:
 
     void append(const element &e);
 
+    friend bool operator==(const element &l, const element &r);
+
 protected:
     const uint64_t m_id{0};
 
diff --git a/src/common/model/nested_trait.h b/src/common/model/nested_trait.h
index 37644d12..55f0e2aa 100644
--- a/src/common/model/nested_trait.h
+++ b/src/common/model/nested_trait.h
@@ -37,7 +37,7 @@ public:
 
     virtual ~nested_trait() = default;
 
-    void add_element(std::unique_ptr<T> p)
+    template <typename V = T> void add_element(std::unique_ptr<V> p)
     {
         auto it = std::find_if(elements_.begin(), elements_.end(),
             [&p](const auto &e) { return *e == *p; });
@@ -50,7 +50,8 @@ public:
         }
     }
 
-    void add_element(std::vector<std::string> path, std::unique_ptr<T> p)
+    template <typename V = T>
+    void add_element(std::vector<std::string> path, std::unique_ptr<V> p)
     {
         assert(p);
 
@@ -64,42 +65,55 @@ public:
 
         auto parent = get_element(path);
 
-        if (parent)
-            parent.value().add_element(std::move(p));
+        if (parent && dynamic_cast<nested_trait<T> *>(&parent.value()))
+            dynamic_cast<nested_trait<T> &>(parent.value())
+                .template add_element<V>(std::move(p));
         else
             spdlog::error(
                 "No parent element found at: {}", fmt::join(path, "::"));
     }
 
-    type_safe::optional_ref<T> get_element(std::vector<std::string> path) const
+    template <typename V = T>
+    auto get_element(std::vector<std::string> path) const
     {
         LOG_DBG("Getting nested element at path: {}", fmt::join(path, "::"));
 
         if (path.empty() || !has_element(path.at(0))) {
             LOG_WARN("Nested element {} not found in element",
                 fmt::join(path, "::"));
-            return {};
+            return type_safe::optional_ref<V>{};
         }
 
-        auto p = get_element(path.at(0));
         if (path.size() == 1)
-            return p;
+            return get_element<V>(path.at(0));
 
-        return p.value().get_element(
-            std::vector<std::string>(path.begin() + 1, path.end()));
+        auto p = get_element<T>(path.at(0));
+
+        if (!p)
+            return type_safe::optional_ref<V>{};
+
+        if (dynamic_cast<nested_trait<T> *>(&p.value()))
+            return dynamic_cast<nested_trait<T> &>(p.value()).get_element<V>(
+                std::vector<std::string>(path.begin() + 1, path.end()));
+
+        return type_safe::optional_ref<V>{};
     }
 
-    type_safe::optional_ref<T> get_element(const std::string &name) const
+    template <typename V = T> auto get_element(const std::string &name) const
     {
         auto it = std::find_if(elements_.cbegin(), elements_.cend(),
             [&](const auto &p) { return name == p->name(); });
 
         if (it == elements_.end())
-            return {};
+            return type_safe::optional_ref<V>{};
 
         assert(it->get() != nullptr);
 
-        return type_safe::ref(*(it->get()));
+        if (dynamic_cast<V *>(it->get()))
+            return type_safe::optional_ref<V>{
+                type_safe::ref<V>(dynamic_cast<V &>(*it->get()))};
+
+        return type_safe::optional_ref<V>{};
     }
 
     bool has_element(const std::string &name) const
diff --git a/src/common/model/package.cc b/src/common/model/package.cc
index c000f843..9d12601d 100644
--- a/src/common/model/package.cc
+++ b/src/common/model/package.cc
@@ -43,11 +43,6 @@ std::string package::full_name(bool relative) const
     return fmt::format("{}", fmt::join(fn, "::"));
 }
 
-bool operator==(const package &l, const package &r)
-{
-    return l.full_name(false) == r.full_name(false);
-}
-
 bool package::is_deprecated() const { return is_deprecated_; }
 
 void package::set_deprecated(bool deprecated) { is_deprecated_ = deprecated; }
diff --git a/src/common/model/package.h b/src/common/model/package.h
index fbab1b2c..e11fafdb 100644
--- a/src/common/model/package.h
+++ b/src/common/model/package.h
@@ -31,9 +31,9 @@
 
 namespace clanguml::common::model {
 
-class package : public common::model::element,
-                public common::model::stylable_element,
-                public common::model::nested_trait<package> {
+class package : public element,
+                public stylable_element,
+                public nested_trait<element> {
 public:
     package(const std::vector<std::string> &using_namespaces);
 
@@ -44,8 +44,6 @@ public:
 
     std::string full_name(bool relative) const override;
 
-    friend bool operator==(const package &l, const package &r);
-
     bool is_deprecated() const;
 
     void set_deprecated(bool deprecated);
diff --git a/src/package_diagram/generators/plantuml/package_diagram_generator.cc b/src/package_diagram/generators/plantuml/package_diagram_generator.cc
index befc0bc4..d4ff8b11 100644
--- a/src/package_diagram/generators/plantuml/package_diagram_generator.cc
+++ b/src/package_diagram/generators/plantuml/package_diagram_generator.cc
@@ -54,8 +54,9 @@ void generator::generate_relationships(
     }
 
     // Process it's subpackages relationships
-    for (const std::unique_ptr<package> &subpackage : p) {
-        generate_relationships(*subpackage, ostr);
+    for (const auto &subpackage : p) {
+        generate_relationships(
+            dynamic_cast<const package &>(*subpackage), ostr);
     }
 }
 
@@ -77,7 +78,7 @@ void generator::generate(const package &p, std::ostream &ostr) const
     ostr << " {" << '\n';
 
     for (const auto &subpackage : p) {
-        generate(*subpackage, ostr);
+        generate(dynamic_cast<const package &>(*subpackage), ostr);
     }
 
     ostr << "}" << '\n';
@@ -93,14 +94,14 @@ void generator::generate(std::ostream &ostr) const
 
     if (m_config.should_include_entities("packages")) {
         for (const auto &p : m_model) {
-            generate(*p, ostr);
+            generate(dynamic_cast<package &>(*p), ostr);
             ostr << '\n';
         }
     }
 
     // Process package relationships
     for (const auto &p : m_model) {
-        generate_relationships(*p, ostr);
+        generate_relationships(dynamic_cast<package &>(*p), ostr);
         ostr << '\n';
     }
 
diff --git a/src/package_diagram/generators/plantuml/package_diagram_generator.h b/src/package_diagram/generators/plantuml/package_diagram_generator.h
index 600417e2..075bcef0 100644
--- a/src/package_diagram/generators/plantuml/package_diagram_generator.h
+++ b/src/package_diagram/generators/plantuml/package_diagram_generator.h
@@ -46,9 +46,9 @@ template <typename C, typename D>
 using common_generator =
     clanguml::common::generators::plantuml::generator<C, D>;
 
+using clanguml::common::model::package;
 using clanguml::common::model::relationship_t;
 using clanguml::common::model::scope_t;
-using clanguml::common::model::package;
 using namespace clanguml::util;
 
 class generator : public common_generator<diagram_config, diagram_model> {
diff --git a/src/package_diagram/model/diagram.cc b/src/package_diagram/model/diagram.cc
index 1d7b1b92..04748b43 100644
--- a/src/package_diagram/model/diagram.cc
+++ b/src/package_diagram/model/diagram.cc
@@ -33,7 +33,7 @@ std::string diagram::to_alias(const std::string &full_name) const
         throw error::uml_alias_missing(
             fmt::format("Missing alias for '{}'", full_name));
 
-    auto package = get_element(fn);
+    auto package = get_element<common::model::package>(fn);
 
     if (!package)
         throw error::uml_alias_missing(
diff --git a/src/package_diagram/model/diagram.h b/src/package_diagram/model/diagram.h
index 3f3319eb..ba092552 100644
--- a/src/package_diagram/model/diagram.h
+++ b/src/package_diagram/model/diagram.h
@@ -29,7 +29,7 @@ namespace clanguml::package_diagram::model {
 
 class diagram : public clanguml::common::model::diagram,
                 public clanguml::common::model::nested_trait<
-                    clanguml::common::model::package> {
+                    clanguml::common::model::element> {
 public:
     diagram() = default;
 
diff --git a/src/package_diagram/visitor/translation_unit_visitor.cc b/src/package_diagram/visitor/translation_unit_visitor.cc
index 2e1b2bf9..f76d97ba 100644
--- a/src/package_diagram/visitor/translation_unit_visitor.cc
+++ b/src/package_diagram/visitor/translation_unit_visitor.cc
@@ -37,11 +37,11 @@ namespace clanguml::package_diagram::visitor {
 
 using clanguml::class_diagram::model::type_alias;
 using clanguml::common::model::access_t;
+using clanguml::common::model::package;
 using clanguml::common::model::relationship;
 using clanguml::common::model::relationship_t;
 using clanguml::common::model::scope_t;
 using clanguml::package_diagram::model::diagram;
-using clanguml::common::model::package;
 
 namespace detail {
 scope_t cpp_access_specifier_to_scope(
@@ -125,7 +125,8 @@ void translation_unit_visitor::operator()(const cppast::cpp_entity &file)
                                 ctx.diagram().add_element(
                                     package_parent, std::move(p));
                                 ctx.set_current_package(
-                                    ctx.diagram().get_element(package_path));
+                                    ctx.diagram().get_element<package>(
+                                        package_path));
                             }
                         }