diff --git a/diagrams/__init__.py b/diagrams/__init__.py index 6e282d9..7e504d2 100644 --- a/diagrams/__init__.py +++ b/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 '<
' + label + '
>' class Diagram: __directions = ("TB", "BT", "LR", "RL") @@ -191,7 +195,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") @@ -216,22 +219,33 @@ class Cluster: label: str = "cluster", direction: str = "LR", graph_attr: dict = {}, + providerIconNode: object = None, + icon_size: int = 30, ): """Cluster represents a cluster context. :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 """ 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') @@ -265,6 +279,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) @@ -284,6 +299,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. diff --git a/tests/test_diagram.py b/tests/test_diagram.py index 00bdacc..c74493a 100644 --- a/tests/test_diagram.py +++ b/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. diff --git a/website/static/img/resources/onprem/registry/jfrog.png b/website/static/img/resources/onprem/registry/jfrog.png new file mode 100644 index 0000000..931bfd8 Binary files /dev/null and b/website/static/img/resources/onprem/registry/jfrog.png differ diff --git a/website/static/img/resources/programming/framework/svelte.png b/website/static/img/resources/programming/framework/svelte.png new file mode 100644 index 0000000..463cc47 Binary files /dev/null and b/website/static/img/resources/programming/framework/svelte.png differ