Merge 3dd51ba713 into 1c380b01b5

diagrams/__init__.py

@@ -36,6 +36,10 @@ def getcluster() -> "Cluster":
 def setcluster(cluster: "Cluster"):
     __cluster.set(cluster)
 
+def getIconLabel(iconPath: str = '', iconSize: int = 30, label: str = ''):
+    return '<<table border="0" width="100%"><tr><td fixedsize="true" width="' + str(
+        iconSize) + '" height="' + str(
+        iconSize) + '"><img src="' + iconPath + '" /></td><td>' + label + '</td></tr></table>>'
 
 class Diagram:
     __directions = ("TB", "BT", "LR", "RL")
@@ -197,7 +201,6 @@ class Diagram:
         else:
             self.dot.render(format=self.outformat, view=self.show, quiet=True)
 
-
 class Cluster:
     __directions = ("TB", "BT", "LR", "RL")
     __bgcolors = ("#E5F5FD", "#EBF3E7", "#ECE8F6", "#FDF7E3")
@@ -228,18 +231,27 @@ class Cluster:
         :param label: Cluster label.
         :param direction: Data flow direction. Default is 'left to right'.
         :param graph_attr: Provide graph_attr dot config attributes.
+        :param providerIconNode: Provide a node to be included as a link
+        :param icon_size: The icon size
         """
         if graph_attr is None:
             graph_attr = {}
         self.label = label
         self.name = "cluster_" + self.label
+        self.providerIconNode = providerIconNode
+        self.icon_size = icon_size
 
         self.dot = Digraph(self.name)
 
-        # Set attributes.
+        # Set attributes
         for k, v in self._default_graph_attrs.items():
             self.dot.graph_attr[k] = v
-        self.dot.graph_attr["label"] = self.label
+
+        if self.providerIconNode:
+                _iconNode = self.providerIconNode() #load the object
+                self.dot.graph_attr["label"] = getIconLabel(_iconNode._load_icon(), self.icon_size, self.label)
+        else:
+            self.dot.graph_attr["label"] = self.label
 
         if not self._validate_direction(direction):
             raise ValueError(f'"{direction}" is not a valid direction')
@@ -273,6 +285,7 @@ class Cluster:
     def _validate_direction(self, direction: str) -> bool:
         return direction.upper() in self.__directions
 
+
     def node(self, nodeid: str, label: str, **attrs) -> None:
         """Create a new node in the cluster."""
         self.dot.node(nodeid, label=label, **attrs)
@@ -292,6 +305,9 @@ class Node:
 
     _height = 1.9
 
+    def __new__(cls, *args, **kwargs):
+        return object.__new__(cls)
+
     def __init__(self, label: str = "", *, nodeid: str = None, **attrs: Dict):
         """Node represents a system component.
 

tests/test_diagram.py

@@ -5,7 +5,7 @@ import pathlib
 
 from diagrams import Cluster, Diagram, Edge, Node
 from diagrams import getcluster, getdiagram, setcluster, setdiagram
-
+from diagrams.aws.network import VPC
 
 class DiagramTest(unittest.TestCase):
     def setUp(self):
@@ -23,6 +23,13 @@ class DiagramTest(unittest.TestCase):
                 os.remove(self.name + ".png")
             except FileNotFoundError:
                 pass
+    
+    def test_cluster_icon(self):
+        self.name = "example-cluster-icon"
+        with Diagram(name=self.name, show=False):
+            cluster = Cluster("example_cluster_icon", providerIconNode=VPC, icon_size=30)
+        self.assertIsNotNone(cluster)
+        self.assertTrue(os.path.exists(f"{self.name}.png"))
 
     def test_validate_direction(self):
         # Normal directions.