All objects are now accessible through Diagram

main.py

@@ -20,7 +20,7 @@ class MainScene(MovingCameraScene):
         parser = PUMLParser()
         d: Diagram = parser.parse_file(self.file)
 
-        layout = SpringLayout(d)
+        layout = HierarchicalLayout(d)
         layout.apply()
         layout.scale(self.scale_x, self.scale_y)
 

pyplaml/diagram.py

@@ -11,55 +11,53 @@ class Diagram:
         self.objects: typing.Dict[str, DiagramObject] = {}
         self.animate = False
 
-    def draw(self):
+    def __setitem__(self, key: str, obj: DiagramObject):
+        self.objects[key] = obj
+
+    def __getitem__(self, key: str):
+        return self.objects[key]
 
+    def objects_by_type(self, obj_type: str = ''):
+        res = {}
         for name, obj in self.objects.items():
+            _type = type(obj).__name__
+
+            if _type in res:
+                res[_type][name] = obj
+            else:
+                res[_type] = {}
+                res[_type][name] = obj
+
+        return res[obj_type] if obj_type != '' else res
+
+    def draw(self):
+        for name, obj in self.objects_by_type('DiagramClass').items():
             self.draw_object(obj)
 
         for name, obj in self.objects.items():
-            for i, edge in enumerate(obj.edges):
-
-                if self.animate:
-                    self.scene.play(Create(edge.draw()))
-                else:
-                    self.scene.add(edge.draw())
+            if hasattr(obj, 'edges'):
+                for i, edge in enumerate(obj.edges):
+                    if self.animate:
+                        self.scene.play(Create(edge.draw()))
+                    else:
+                        self.scene.add(edge.draw())
 
     def draw_object(self, obj: DiagramObject):
-
         if obj.mobject is None:
             mobject = obj.draw()
+            if mobject:
+                mobject.to_edge(UP)
+                mobject.shift(RIGHT * obj.x)
+                mobject.shift(DOWN * obj.y)
 
-            mobject.to_edge(UP)
-
-            mobject.shift(RIGHT * obj.x)
-            mobject.shift(DOWN * obj.y)
-
-            if self.animate:
-                self.scene.play(Create(mobject))
-                self.scene.play(self.scene.camera.auto_zoom(
-                    self.scene.mobjects, margin=0.5))
-            else:
-                self.scene.add(mobject)
-                self.scene.camera.auto_zoom(
-                    self.scene.mobjects, margin=0.5, animate=False)
-
-    def add_object(self, obj: DiagramObject):
-
-        if obj.name not in self.objects:
-            self.objects[obj.name] = obj
-
-            for edge in obj.edges:
-                if edge.target.name in self.objects:
-                    edge.target = self.objects[edge.target.name]
+                if self.animate:
+                    self.scene.play(Create(mobject))
+                    self.scene.play(self.scene.camera.auto_zoom(
+                        self.scene.mobjects, margin=0.5))
                 else:
-                    self.add_object(edge.target)
-        else:
-            for edge in obj.edges:
-                if edge.target not in self.objects:
-                    self.add_object(edge.target)
-                edge.target = self.objects[edge.target.name]
-
-            self.objects[obj.name].edges += obj.edges
+                    self.scene.add(mobject)
+                    self.scene.camera.auto_zoom(
+                        self.scene.mobjects, margin=0.5, animate=False)
 
     def set_scene(self, scene: Scene):
         self.scene = scene

pyplaml/diagram_class.py

 from .diagram_object import DiagramObject
 from .diagram_edge import DiagramEdge
 from .class_attribute import ClassAttribute
+from .diagram import Diagram
 
 from manim import *
 
@@ -14,8 +15,24 @@ class DiagramClass(DiagramObject):
         self.attributes: List[ClassAttribute] = []
         self.methods: List[ClassAttribute] = []
 
-    def draw(self):
+    def append_to_diagram(self, diagram: Diagram) -> DiagramObject:
+        if self.name not in diagram.objects:
+            diagram[self.name] = self
+            for edge in self.edges:
+                if edge.target.name in diagram.objects:
+                    edge.target = diagram[edge.target.name]
+                else:
+                    edge.target.append_to_diagram(diagram)
+        else:
+            for edge in self.edges:
+                if edge.target not in diagram:
+                    edge.target.append_to_diagram(diagram)
+                edge.target = diagram[edge.target.name]
+            diagram[self.name].edges += self.edges
+
+        return diagram[self.name]
 
+    def draw(self):
         header = Rectangle(color=GRAY)
         text = Text(self.name, color=BLACK)
         header.surround(text)

pyplaml/diagram_edge.py

 from .diagram_object import DiagramObject
 from .relation import Relation
+import warnings
 
 from manim import *
 
@@ -25,11 +26,12 @@ class DiagramEdge(DiagramObject):
         self.target_text = ""
 
     def draw(self):
-
         if self.source.mobject is None:
-            raise Exception("Start object \"{}\" has not been drawn.".format(self.source.name))
+            warnings.warn("Start object \"{}\" has not been drawn.".format(self.source.name))
+            return
         elif self.target.mobject is None:
-            raise Exception("Target object \"{}\" has not been drawn.".format(self.target.name))
+            warnings.warn("Target object \"{}\" has not been drawn.".format(self.target.name))
+            return
 
         start = self.source.mobject.get_top()
         target = self.target.mobject.get_bottom()
@@ -84,22 +86,17 @@ class DiagramEdge(DiagramObject):
     @staticmethod
     def get_line_tip(rel: Relation):
         if rel == Relation.EXTENSION:
-
             return ArrowTriangleTip(color=BLACK, stroke_width=2, length=0.2, width=0.2)
 
         elif rel == Relation.ASSOCIATION:
-
             return StealthTip(color=BLACK, stroke_width=2, length=0.2)
 
         elif rel == Relation.AGGREGATION:
-
             return ArrowSquareTip(color=BLACK, stroke_width=2,
                                   length=0.15)
 
         elif rel == Relation.COMPOSITION:
-
             return ArrowSquareFilledTip(color=BLACK, stroke_width=2, length=0.15)
 
         elif rel == Relation.HASH:
-
             return ArrowSquareTip(color=BLACK, stroke_width=2, length=0.15)

pyplaml/diagram_object.py

@@ -4,6 +4,8 @@ from abc import ABC, abstractmethod
 
 from manim import *
 
+import pyplaml
+
 
 class DiagramObject(ABC):
 
@@ -17,3 +19,8 @@ class DiagramObject(ABC):
     @abstractmethod
     def draw(self) -> VMobject:
         pass
+
+    def append_to_diagram(self, diagram: pyplaml.Diagram) -> DiagramObject:
+        if self.name not in diagram.objects:
+            diagram[self.name] = self
+        return diagram[self.name]

pyplaml/layout/diagram_layout.py

@@ -15,13 +15,11 @@ class DiagramLayout(ABC):
 
     def get_graph(self) -> nx.DiGraph:
         g = nx.DiGraph()
-
         for name, obj in self.diagram.objects.items():
             g.add_node(name)
-
-            for e in obj.edges:
-                g.add_edge(name, e.target.name)
-
+            if hasattr(obj, 'edges'):
+                for e in obj.edges:
+                    g.add_edge(name, e.target.name)
         return g
 
     def scale(self, x: float, y: float) -> None:

pyplaml/layout/hierarchical_layout.py

@@ -4,28 +4,27 @@ from .diagram_layout import DiagramLayout
 class HierarchicalLayout(DiagramLayout):
 
     def apply(self) -> None:
-        first_level_nodes = list(self.diagram.objects.keys())
-        for name, obj in self.diagram.objects.items():
+        first_level_nodes = list(self.diagram.objects_by_type('DiagramClass').keys())
+        for name, obj in self.diagram.objects_by_type('DiagramClass').items():
             if len(obj.edges) != 0:
                 first_level_nodes.remove(obj.name)
-
         self.assign_object_coordinates(first_level_nodes, 0)
 
     def assign_object_coordinates(self, nodes: list, y: int):
-
         if len(nodes) == 0:
             return
 
         next_nodes = []
         x_range = self.range_around_zero(len(nodes))
         for i, name in enumerate(nodes):
-            self.diagram.objects[name].x = x_range[i]
-            self.diagram.objects[name].y = y
-
-            for _, obj in self.diagram.objects.items():
-                for edge in obj.edges:
-                    if edge.target.name == name:
-                        next_nodes.append(obj.name)
+            self.diagram[name].x = x_range[i]
+            self.diagram[name].y = y
+
+            for _, obj in self.diagram.objects_by_type('DiagramClass').items():
+                if hasattr(obj, 'edges'):
+                    for edge in obj.edges:
+                        if edge.target.name == name:
+                            next_nodes.append(obj.name)
 
         self.assign_object_coordinates(next_nodes, y + 1)
 

pyplaml/puml_parser.py

@@ -43,10 +43,8 @@ class PUMLParser(object):
         (edge, is_source_on_left) = p[2]
         edge: DiagramEdge
 
-        l_class = DiagramClass(left_class_name, 'class')
-        r_class = DiagramClass(right_class_name, 'class')
-
-        edge.name = left_class_name + "-" + edge.source_rel_type.name + "-" + edge.target_rel_type.name + "-" + right_class_name
+        l_class = DiagramClass(left_class_name, 'class').append_to_diagram(self.diagram)
+        r_class = DiagramClass(right_class_name, 'class').append_to_diagram(self.diagram)
 
         if is_source_on_left:
             edge.source = l_class
@@ -57,11 +55,12 @@ class PUMLParser(object):
             edge.target = l_class
             r_class.add_edge(edge)
 
+        edge.name = edge.source.name + "-" + edge.source_rel_type.name + "-" + edge.target_rel_type.name + "-" + edge.target.name
+
         if len(p) == 5:
             edge.text = p[4]
 
-        self.diagram.add_object(l_class)
-        self.diagram.add_object(r_class)
+        edge.append_to_diagram(self.diagram)
 
     @staticmethod
     def p_rel_line(p):
@@ -156,9 +155,7 @@ class PUMLParser(object):
             class_type = "abstract_class"
             name = str(p[3])
 
-        class_obj = DiagramClass(name, class_type)
-
-        self.diagram.add_object(class_obj)
+        DiagramClass(name, class_type).append_to_diagram(self.diagram)
 
     def p_class_attr(self, p):
         """