refactor

cuongnb14 · Oct 10, 2024 · 82f53e3 · 82f53e3
1 parent fd4da95
commit 82f53e3
Show file tree

Hide file tree

Showing 2 changed files with 63 additions and 52 deletions.
diff --git a/examples/test_flowchart.py b/examples/test_flowchart.py
@@ -10,19 +10,21 @@ def test_flowchart():
     chart = Flowchart(title="Simple Flowchart")
 
     # Declare nodes
+    subgraph = Subgraph(name="Storage")
+
     user_node = Node(name="User", shape=NodeShape.ROUNDED_RECTANGLE, icon=Icon.USER)
     elb_node = Node(name="ELB", shape=NodeShape.ROUNDED_RECTANGLE, icon=Icon.LOAD_BALANCER, class_name="gateway")
+
     order_node = Node(name="Order Service", shape=NodeShape.ROUNDED_RECTANGLE, class_name="api")
     payment_node = Node(name="Payment Service", shape=NodeShape.ROUNDED_RECTANGLE, class_name="api")
-    redis_node = Node(name="Redis", shape=NodeShape.ROUNDED_RECTANGLE, class_name="storage")
-    db_node = Node(name="DB", shape=NodeShape.CYLINDER, icon=Icon.DATABASE, class_name="storage")
+
+    redis_node = Node(name="Redis", shape=NodeShape.ROUNDED_RECTANGLE, class_name="storage", subgraph=subgraph)
+    db_node = Node(name="DB", shape=NodeShape.CYLINDER, icon=Icon.DATABASE, class_name="storage", subgraph=subgraph)
 
     # Declare node connections (edges)
     user_node.add_connections(Link(text="call"), elb_node)
     elb_node.add_connections(Link(), order_node, payment_node)
     order_node.add_connections(Link(), redis_node, db_node)
-    subgraph = Subgraph(name="Storage")
-    subgraph.add_nodes(redis_node, db_node)
     payment_node.add_connections(Link(), subgraph)
 
     # Add items to chart (nodes or subgraph)

diff --git a/mermaid_builder/flowchart.py b/mermaid_builder/flowchart.py
@@ -64,14 +64,69 @@ class LinkShape:
     THICK = "==>"
 
 
-@dataclass
+class Subgraph:
+    def __init__(self, name, direction=Direction.TB):
+        self.name = name
+        self.direction = direction
+        self.nodes = []
+        self.connections = []
+
+    def add_connections(self, link, *nodes):
+        for node in nodes:
+            self.add_connection(next_node=node, link=link)
+
+    def add_connection(self, next_node, link):
+        if not self.connections:
+            self.connections = []
+
+        self.connections.append(Connection(link=link, next_node=next_node))
+
+    def get_id(self) -> str:
+        return self.name.lower().replace(" ", "_")
+
+    def add_nodes(self, *nodes):
+        self.nodes.extend(nodes)
+
+    def draw_node(self):
+        return self.get_id()
+
+    def get_lines(self, indent=0):
+        result = [
+            f"{INTENT_CHAR * indent}style {self.get_id()} rx:10,ry:10",
+            f"{INTENT_CHAR * indent}subgraph {self.get_id()} [{self.name}]",
+            f"{INTENT_CHAR * (indent + 1)}direction {self.direction}",
+        ]
+        for node in self.nodes:
+            if isinstance(node, Subgraph):
+                result.extend(node.get_lines(indent + 1))
+            else:
+                result.append(f"{INTENT_CHAR * (indent + 1)}{node.draw_node()}")
+        result.append(f"{INTENT_CHAR * indent}end")
+
+        for connection in self.connections:
+            result.append(
+                f"{indent * INTENT_CHAR}{self.draw_node()}{connection.link.draw()}{connection.next_node.draw_node()}"
+            )
+        return result
+
+
 class Node:
     name: str
-    shape: NodeShape
+    shape: NodeShape = NodeShape.ROUNDED_RECTANGLE
     connections: list = ()
     class_name: str = ""
     icon: str = ""
 
+    def __init__(self, name: str, shape: NodeShape = NodeShape.ROUNDED_RECTANGLE, class_name: str = "", icon: str = "",
+                 connections: list = (), subgraph: Subgraph = None):
+        self.name = name
+        self.shape = shape
+        self.class_name = class_name
+        self.icon = icon
+        self.connections = connections
+        if subgraph:
+            subgraph.add_nodes(self)
+
     def get_id(self) -> str:
         return self.name.lower().replace(" ", "_")
 
@@ -136,52 +191,6 @@ class Connection:
     next_node: Node
 
 
-class Subgraph:
-    def __init__(self, name, direction=Direction.TB):
-        self.name = name
-        self.direction = direction
-        self.nodes = []
-        self.connections = []
-
-    def add_connections(self, link, *nodes):
-        for node in nodes:
-            self.add_connection(next_node=node, link=link)
-
-    def add_connection(self, next_node, link):
-        if not self.connections:
-            self.connections = []
-
-        self.connections.append(Connection(link=link, next_node=next_node))
-
-    def get_id(self) -> str:
-        return self.name.lower().replace(" ", "_")
-
-    def add_nodes(self, *nodes):
-        self.nodes.extend(nodes)
-
-    def draw_node(self):
-        return self.get_id()
-
-    def get_lines(self, indent=0):
-        result = [
-            f"{INTENT_CHAR * indent}style {self.get_id()} rx:10,ry:10",
-            f"{INTENT_CHAR * indent}subgraph {self.get_id()} [{self.name}]",
-            f"{INTENT_CHAR * (indent + 1)}direction {self.direction}",
-        ]
-        for node in self.nodes:
-            if isinstance(node, Subgraph):
-                result.extend(node.get_lines(indent + 1))
-            else:
-                result.append(f"{INTENT_CHAR * (indent + 1)}{node.draw_node()}")
-        result.append(f"{INTENT_CHAR * indent}end")
-
-        for connection in self.connections:
-            result.append(
-                f"{indent * INTENT_CHAR}{self.draw_node()}{connection.link.draw()}{connection.next_node.draw_node()}"
-            )
-        return result
-
-
 @dataclass
 class Style:
     class_name: str