feat: added styling; support multiple files on read

tsypuk · Sep 27, 2023 · 5a5736c · 5a5736c
1 parent 92a122a
commit 5a5736c
Show file tree

Hide file tree

Showing 5 changed files with 76 additions and 31 deletions.
diff --git a/docs/docs/aws-components/output/jpg/certificate_manager.jpg b/docs/docs/aws-components/output/jpg/certificate_manager.jpg
diff --git a/multicloud_diagrams/__init__.py b/multicloud_diagrams/__init__.py
@@ -56,6 +56,17 @@ class Distribution:
     columns: int = 1
 
 
+def extract_actors(sequence_diagram):
+    actors = []
+    lines = sequence_diagram.split('\n')
+    for line in lines:
+        strip = line.strip()
+        if (strip.startswith('participant')) | (strip.startswith('actor')):
+            actor = line.strip().split()[1]
+            actors.append(actor)
+    return actors
+
+
 class MultiCloudDiagrams:
     def __init__(self, debug_mode=False, shadow=True, layer_name=''):
         self.actors_to_nodes = {}
@@ -313,7 +324,7 @@ def add_vertex_list(self, vertexes, distribution: Distribution = None):
             for vertex in vertexes:
                 self.add_vertex(**vertex)
 
-    def add_connection(self, src_node_id, dst_node_id, edge_style=None, labels=None, label_style=None, layer_name=None, layer_id=None):
+    def add_connection(self, src_node_id, dst_node_id, edge_style=None, labels=None, label_style=None, layer_name=None, layer_id=None, prefix=None):
         if labels is None:
             labels = []
         if edge_style is None:
@@ -332,15 +343,24 @@ def add_connection(self, src_node_id, dst_node_id, edge_style=None, labels=None,
             if found == 2:
                 break
 
+        edge_id = f'edge:{src_node_id}:to:{dst_node_id}'
+        label_id = f'label:{src_node_id}:to:{dst_node_id}'
+        parent_edge_id = f'edge:{src_node_id}:to:{dst_node_id}'
+
+        if prefix is not None:
+            edge_id = f'edge_{prefix}:{src_node_id}:to:{dst_node_id}'
+            label_id = f'label_{prefix}:{src_node_id}:to:{dst_node_id}'
+            parent_edge_id = f'edge_{prefix}:{src_node_id}:to:{dst_node_id}'
+
         if found == 2:
             # check that Edge does not exist
             edge_exist = False
             for mx_cell in self.root:
-                if mx_cell.attrib['id'] == f'edge:{src_node_id}:to:{dst_node_id}':
+                if mx_cell.attrib['id'] == edge_id:
                     edge_exist = True
                     # update the labels
                     for mxLabel in self.root:
-                        if mxLabel.attrib['id'] == f'label:{src_node_id}:to:{dst_node_id}':
+                        if mxLabel.attrib['id'] == f'label_{prefix}:{src_node_id}:to:{dst_node_id}':
                             if 'value' in mxLabel.attrib:
 
                                 for label in labels:
@@ -359,7 +379,7 @@ def add_connection(self, src_node_id, dst_node_id, edge_style=None, labels=None,
                 parent_id = str(self.get_layer_id(layer_name, layer_id))
                 mx_cell = Et.SubElement(self.root,
                                         'mxCell',
-                                        id=f'edge:{src_node_id}:to:{dst_node_id}',
+                                        id=edge_id,
                                         style=node_template['style'],
                                         parent=parent_id,
                                         source=f'vertex:{src_node_id}',
@@ -377,10 +397,10 @@ def add_connection(self, src_node_id, dst_node_id, edge_style=None, labels=None,
                     customize(node_template=node_template, style=label_style)
                     mx_cell = Et.SubElement(self.root,
                                             'mxCell',
-                                            id=f'label:{src_node_id}:to:{dst_node_id}',
+                                            id=label_id,
                                             value=stringify_labels(labels),
                                             style=node_template['style'],
-                                            parent=f'edge:{src_node_id}:to:{dst_node_id}',
+                                            parent=parent_edge_id,
                                             vertex="1",
                                             connectable="0")
                     # <mxGeometry relative="1" as="geometry">
@@ -412,6 +432,9 @@ def add_link(self, src_node_id, dst_node_id, action=None, layer_name=None, layer
         }
         self.add_connection(src_node_id=src_node_id, dst_node_id=dst_node_id, labels=action, edge_style=style, layer_name=layer_name, layer_id=layer_id)
 
+    def add_link_uml(self, src_node_id, dst_node_id, action=None, layer_name=None, layer_id=None, style=None):
+        self.add_connection(src_node_id=src_node_id, dst_node_id=dst_node_id, labels=action, edge_style=style, layer_name=layer_name, layer_id=layer_id, prefix=layer_name)
+
     def add_bidirectional_link(self, src_node_id, dst_node_id, action=None):
         style = {
             'startArrow': 'classic',
@@ -529,45 +552,39 @@ def export_to_file(self, file_path):
         with open(file_path, 'w', encoding="utf-8") as file:
             file.write(resulting_xml)
 
-    def read_uml_from_file(self, file_name):
+    def read_uml_from_file(self, file_name, style=None):
         with open(file_name, 'r') as file:
             sequence_diagram = file.read()
 
-        actors = self.extract_actors(sequence_diagram)
+        actors = extract_actors(sequence_diagram)
 
         print("Actors:")
         for actor in actors:
             print(actor)
 
         # create Layer with UML file name
-        self.add_layer(file_name)
-        self.extract_messages_from_uml(sequence_diagram, actors, layer_name=file_name)
-
-    def extract_actors(self, sequence_diagram):
-        actors = []
-        lines = sequence_diagram.split('\n')
-        for line in lines:
-            strip = line.strip()
-            if (strip.startswith('participant')) | (strip.startswith('actor')):
-                actor = line.strip().split()[1]
-                actors.append(actor)
-        return actors
+        base_name = os.path.splitext(os.path.basename(file_name))[0]
+        self.add_layer(base_name)
+        self.extract_messages_from_uml(sequence_diagram, actors, layer_name=base_name, style=style)
 
-    def extract_messages_from_uml(self, sequence_diagram, actors, layer_name):
+    def extract_messages_from_uml(self, sequence_diagram, actors, layer_name, style):
         lines = sequence_diagram.split('\n')
+        action_id = 0
         for line in lines:
             strip = line.strip()
             if any(strip.startswith(actor) for actor in actors):
                 data = self.extract_info(line)
                 print(data)
                 try:
                     # connect vertex of actor1 actor2 using arrow and message
-                    self.add_link(
+                    action_id = action_id + 1
+                    self.add_link_uml(
                         self.actors_to_nodes[data[0]],
                         self.actors_to_nodes[data[1]],
-                        action=[data[2]],
-                        layer_name=layer_name)
-                except:
+                        action=[f'{action_id}: {data[2]}'],
+                        layer_name=layer_name,
+                        style=style)
+                except KeyError:
                     print('No such node')
 
     # [Actor][Arrow][Actor]:Message text

diff --git a/samples/samples/aws_mermaid_uml.py b/samples/samples/aws_mermaid_uml.py
@@ -13,7 +13,32 @@ def main():
     prepare_end2end(mcd)
 
     mcd.read_uml_mappings('uml_mapping.yml')
-    mcd.read_uml_from_file('uml.mermaid')
+    style1 = {
+        'orthogonalLoop': '1',
+        'edgeStyle': 'orthogonalEdgeStyle',
+        'curved': '1',
+        'startArrow': 'oval',
+        'endArrow': 'classicThin',
+        'dashed': '1',
+        'strokeColor': '#FF3333',
+        'strokeWidth': '3',
+        'fontSize': '22'
+    }
+    mcd.read_uml_from_file('file_upload.mermaid', style=style1)
+
+    style2 = {
+        'orthogonalLoop': '1',
+        'edgeStyle': 'orthogonalEdgeStyle',
+        'curved': '1',
+        'startArrow': 'oval',
+        'endArrow': 'classicThin',
+        'dashed': '1',
+        'strokeColor': '#0000FF',
+        'strokeWidth': '3',
+        'fontSize': '22'
+    }
+
+    mcd.read_uml_from_file('process.mermaid', style=style2)
 
     mcd.export_to_file(result_file)
 

diff --git a/samples/samples/uml.mermaid → samples/samples/file_upload.mermaid b/samples/samples/uml.mermaid → samples/samples/file_upload.mermaid
@@ -2,14 +2,10 @@ sequenceDiagram
     actor Donald
     participant S3Bucket
     participant LambdaFunction1
-    participant LambdaFunction2
     participant SNS
     participant SQS
-    participant Storage
 
     Donald->>S3Bucket: Put new file into S3 bucket
     S3Bucket->>LambdaFunction1: LifeConfig Rule invokes lambda function
     LambdaFunction1->>SNS: Send Event to SNS that file processing is finished
-    SNS->>SQS: Fanout Event to SQS
-    SQS->>LambdaFunction2: poll message from SQS and invoke consumer lambda function
-    LambdaFunction2->>Storage: Update DynamoDB Table. Save image status.
+    SNS->>SQS: Fanout Event to SQS
diff --git a/samples/samples/process.mermaid b/samples/samples/process.mermaid
@@ -0,0 +1,7 @@
+sequenceDiagram
+    participant SQS
+    participant LambdaFunction2
+    participant Storage
+
+    SQS->>LambdaFunction2: poll message from SQS and invoke consumer lambda function
+    LambdaFunction2->>Storage: Update DynamoDB Table. Save image status.