From 45ffea92543d1358b87cd759b35b766cc5086d0e Mon Sep 17 00:00:00 2001
From: Georgi Ivanov <givanov@hotels.com>
Date: Sat, 4 Apr 2020 01:15:28 +0100
Subject: [PATCH] ability to add icons to the cluster label

---
 diagrams/__init__.py   | 33 +++++++++++++++++++++++++++++++--
 docs/guides/cluster.md | 25 ++++++++++++++++++++++++-
 2 files changed, 55 insertions(+), 3 deletions(-)

diff --git a/diagrams/__init__.py b/diagrams/__init__.py
index 9626ea6d..f2f034fb 100644
--- a/diagrams/__init__.py
+++ b/diagrams/__init__.py
@@ -37,6 +37,10 @@ def getcluster():
 def setcluster(cluster):
     __cluster.set(cluster)
 
+def new_init(cls, init):
+    def reset_init(*args, **kwargs):
+        cls.__init__ = init
+    return reset_init
 
 class Diagram:
     __directions = ("TB", "BT", "LR", "RL")
@@ -194,7 +198,13 @@ class Cluster:
     # FIXME:
     #  Cluster direction does not work now. Graphviz couldn't render
     #  correctly for a subgraph that has a different rank direction.
-    def __init__(self, label: str = "cluster", direction: str = "LR"):
+    def __init__(
+        self,
+        label: str = "cluster", 
+        direction: str = "LR",
+        icon: object = None,
+        icon_size: int = 30
+        ):
         """Cluster represents a cluster context.
 
         :param label: Cluster label.
@@ -202,14 +212,25 @@ class Cluster:
         """
         self.label = label
         self.name = "cluster_" + self.label
+        self.icon = icon
+        self.icon_size = icon_size
 
         self.dot = Digraph(self.name)
 
         # 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 an icon is set, try to find and instantiate a Node without calling __init__()
+        # then find it's icon by calling _load_icon()
+        if self.icon:
+            _node = self.icon(_no_init=True)
+            if isinstance(_node,Node):
+                self.icon_label = '<<TABLE border="0"><TR><TD fixedsize="true" width="' + str(self.icon_size) +'" height="' + str(self.icon_size) +'"><IMG SRC="' + _node._load_icon() + '"></IMG></TD><TD>' + self.label + '</TD></TR></TABLE>>'
+                self.dot.graph_attr["label"] = self.icon_label
+        else:
+            self.dot.graph_attr["label"] = self.label
+ 
         if not self._validate_direction(direction):
             raise ValueError(f'"{direction}" is not a valid direction')
         self.dot.graph_attr["rankdir"] = direction
@@ -262,6 +283,14 @@ class Node:
 
     _height = 1.9
 
+    def __new__(cls, *args, **kwargs):
+        instance = object.__new__(cls)
+        lazy = kwargs.pop('_no_init', False)
+        if not lazy:
+            return instance
+        cls.__init__ = new_init(cls, cls.__init__)
+        return instance
+
     def __init__(self, label: str = ""):
         """Node represents a system component.
 
diff --git a/docs/guides/cluster.md b/docs/guides/cluster.md
index 6cbf820e..5001597a 100644
--- a/docs/guides/cluster.md
+++ b/docs/guides/cluster.md
@@ -66,6 +66,29 @@ with Diagram("Event Processing", show=False):
     handlers >> dw
 ```
 
+## Clusters with icons in the label
+
+You can add a Node icon before the cluster label (and specify its size as well).  You need to import the used Node class first.
+
+```python
+from diagrams import Cluster, Diagram
+from diagrams.aws.compute import ECS
+from diagrams.aws.database import RDS, Aurora
+from diagrams.aws.network import Route53, VPC
+
+with Diagram("Simple Web Service with DB Cluster", show=False):
+    dns = Route53("dns")
+    web = ECS("service")
+
+    with Cluster(label='VPC',icon=VPC):
+        with Cluster("DB Cluster",icon=Aurora,icon_size=30):
+            db_master = RDS("master")
+            db_master - [RDS("slave1"),
+                         RDS("slave2")]
+
+        dns >> web >> db_master
+```
+
 ![event processing diagram](/img/event_processing_diagram.png)
 
-> There is no depth limit of nesting. Feel free to create nested clusters as deep as you want.
\ No newline at end of file
+> There is no depth limit of nesting. Feel free to create nested clusters as deep as you want.