From 678587a2bb15c78c0b5ce9734ce17e977cd42342 Mon Sep 17 00:00:00 2001
From: Bruno Meneguello <1322552+bkmeneguello@users.noreply.github.com>
Date: Tue, 29 Dec 2020 15:53:35 -0300
Subject: [PATCH] Allow Node icon to be modified

---
 diagrams/__init__.py | 21 +++++++++++++--------
 1 file changed, 13 insertions(+), 8 deletions(-)

diff --git a/diagrams/__init__.py b/diagrams/__init__.py
index df6a8f1..c307cb2 100644
--- a/diagrams/__init__.py
+++ b/diagrams/__init__.py
@@ -283,11 +283,15 @@ class Node(_Cluster):
     def __init__(
         self,
         label: str = "",
+        icon: object = None,
+        icon_size: int = None,
         **attrs: Dict
         ):
         """Node represents a system component.
 
         :param label: Node label.
+        :param icon: Custom icon for tihs cluster. Must be a node class or reference.
+        :param icon_size: The icon size when used as a Cluster. Default is 30.
         """
         # Generates an ID for identifying a node.
         self._id = self._rand_id()
@@ -295,6 +299,13 @@ class Node(_Cluster):
 
         super().__init__()
 
+        if icon:
+            _node = icon(_no_init=True)
+            self._icon = _node._icon
+            self._icon_dir = _node._icon_dir
+        if icon_size:
+            self._icon_size = icon_size
+
         # fmt: off
         # If a node has an icon, increase the height slightly to avoid
         # that label being spanned between icon image and white space.
@@ -458,7 +469,7 @@ class Cluster(Node):
         label: str = "",
         direction: str = "LR",
         icon: object = None,
-        icon_size: int = 30,
+        icon_size: int = None,
         **attrs: Dict
         ):
         """Cluster represents a cluster context.
@@ -469,13 +480,7 @@ class Cluster(Node):
         :param icon_size: The icon size. Default is 30.
         """
         self._direction = direction
-        if icon:
-            _node = icon(_no_init=True)
-            self._icon = _node._icon
-            self._icon_dir = _node._icon_dir
-        if icon_size:
-            self._icon_size = icon_size
-        super().__init__(label, **attrs)
+        super().__init__(label, icon, icon_size, **attrs)
 
 
 class Edge: