diff --git a/src/common/model/element.cc b/src/common/model/element.cc
index b5c7fe25..ec1b2549 100644
--- a/src/common/model/element.cc
+++ b/src/common/model/element.cc
@@ -1,5 +1,5 @@
 /**
- * src/class_diagram/model/element.cc
+ * src/common/model/element.cc
  *
  * Copyright (c) 2021-2022 Bartek Kryza <bkryza@gmail.com>
  *
diff --git a/src/common/model/element.h b/src/common/model/element.h
index 1fad1f0f..6704a124 100644
--- a/src/common/model/element.h
+++ b/src/common/model/element.h
@@ -1,5 +1,5 @@
 /**
- * src/class_diagram/model/element.h
+ * src/common/model/element.h
  *
  * Copyright (c) 2021-2022 Bartek Kryza <bkryza@gmail.com>
  *
diff --git a/src/common/model/nested_trait.h b/src/common/model/nested_trait.h
new file mode 100644
index 00000000..37644d12
--- /dev/null
+++ b/src/common/model/nested_trait.h
@@ -0,0 +1,125 @@
+/**
+ * src/common/model/element.h
+ *
+ * Copyright (c) 2021-2022 Bartek Kryza <bkryza@gmail.com>
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+#pragma once
+
+#include "util/util.h"
+
+#include <type_safe/optional_ref.hpp>
+
+#include <string>
+#include <vector>
+
+namespace clanguml::common::model {
+template <typename T> class nested_trait {
+public:
+    nested_trait() = default;
+
+    nested_trait(const nested_trait &) = delete;
+    nested_trait(nested_trait &&) = default;
+
+    nested_trait &operator=(const nested_trait &) = delete;
+    nested_trait &operator=(nested_trait &&) = default;
+
+    virtual ~nested_trait() = default;
+
+    void add_element(std::unique_ptr<T> p)
+    {
+        auto it = std::find_if(elements_.begin(), elements_.end(),
+            [&p](const auto &e) { return *e == *p; });
+
+        if (it != elements_.end()) {
+            (*it)->append(*p);
+        }
+        else {
+            elements_.emplace_back(std::move(p));
+        }
+    }
+
+    void add_element(std::vector<std::string> path, std::unique_ptr<T> p)
+    {
+        assert(p);
+
+        LOG_DBG("Adding nested element {} at path '{}'", p->name(),
+            fmt::join(path, "::"));
+
+        if (path.empty()) {
+            add_element(std::move(p));
+            return;
+        }
+
+        auto parent = get_element(path);
+
+        if (parent)
+            parent.value().add_element(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
+    {
+        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 {};
+        }
+
+        auto p = get_element(path.at(0));
+        if (path.size() == 1)
+            return p;
+
+        return p.value().get_element(
+            std::vector<std::string>(path.begin() + 1, path.end()));
+    }
+
+    type_safe::optional_ref<T> 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 {};
+
+        assert(it->get() != nullptr);
+
+        return type_safe::ref(*(it->get()));
+    }
+
+    bool has_element(const std::string &name) const
+    {
+        return std::find_if(elements_.cbegin(), elements_.cend(),
+                   [&](const auto &p) { return name == p->name(); }) !=
+            elements_.end();
+    }
+
+    auto begin() { return elements_.begin(); }
+    auto end() { return elements_.end(); }
+
+    auto cbegin() const { return elements_.cbegin(); }
+    auto cend() const { return elements_.cend(); }
+
+    auto begin() const { return elements_.begin(); }
+    auto end() const { return elements_.end(); }
+
+private:
+    std::vector<std::unique_ptr<T>> elements_;
+};
+
+}
diff --git a/src/package_diagram/model/diagram.cc b/src/package_diagram/model/diagram.cc
index 00a70863..1d7b1b92 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_package(fn);
+    auto package = get_element(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 c4d9ef9a..2a1d8ab6 100644
--- a/src/package_diagram/model/diagram.h
+++ b/src/package_diagram/model/diagram.h
@@ -28,7 +28,7 @@
 namespace clanguml::package_diagram::model {
 
 class diagram : public clanguml::common::model::diagram,
-                public detail::package_trait<package> {
+                public clanguml::common::model::nested_trait<package> {
 public:
     diagram() = default;
 
diff --git a/src/package_diagram/model/package.h b/src/package_diagram/model/package.h
index 83e8646e..6f7ed77f 100644
--- a/src/package_diagram/model/package.h
+++ b/src/package_diagram/model/package.h
@@ -18,6 +18,7 @@
 #pragma once
 
 #include "common/model/element.h"
+#include "common/model/nested_trait.h"
 #include "common/model/stylable_element.h"
 #include "util/util.h"
 
@@ -30,108 +31,9 @@
 
 namespace clanguml::package_diagram::model {
 
-namespace detail {
-template <typename T> class package_trait {
-public:
-    package_trait() = default;
-
-    package_trait(const package_trait &) = delete;
-    package_trait(package_trait &&) = default;
-
-    package_trait &operator=(const package_trait &) = delete;
-    package_trait &operator=(package_trait &&) = default;
-
-    virtual ~package_trait() = default;
-
-    void add_package(std::unique_ptr<T> p)
-    {
-        auto it = std::find_if(packages_.begin(), packages_.end(),
-            [&p](const auto &e) { return *e == *p; });
-
-        if (it != packages_.end()) {
-            (*it)->append(*p);
-        }
-        else {
-            packages_.emplace_back(std::move(p));
-        }
-    }
-
-    void add_package(std::vector<std::string> path, std::unique_ptr<T> p)
-    {
-        assert(p);
-
-        LOG_DBG(
-            "Adding package {} at path '{}'", p->name(), fmt::join(path, "::"));
-
-        if (path.empty()) {
-            add_package(std::move(p));
-            return;
-        }
-
-        auto parent = get_package(path);
-
-        if (parent)
-            parent.value().add_package(std::move(p));
-        else
-            spdlog::error(
-                "No parent package found at: {}", fmt::join(path, "::"));
-    }
-
-    type_safe::optional_ref<T> get_package(std::vector<std::string> path) const
-    {
-        LOG_DBG("Getting package at path: {}", fmt::join(path, "::"));
-
-        if (path.empty() || !has_package(path.at(0))) {
-            LOG_WARN(
-                "Sub package {} not found in package", fmt::join(path, "::"));
-            return {};
-        }
-
-        auto p = get_package(path.at(0));
-        if (path.size() == 1)
-            return p;
-
-        return p.value().get_package(
-            std::vector<std::string>(path.begin() + 1, path.end()));
-    }
-
-    type_safe::optional_ref<T> get_package(const std::string &name) const
-    {
-        auto it = std::find_if(packages_.cbegin(), packages_.cend(),
-            [&](const auto &p) { return name == p->name(); });
-
-        if (it == packages_.end())
-            return {};
-
-        assert(it->get() != nullptr);
-
-        return type_safe::ref(*(it->get()));
-    }
-
-    bool has_package(const std::string &name) const
-    {
-        return std::find_if(packages_.cbegin(), packages_.cend(),
-                   [&](const auto &p) { return name == p->name(); }) !=
-            packages_.end();
-    }
-
-    auto begin() { return packages_.begin(); }
-    auto end() { return packages_.end(); }
-
-    auto cbegin() const { return packages_.cbegin(); }
-    auto cend() const { return packages_.cend(); }
-
-    auto begin() const { return packages_.begin(); }
-    auto end() const { return packages_.end(); }
-
-private:
-    std::vector<std::unique_ptr<T>> packages_;
-};
-}
-
 class package : public common::model::element,
                 public common::model::stylable_element,
-                public detail::package_trait<package> {
+                public common::model::nested_trait<package> {
 public:
     package(const std::vector<std::string> &using_namespaces);
 
diff --git a/src/package_diagram/visitor/translation_unit_visitor.cc b/src/package_diagram/visitor/translation_unit_visitor.cc
index b16cd376..7906df1b 100644
--- a/src/package_diagram/visitor/translation_unit_visitor.cc
+++ b/src/package_diagram/visitor/translation_unit_visitor.cc
@@ -122,10 +122,10 @@ void translation_unit_visitor::operator()(const cppast::cpp_entity &file)
                             }
 
                             if (!p->skip()) {
-                                ctx.diagram().add_package(
+                                ctx.diagram().add_element(
                                     package_parent, std::move(p));
                                 ctx.set_current_package(
-                                    ctx.diagram().get_package(package_path));
+                                    ctx.diagram().get_element(package_path));
                             }
                         }