From 922fc14fae87a31193b9e9913c977627a233b758 Mon Sep 17 00:00:00 2001
From: Chuan Ou Yang <Chuan.Ou.Yang@ibm.com>
Date: Wed, 25 May 2022 13:16:56 +0800
Subject: [PATCH] chore: update full support graph renderer

---
 compose_viz/graph.py | 38 +++++++++++++++++++++++++++-----------
 1 file changed, 27 insertions(+), 11 deletions(-)

diff --git a/compose_viz/graph.py b/compose_viz/graph.py
index 1e71be4..992716f 100644
--- a/compose_viz/graph.py
+++ b/compose_viz/graph.py
@@ -32,10 +32,16 @@ def apply_edge_style(type) -> dict:
         },
         "volumes": {
             "style": "dashed",
+            "dir": "both",
         },
         "depends_on": {
             "style": "dotted",
         },
+        "extends":{
+            "dir": "both",
+            "arrowhead": "inv",
+            "arrowtail": "dot",
+        }
     }
 
     return style[type]
@@ -48,25 +54,35 @@ class Graph:
         self.compose = compose
         self.filename = filename
 
-    def add_vertex(self, name: str, type: str) -> None:
-        self.dot.node(name, **apply_vertex_style(type))
+    def validate_name(self, name: str) -> str:
+        # graphviz does not allow ':' in node name
+        transTable = name.maketrans({":": ""})
+        return name.translate(transTable)
 
-    def add_edge(self, head: str, tail: str, type: str) -> None:
-        self.dot.edge(head, tail, **apply_edge_style(type))
+    def add_vertex(self, name: str, type: str, lable: str = None) -> None:
+        self.dot.node(self.validate_name(name), lable, **apply_vertex_style(type))
+
+    def add_edge(self, head: str, tail: str, type: str, lable: str = None) -> None:
+        self.dot.edge(self.validate_name(head), self.validate_name(tail), lable, **apply_edge_style(type))
 
     def render(self, format: str, cleanup: bool = True) -> None:
         for service in self.compose.services:
-            self.add_vertex(service.name, "service")
+            if service.image is not None:
+                self.add_vertex(service.name, "service", lable=f"{service.name}\n({service.image})")
+            if service.extends is not None:
+                self.add_edge(service.extends.service_name, service.name, "extends")
             for network in service.networks:
-                self.add_vertex("net#" + network, "network")
-                self.add_edge(service.name, "net#" + network, "links")
+                self.add_vertex(network, "network", lable=f"net:{network}")
+                self.add_edge(service.name, network, "links")
             for volume in service.volumes:
                 self.add_vertex(volume.source, "volume")
-                self.add_edge(service.name, volume.source, "links")
+                self.add_edge(service.name, volume.source, "volumes", lable=volume.target)
             for port in service.ports:
-                self.add_vertex(port.host_port, "port")
-                self.add_edge(service.name, port.host_port, "ports")
+                self.add_vertex(port.host_port, "port", lable=port.host_port)
+                self.add_edge(port.host_port, service.name, "ports", lable=port.container_port)
+            for link in service.links:
+                self.add_edge(link.split(":")[0], service.name, "links", link.split(":")[1])
             for depends_on in service.depends_on:
-                self.dot.edge(depends_on, service.name, "depends_on")
+                self.add_edge(service.name, depends_on, "depends_on")
 
         self.dot.render(outfile=self.filename, format=format, cleanup=cleanup)