Merge pull request #55 from ouhammmourachid/add-enum-in-flowchart

add LinkShape and LinkHead enum to the flowchart diagram
ouhammmourachid · Nov 11, 2023 · 0f7e9e3 · 0f7e9e3
2 parents a48bb9e + 7657ab9
commit 0f7e9e3
Show file tree

Hide file tree

Showing 3 changed files with 42 additions and 9 deletions.
diff --git a/mermaid/flowchart/__init__.py b/mermaid/flowchart/__init__.py
@@ -1,6 +1,6 @@
 from mermaid.graph import Graph
 
-from .link import Link
+from .link import Link, LinkHead, LinkShape
 from .node import Node
 
 

diff --git a/mermaid/flowchart/link.py b/mermaid/flowchart/link.py
@@ -1,3 +1,6 @@
+from enum import Enum
+from typing import Union
+
 from .node import Node
 
 LINK_SHAPES: dict[str, str] = {
@@ -16,20 +19,38 @@
 }
 
 
+class LinkShape(Enum):
+    NORMAL = 'normal'
+    DOTTED = 'dotted'
+    THICK = 'thick'
+    HIDEN = 'hiden'
+
+
+class LinkHead(Enum):
+    NONE = 'none'
+    ARROW = 'arrow'
+    LEFT_ARROW = 'left-arrow'
+    BULLET = 'bullet'
+    CROSS = 'cross'
+
+
 class Link:
     def __init__(self,
                  origin: Node,
                  end: Node,
-                 shape: str = 'normal',
-                 head_left: str = 'none',
-                 head_right: str = 'arrow',
+                 shape: Union[str, LinkShape] = 'normal',
+                 head_left: Union[str, LinkHead] = 'none',
+                 head_right: Union[str, LinkHead] = 'arrow',
                  message: str = '') -> None:
         self.oigin: Node = origin
         self.end: Node = end
-        self.head_left = LINK_HEADS[head_left]
-        self.head_right = LINK_HEADS[head_right]
-        self.shape = LINK_SHAPES[shape]
-        self.message = f'|{message}|' if message else message
+        self.head_left: str = LINK_HEADS[
+            head_left if isinstance(head_left, str) else head_left.value]
+        self.head_right: str = LINK_HEADS[
+            head_right if isinstance(head_right, str) else head_right.value]
+        self.shape: str = LINK_SHAPES[shape if isinstance(shape, str
+                                                          ) else shape.value]
+        self.message: str = f'|{message}|' if message else message
 
     def __str__(self) -> str:
         if self.message:

diff --git a/mermaid/tests/test_flowchart.py b/mermaid/tests/test_flowchart.py
@@ -1,6 +1,6 @@
 import unittest
 
-from mermaid.flowchart import FlowChart, Link, Node
+from mermaid.flowchart import FlowChart, Link, LinkHead, LinkShape, Node
 from mermaid.flowchart.node import NodeShape
 
 
@@ -81,6 +81,18 @@ def test_str_link_with_no_default_value(self):
                           message=message)
         expect_string: str = f'first_node o-.-x|{message}| second_node'
         self.assertEqual(expect_string, str(link))
+
+    def test_str_link_with_enum_shape(self):
+        link: Link = Link(self.node_1, self.node_2, shape=LinkShape.DOTTED)
+        expect_string: str = 'first_node -.-> second_node'
+        self.assertEqual(expect_string, str(link))
+
+    def test_str_link_with_enum_link_head(self):
+        link: Link = Link(self.node_1,
+                          self.node_2,
+                          head_left=LinkHead.BULLET,
+                          head_right=LinkHead.CROSS)
+        expect_string: str = 'first_node o--x second_node'
         self.assertEqual(expect_string, str(link))