From d9a096facfde93a78f1ce7546bb4f34f4e3cbde1 Mon Sep 17 00:00:00 2001
From: Patrick Lehmann <Patrick.Lehmann@plc2.de>
Date: Sat, 19 Jun 2021 02:22:36 +0200
Subject: Improvements to pyGHDL.dom.

---
 pyGHDL/dom/DesignUnit.py | 10 ++++++----
 1 file changed, 6 insertions(+), 4 deletions(-)

(limited to 'pyGHDL/dom/DesignUnit.py')

diff --git a/pyGHDL/dom/DesignUnit.py b/pyGHDL/dom/DesignUnit.py
index 3619c47c5..534149677 100644
--- a/pyGHDL/dom/DesignUnit.py
+++ b/pyGHDL/dom/DesignUnit.py
@@ -41,7 +41,8 @@ This module contains all DOM classes for VHDL's design units (:class:`context <E
 """
 from pydecor import export
 
-from pyVHDLModel.VHDLModel import Entity as VHDLModel_Entity
+from pyGHDL.dom.Symbol import EntitySymbol
+from pyVHDLModel.VHDLModel import Entity as VHDLModel_Entity, EntityOrSymbol
 from pyVHDLModel.VHDLModel import Architecture as VHDLModel_Architecture
 from pyVHDLModel.VHDLModel import Package as VHDLModel_Package
 from pyVHDLModel.VHDLModel import PackageBody as VHDLModel_PackageBody
@@ -87,17 +88,18 @@ class Entity(VHDLModel_Entity, GHDLMixin):
 
 @export
 class Architecture(VHDLModel_Architecture, GHDLMixin):
-    def __init__(self, name: str, entityName: str):
+    def __init__(self, name: str, entity: EntityOrSymbol):
         super().__init__(name)
 
-        self.__entityName = entityName
+        self._entity = entity
 
     @classmethod
     def parse(cls, libraryUnit):
         name = NodeToName(libraryUnit)
         entityName = NodeToName(nodes.Get_Entity_Name(libraryUnit))
+        entity = EntitySymbol(entityName)
 
-        architecture = cls(name, entityName)
+        architecture = cls(name, entity)
 
         for item in GetDeclaredItemsFromChainedNodes(
             nodes.Get_Declaration_Chain(libraryUnit), "architecture", name
-- 
cgit v1.2.3