Added layouts

main.py

@@ -4,15 +4,20 @@ from typing import Type
 from pyPlantUML import *
 
 import argparse
+from pyPlantUML.layout.KamadaKawai import KamadaKawai
+from pyPlantUML.layout.DotLayout import DotLayout
+from pyPlantUML.layout.SpringLayout import SpringLayout
 
 
 class MainScene(MovingCameraScene):
-
     def construct(self):
-
         parser = PUMLParser()
         diagram: Type[Diagram] = parser.parseFile(self.file)
 
+        layout = SpringLayout(diagram)
+        layout.apply()
+        layout.scale(7, 5)
+
         self.camera.background_color = WHITE
         Text.set_default(font_size=16)
 
@@ -34,13 +39,12 @@ if __name__ == "__main__":
     config.cairo_path = "media/images/"
     config.ffmpeg_path = "media/videos/"
 
-    argparser = argparse.ArgumentParser(
-        prog="pyPlantUML"
-    )
+    argparser = argparse.ArgumentParser(prog="pyPlantUML")
 
-    argparser.add_argument('-f', '--file', required=True,
-                           type=str, help="Path to source plantUml file.")
-    argparser.add_argument('-a', '--animate', action="store_true")
+    argparser.add_argument(
+        "-f", "--file", required=True, type=str, help="Path to source plantUml file."
+    )
+    argparser.add_argument("-a", "--animate", action="store_true")
 
     args = argparser.parse_args()
 

pyPlantUML/Diagram.py

@@ -2,7 +2,6 @@ from .DiagramObject import DiagramObject
 
 from manim import *
 
-
 class Diagram():
 
     def __init__(self, name: str):
@@ -10,36 +9,8 @@ class Diagram():
         self.objects: typing.Dict[str, DiagramObject] = {}
         self.animate = False
 
-    def setLayout(self):
-        firstLevelNodes = list(self.objects.keys())
-        for name, obj in self.objects.items():
-            if len(obj.edges) != 0:
-                firstLevelNodes.remove(obj.name)
-
-        self.assignObjectCoordinates(firstLevelNodes, 0)
-
-    def assignObjectCoordinates(self, nodes: list, y: int):
-
-        if len(nodes) == 0:
-            return
-
-        nextNodes = []
-        xRange = self.rangeAroundZero(len(nodes))
-        for i, name in enumerate(nodes):
-            self.objects[name].x = xRange[i]
-            self.objects[name].y = y
-
-            for _, obj in self.objects.items():
-                for edge in obj.edges:
-                    if edge.target.name == name:
-                        nextNodes.append(obj.name)
-
-        self.assignObjectCoordinates(nextNodes, y+1)
-
     def draw(self):
 
-        self.setLayout()
-
         for name, obj in self.objects.items():
             self.drawObject(obj)
 
@@ -58,8 +29,8 @@ class Diagram():
 
             mobj.to_edge(UP)
 
-            mobj.shift(RIGHT * obj.x * 2)
-            mobj.shift(DOWN * obj.y * 2)
+            mobj.shift(RIGHT * obj.x)
+            mobj.shift(DOWN * obj.y)
 
             if self.animate:
                 self.scene.play(Create(mobj))

pyPlantUML/layout/DiagramLayout.py

+from abc import ABC, abstractmethod
+import networkx as nx
+
+from .. import Diagram
+
+class DiagramLayout(ABC):
+
+    def __init__(self, diagram: Diagram) -> None:
+        self.diagram = diagram
+
+    @abstractmethod
+    def apply(self) -> None:
+        pass
+
+    def getGraph(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)
+        
+        return G
+
+    def scale(self, x: float, y: float) -> None:
+        for name, obj in self.diagram.objects.items():
+            self.diagram.objects[name].x = self.diagram.objects[name].x * x
+            self.diagram.objects[name].y = self.diagram.objects[name].y * y

pyPlantUML/layout/DotLayout.py

+from networkx.drawing.nx_agraph import graphviz_layout
+
+from .DiagramLayout import DiagramLayout
+from .. import Diagram
+
+class DotLayout(DiagramLayout):
+
+    def __init__(self, diagram: Diagram) -> None:
+        self.diagram = diagram
+
+    def apply(self) -> None:
+
+        layout = graphviz_layout(self.getGraph(), prog='dot')
+
+        for key, pos in layout.items():
+            self.diagram.objects[key].x = pos[0]
+            self.diagram.objects[key].y = pos[1]

pyPlantUML/layout/KamadaKawai.py

+import networkx as nx
+
+from .DiagramLayout import DiagramLayout
+from .. import Diagram
+
+class KamadaKawai(DiagramLayout):
+
+    def __init__(self, diagram: Diagram) -> None:
+        self.diagram = diagram
+
+    def apply(self) -> None:
+
+        layout = nx.kamada_kawai_layout(self.getGraph())
+
+        for key, pos in layout.items():
+            self.diagram.objects[key].x = pos[0]
+            self.diagram.objects[key].y = pos[1]

pyPlantUML/layout/SpringLayout.py

+import networkx as nx
+
+from .DiagramLayout import DiagramLayout
+from .. import Diagram
+
+class SpringLayout(DiagramLayout):
+
+    def __init__(self, diagram: Diagram) -> None:
+        self.diagram = diagram
+
+    def apply(self) -> None:
+
+        layout = nx.spring_layout(self.getGraph())
+
+        for key, pos in layout.items():
+            self.diagram.objects[key].x = pos[0]
+            self.diagram.objects[key].y = pos[1]