From dc330102d569671cf0f117739962de8bd75f9cf6 Mon Sep 17 00:00:00 2001 From: Bruno Meneguello <1322552+bkmeneguello@users.noreply.github.com> Date: Thu, 24 Dec 2020 19:05:09 -0300 Subject: [PATCH 1/9] Allow nodes to be user as cluster --- diagrams/__init__.py | 103 +++++++++++++++++++++++++++++++++++++++-- docs/guides/cluster.md | 7 +++ 2 files changed, 105 insertions(+), 5 deletions(-) diff --git a/diagrams/__init__.py b/diagrams/__init__.py index 8829a012..556e0438 100644 --- a/diagrams/__init__.py +++ b/diagrams/__init__.py @@ -112,7 +112,9 @@ class Diagram: elif not filename: filename = "_".join(self.name.split()).lower() self.filename = filename + self.dot = Digraph(self.name, filename=self.filename) + self._nodes = {} # Set attributes. for k, v in self._default_graph_attrs.items(): @@ -150,6 +152,9 @@ class Diagram: return self def __exit__(self, exc_type, exc_value, traceback): + for nodeid, node in self._nodes.items(): + self.dot.node(nodeid, label=node['label'], **node['attrs']) + self.render() # Remove the graphviz file leaving only the image. os.remove(self.filename) @@ -181,7 +186,10 @@ class Diagram: def node(self, nodeid: str, label: str, **attrs) -> None: """Create a new node.""" - self.dot.node(nodeid, label=label, **attrs) + self._nodes[nodeid] = {'label': label, 'attrs': attrs} + + def remove_node(self, nodeid: str) -> None: + del self._nodes[nodeid] def connect(self, node: "Node", node2: "Node", edge: "Edge") -> None: """Connect the two Nodes.""" @@ -239,6 +247,7 @@ class Cluster: self._icon_size = icon_size self.dot = Digraph(self.name) + self._nodes = {} # Set attributes. for k, v in self._default_graph_attrs.items(): @@ -277,13 +286,16 @@ class Cluster: return self def __exit__(self, exc_type, exc_value, traceback): + for nodeid, node in self._nodes.items(): + self.dot.node(nodeid, label=node['label'], **node['attrs']) + if self._parent: self._parent.subgraph(self.dot) else: self._diagram.subgraph(self.dot) setcluster(self._parent) - def _validate_direction(self, direction: str): + def _validate_direction(self, direction: str) -> bool: direction = direction.upper() for v in self.__directions: if v == direction: @@ -292,7 +304,10 @@ class Cluster: def node(self, nodeid: str, label: str, **attrs) -> None: """Create a new node in the cluster.""" - self.dot.node(nodeid, label=label, **attrs) + self._nodes[nodeid] = {'label': label, 'attrs': attrs} + + def remove_node(self, nodeid: str) -> None: + del self._nodes[nodeid] def subgraph(self, dot: Digraph) -> None: self.dot.subgraph(dot) @@ -300,15 +315,30 @@ class Cluster: class Node: """Node represents a node for a specific backend service.""" + __directions = ("TB", "BT", "LR", "RL") + __bgcolors = ("#E5F5FD", "#EBF3E7", "#ECE8F6", "#FDF7E3") + + # fmt: off + _default_graph_attrs = { + "shape": "box", + "style": "rounded", + "labeljust": "l", + "pencolor": "#AEB6BE", + "fontname": "Sans-Serif", + "fontsize": "12", + } _provider = None _type = None _icon_dir = None _icon = None - + _icon_size = 30 + _direction = "TB" _height = 1.9 + # fmt: on + def __new__(cls, *args, **kwargs): instance = object.__new__(cls) lazy = kwargs.pop('_no_init', False) @@ -317,7 +347,11 @@ class Node: cls.__init__ = new_init(cls, cls.__init__) return instance - def __init__(self, label: str = "", **attrs: Dict): + def __init__( + self, + label: str = "", + **attrs: Dict + ): """Node represents a system component. :param label: Node label. @@ -352,6 +386,65 @@ class Node: else: self._diagram.node(self._id, self.label, **self._attrs) + def __enter__(self): + setcluster(self) + self.name = "cluster_" + self.label + self.dot = Digraph(self.name) + self._nodes = {} + + if self._cluster: + self._cluster.remove_node(self._id) + else: + self._diagram.remove_node(self._id) + + # Set attributes. + for k, v in self._default_graph_attrs.items(): + self.dot.graph_attr[k] = v + + if self._icon: + self.dot.graph_attr["label"] = '<'\ + ''\ + '
'\ + '' + self.label + '
>' + + if not self._validate_direction(self._direction): + raise ValueError(f'"{self._direction}" is not a valid direction') + self.dot.graph_attr["rankdir"] = self._direction + + # Set cluster depth for distinguishing the background color + self.depth = self._cluster.depth + 1 if self._cluster else 0 + coloridx = self.depth % len(self.__bgcolors) + self.dot.graph_attr["bgcolor"] = self.__bgcolors[coloridx] + + return self + + def __exit__(self, exc_type, exc_value, traceback): + for nodeid, node in self._nodes.items(): + self.dot.node(nodeid, label=node['label'], **node['attrs']) + + if self._cluster: + self._cluster.subgraph(self.dot) + else: + self._diagram.subgraph(self.dot) + setcluster(self._cluster) + + def _validate_direction(self, direction: str): + direction = direction.upper() + for v in self.__directions: + if v == direction: + return True + return False + + def node(self, nodeid: str, label: str, **attrs) -> None: + """Create a new node in the cluster.""" + self._nodes[nodeid] = {'label': label, 'attrs': attrs} + + def remove_node(self, nodeid: str) -> None: + del self._nodes[nodeid] + + def subgraph(self, dot: Digraph) -> None: + self.dot.subgraph(dot) + def __repr__(self): _name = self.__class__.__name__ return f"<{self._provider}.{self._type}.{_name}>" diff --git a/docs/guides/cluster.md b/docs/guides/cluster.md index 5001597a..7594e284 100644 --- a/docs/guides/cluster.md +++ b/docs/guides/cluster.md @@ -70,6 +70,9 @@ with Diagram("Event Processing", show=False): 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. +It's also possible to use the node in the `with` context adding `cluster=True` to +make it behave like a cluster. + ```python from diagrams import Cluster, Diagram from diagrams.aws.compute import ECS @@ -85,6 +88,10 @@ with Diagram("Simple Web Service with DB Cluster", show=False): db_master = RDS("master") db_master - [RDS("slave1"), RDS("slave2")] + with Aurora("DB Cluster", cluster=True): + db_master = RDS("master") + db_master - [RDS("slave1"), + RDS("slave2")] dns >> web >> db_master ``` From 496a69cb4f8ec8cfe5a6d6f4d82b9db14215682a Mon Sep 17 00:00:00 2001 From: Bruno Meneguello <1322552+bkmeneguello@users.noreply.github.com> Date: Mon, 28 Dec 2020 16:52:17 -0300 Subject: [PATCH 2/9] Allow node to cluster edges --- diagrams/__init__.py | 27 ++++++++++++++++++++------- 1 file changed, 20 insertions(+), 7 deletions(-) diff --git a/diagrams/__init__.py b/diagrams/__init__.py index 556e0438..4b19d9bf 100644 --- a/diagrams/__init__.py +++ b/diagrams/__init__.py @@ -115,7 +115,9 @@ class Diagram: self.dot = Digraph(self.name, filename=self.filename) self._nodes = {} + self._edges = {} + self.dot.attr(compound="true") # Set attributes. for k, v in self._default_graph_attrs.items(): self.dot.graph_attr[k] = v @@ -155,6 +157,17 @@ class Diagram: for nodeid, node in self._nodes.items(): self.dot.node(nodeid, label=node['label'], **node['attrs']) + for nodes, edge in self._edges.items(): + node1, node2 = nodes + nodeid1, nodeid2 = node1.nodeid, node2.nodeid + if hasattr(node1, '_nodes') and node1._nodes: + edge._attrs['ltail'] = nodeid1 + nodeid1 = next(iter(node1._nodes.keys())) + if hasattr(node2, '_nodes') and node2._nodes: + edge._attrs['lhead'] = nodeid2 + nodeid2 = next(iter(node2._nodes.keys())) + self.dot.edge(nodeid1, nodeid2, **edge.attrs) + self.render() # Remove the graphviz file leaving only the image. os.remove(self.filename) @@ -193,7 +206,7 @@ class Diagram: def connect(self, node: "Node", node2: "Node", edge: "Edge") -> None: """Connect the two Nodes.""" - self.dot.edge(node.nodeid, node2.nodeid, **edge.attrs) + self._edges[(node, node2)] = edge def subgraph(self, dot: Digraph) -> None: """Create a subgraph for clustering""" @@ -387,16 +400,16 @@ class Node: self._diagram.node(self._id, self.label, **self._attrs) def __enter__(self): - setcluster(self) - self.name = "cluster_" + self.label - self.dot = Digraph(self.name) - self._nodes = {} - if self._cluster: self._cluster.remove_node(self._id) else: self._diagram.remove_node(self._id) + setcluster(self) + self._id = "cluster_" + self.label + self.dot = Digraph(self._id) + self._nodes = {} + # Set attributes. for k, v in self._default_graph_attrs.items(): self.dot.graph_attr[k] = v @@ -421,7 +434,7 @@ class Node: def __exit__(self, exc_type, exc_value, traceback): for nodeid, node in self._nodes.items(): self.dot.node(nodeid, label=node['label'], **node['attrs']) - + if self._cluster: self._cluster.subgraph(self.dot) else: From 0115211b3440bb732bf7552db0532ff2cbd8388d Mon Sep 17 00:00:00 2001 From: Bruno Meneguello <1322552+bkmeneguello@users.noreply.github.com> Date: Tue, 29 Dec 2020 10:27:24 -0300 Subject: [PATCH 3/9] Allow Node as Cluster with same label --- diagrams/__init__.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/diagrams/__init__.py b/diagrams/__init__.py index 4b19d9bf..55f6389d 100644 --- a/diagrams/__init__.py +++ b/diagrams/__init__.py @@ -406,7 +406,7 @@ class Node: self._diagram.remove_node(self._id) setcluster(self) - self._id = "cluster_" + self.label + self._id = "cluster_" + self._id self.dot = Digraph(self._id) self._nodes = {} From fc9ac5bfae369ad7c2805951d4b3fd81fdd784e6 Mon Sep 17 00:00:00 2001 From: Bruno Meneguello <1322552+bkmeneguello@users.noreply.github.com> Date: Tue, 29 Dec 2020 11:33:21 -0300 Subject: [PATCH 4/9] Allow for empty "Node as Cluster" to render as Node --- diagrams/__init__.py | 39 ++++++++++++++++++++++++++++----------- 1 file changed, 28 insertions(+), 11 deletions(-) diff --git a/diagrams/__init__.py b/diagrams/__init__.py index 55f6389d..e9da8039 100644 --- a/diagrams/__init__.py +++ b/diagrams/__init__.py @@ -116,9 +116,10 @@ class Diagram: self.dot = Digraph(self.name, filename=self.filename) self._nodes = {} self._edges = {} + self._subgraphs = [] - self.dot.attr(compound="true") # Set attributes. + self.dot.attr(compound="true") for k, v in self._default_graph_attrs.items(): self.dot.graph_attr[k] = v self.dot.graph_attr["label"] = self.name @@ -156,6 +157,9 @@ class Diagram: def __exit__(self, exc_type, exc_value, traceback): for nodeid, node in self._nodes.items(): self.dot.node(nodeid, label=node['label'], **node['attrs']) + + for dot in self._subgraphs: + self.dot.subgraph(dot) for nodes, edge in self._edges.items(): node1, node2 = nodes @@ -210,7 +214,7 @@ class Diagram: def subgraph(self, dot: Digraph) -> None: """Create a subgraph for clustering""" - self.dot.subgraph(dot) + self._subgraphs.append(dot) def render(self) -> None: self.dot.render(format=self.outformat, view=self.show, quiet=True) @@ -261,6 +265,7 @@ class Cluster: self.dot = Digraph(self.name) self._nodes = {} + self._subgraphs = [] # Set attributes. for k, v in self._default_graph_attrs.items(): @@ -301,6 +306,9 @@ class Cluster: def __exit__(self, exc_type, exc_value, traceback): for nodeid, node in self._nodes.items(): self.dot.node(nodeid, label=node['label'], **node['attrs']) + + for dot in self._subgraphs: + self.dot.subgraph(dot) if self._parent: self._parent.subgraph(self.dot) @@ -323,7 +331,7 @@ class Cluster: del self._nodes[nodeid] def subgraph(self, dot: Digraph) -> None: - self.dot.subgraph(dot) + self._subgraphs.append(dot) class Node: @@ -400,15 +408,10 @@ class Node: self._diagram.node(self._id, self.label, **self._attrs) def __enter__(self): - if self._cluster: - self._cluster.remove_node(self._id) - else: - self._diagram.remove_node(self._id) - setcluster(self) - self._id = "cluster_" + self._id - self.dot = Digraph(self._id) + self.dot = Digraph() self._nodes = {} + self._subgraphs = [] # Set attributes. for k, v in self._default_graph_attrs.items(): @@ -432,8 +435,22 @@ class Node: return self def __exit__(self, exc_type, exc_value, traceback): + if not (self._nodes or self._subgraphs): + return + + if self._cluster: + self._cluster.remove_node(self._id) + else: + self._diagram.remove_node(self._id) + + self._id = "cluster_" + self._id + self.dot.name = self._id + for nodeid, node in self._nodes.items(): self.dot.node(nodeid, label=node['label'], **node['attrs']) + + for dot in self._subgraphs: + self.dot.subgraph(dot) if self._cluster: self._cluster.subgraph(self.dot) @@ -456,7 +473,7 @@ class Node: del self._nodes[nodeid] def subgraph(self, dot: Digraph) -> None: - self.dot.subgraph(dot) + self._subgraphs.append(dot) def __repr__(self): _name = self.__class__.__name__ From bf6a41a512caba6efa5d54025677c7cf2f5321c1 Mon Sep 17 00:00:00 2001 From: Bruno Meneguello <1322552+bkmeneguello@users.noreply.github.com> Date: Tue, 29 Dec 2020 15:25:55 -0300 Subject: [PATCH 5/9] Convert the "Cluster" class to a "Node" and extract commons to a base class --- diagrams/__init__.py | 322 +++++++++++++++---------------------------- 1 file changed, 113 insertions(+), 209 deletions(-) diff --git a/diagrams/__init__.py b/diagrams/__init__.py index e9da8039..df6a8f17 100644 --- a/diagrams/__init__.py +++ b/diagrams/__init__.py @@ -19,7 +19,7 @@ def getdiagram(): try: return __diagram.get() except LookupError: - return None + raise EnvironmentError("Global diagrams context not set up") def setdiagram(diagram): @@ -41,8 +41,60 @@ def new_init(cls, init): cls.__init__ = init return reset_init -class Diagram: +class _Cluster: __directions = ("TB", "BT", "LR", "RL") + + def __init__(self, name=None, **kwargs): + self.dot = Digraph(name, **kwargs) + self.depth = 0 + self.nodes = {} + self.subgraphs = [] + + try: + self._parent = getcluster() or getdiagram() + except EnvironmentError: + self._parent = None + + + def __enter__(self): + setcluster(self) + return self + + def __exit__(self, exc_type, exc_value, traceback): + setcluster(self._parent) + + for nodeid, node in self.nodes.items(): + self.dot.node(nodeid, label=node['label'], **node['attrs']) + + for dot in self.subgraphs: + self.dot.subgraph(dot) + + if self._parent: + self._parent.subgraph(self.dot) + + def node(self, nodeid: str, label: str, **attrs) -> None: + """Create a new node.""" + self.nodes[nodeid] = {'label': label, 'attrs': attrs} + + def remove_node(self, nodeid: str) -> None: + del self.nodes[nodeid] + + def subgraph(self, dot: Digraph) -> None: + """Create a subgraph for clustering""" + self.subgraphs.append(dot) + + def _validate_direction(self, direction: str): + direction = direction.upper() + for v in self.__directions: + if v == direction: + return True + return False + + def __str__(self) -> str: + return str(self.dot) + + +class Diagram(_Cluster): __curvestyles = ("ortho", "curved") __outformats = ("png", "jpg", "svg", "pdf") @@ -106,6 +158,7 @@ class Diagram: :param node_attr: Provide node_attr dot config attributes. :param edge_attr: Provide edge_attr dot config attributes. """ + self.name = name if not name and not filename: filename = "diagrams_image" @@ -113,10 +166,8 @@ class Diagram: filename = "_".join(self.name.split()).lower() self.filename = filename - self.dot = Digraph(self.name, filename=self.filename) - self._nodes = {} - self._edges = {} - self._subgraphs = [] + super().__init__(self.name, filename=self.filename) + self.edges = {} # Set attributes. self.dot.attr(compound="true") @@ -147,46 +198,33 @@ class Diagram: self.show = show - def __str__(self) -> str: - return str(self.dot) - def __enter__(self): setdiagram(self) + super().__enter__() return self def __exit__(self, exc_type, exc_value, traceback): - for nodeid, node in self._nodes.items(): - self.dot.node(nodeid, label=node['label'], **node['attrs']) - - for dot in self._subgraphs: - self.dot.subgraph(dot) + super().__exit__(exc_type, exc_value, traceback) + setdiagram(None) - for nodes, edge in self._edges.items(): + for nodes, edge in self.edges.items(): node1, node2 = nodes nodeid1, nodeid2 = node1.nodeid, node2.nodeid - if hasattr(node1, '_nodes') and node1._nodes: + if node1.nodes: edge._attrs['ltail'] = nodeid1 - nodeid1 = next(iter(node1._nodes.keys())) - if hasattr(node2, '_nodes') and node2._nodes: + nodeid1 = next(iter(node1.nodes.keys())) + if node2.nodes: edge._attrs['lhead'] = nodeid2 - nodeid2 = next(iter(node2._nodes.keys())) + nodeid2 = next(iter(node2.nodes.keys())) self.dot.edge(nodeid1, nodeid2, **edge.attrs) self.render() # Remove the graphviz file leaving only the image. os.remove(self.filename) - setdiagram(None) def _repr_png_(self): return self.dot.pipe(format="png") - def _validate_direction(self, direction: str) -> bool: - direction = direction.upper() - for v in self.__directions: - if v == direction: - return True - return False - def _validate_curvestyle(self, curvestyle: str) -> bool: curvestyle = curvestyle.lower() for v in self.__curvestyles: @@ -201,142 +239,16 @@ class Diagram: return True return False - def node(self, nodeid: str, label: str, **attrs) -> None: - """Create a new node.""" - self._nodes[nodeid] = {'label': label, 'attrs': attrs} - - def remove_node(self, nodeid: str) -> None: - del self._nodes[nodeid] - def connect(self, node: "Node", node2: "Node", edge: "Edge") -> None: """Connect the two Nodes.""" - self._edges[(node, node2)] = edge - - def subgraph(self, dot: Digraph) -> None: - """Create a subgraph for clustering""" - self._subgraphs.append(dot) + self.edges[(node, node2)] = edge def render(self) -> None: self.dot.render(format=self.outformat, view=self.show, quiet=True) -class Cluster: - __directions = ("TB", "BT", "LR", "RL") - __bgcolors = ("#E5F5FD", "#EBF3E7", "#ECE8F6", "#FDF7E3") - - # fmt: off - _default_graph_attrs = { - "shape": "box", - "style": "rounded", - "labeljust": "l", - "pencolor": "#AEB6BE", - "fontname": "Sans-Serif", - "fontsize": "12", - } - - _icon = None - _icon_size = 0 - - # fmt: on - - # 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", - graph_attr: dict = {}, - icon: 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. - """ - self.label = label - self.name = "cluster_" + self.label - if not self._icon: - self._icon = icon - if not self._icon_size: - self._icon_size = icon_size - - self.dot = Digraph(self.name) - self._nodes = {} - self._subgraphs = [] - - # Set attributes. - for k, v in self._default_graph_attrs.items(): - self.dot.graph_attr[k] = v - - # 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 = '<
' + self.label + '
>' - 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 - - # Node must be belong to a diagrams. - self._diagram = getdiagram() - if self._diagram is None: - raise EnvironmentError("Global diagrams context not set up") - self._parent = getcluster() - - # Set cluster depth for distinguishing the background color - self.depth = self._parent.depth + 1 if self._parent else 0 - coloridx = self.depth % len(self.__bgcolors) - self.dot.graph_attr["bgcolor"] = self.__bgcolors[coloridx] - - # Merge passed in attributes - self.dot.graph_attr.update(graph_attr) - - def __enter__(self): - setcluster(self) - return self - - def __exit__(self, exc_type, exc_value, traceback): - for nodeid, node in self._nodes.items(): - self.dot.node(nodeid, label=node['label'], **node['attrs']) - - for dot in self._subgraphs: - self.dot.subgraph(dot) - - if self._parent: - self._parent.subgraph(self.dot) - else: - self._diagram.subgraph(self.dot) - setcluster(self._parent) - - def _validate_direction(self, direction: str) -> bool: - direction = direction.upper() - for v in self.__directions: - if v == direction: - return True - return False - - def node(self, nodeid: str, label: str, **attrs) -> None: - """Create a new node in the cluster.""" - self._nodes[nodeid] = {'label': label, 'attrs': attrs} - - def remove_node(self, nodeid: str) -> None: - del self._nodes[nodeid] - - def subgraph(self, dot: Digraph) -> None: - self._subgraphs.append(dot) - - -class Node: +class Node(_Cluster): """Node represents a node for a specific backend service.""" - __directions = ("TB", "BT", "LR", "RL") __bgcolors = ("#E5F5FD", "#EBF3E7", "#ECE8F6", "#FDF7E3") # fmt: off @@ -381,46 +293,39 @@ class Node: self._id = self._rand_id() self.label = label + super().__init__() + # fmt: off # If a node has an icon, increase the height slightly to avoid # that label being spanned between icon image and white space. # Increase the height by the number of new lines included in the label. padding = 0.4 * (label.count('\n')) + icon = self._load_icon() self._attrs = { "shape": "none", "height": str(self._height + padding), - "image": self._load_icon(), - } if self._icon else {} + "image": icon, + } if icon else {} # fmt: on self._attrs.update(attrs) - # Node must be belong to a diagrams. - self._diagram = getdiagram() - if self._diagram is None: - raise EnvironmentError("Global diagrams context not set up") - self._cluster = getcluster() - # If a node is in the cluster context, add it to cluster. - if self._cluster: - self._cluster.node(self._id, self.label, **self._attrs) - else: - self._diagram.node(self._id, self.label, **self._attrs) + self._parent.node(self._id, self.label, **self._attrs) def __enter__(self): + super().__enter__() setcluster(self) - self.dot = Digraph() - self._nodes = {} - self._subgraphs = [] # Set attributes. for k, v in self._default_graph_attrs.items(): self.dot.graph_attr[k] = v - if self._icon: + icon = self._load_icon() + if icon: self.dot.graph_attr["label"] = '<'\ ''\ + ''\ '
'\ - '' + self.label + '
>' if not self._validate_direction(self._direction): @@ -428,52 +333,23 @@ class Node: self.dot.graph_attr["rankdir"] = self._direction # Set cluster depth for distinguishing the background color - self.depth = self._cluster.depth + 1 if self._cluster else 0 + self.depth = self._parent.depth + 1 coloridx = self.depth % len(self.__bgcolors) self.dot.graph_attr["bgcolor"] = self.__bgcolors[coloridx] return self def __exit__(self, exc_type, exc_value, traceback): - if not (self._nodes or self._subgraphs): + if not (self.nodes or self.subgraphs): return - if self._cluster: - self._cluster.remove_node(self._id) - else: - self._diagram.remove_node(self._id) - + self._parent.remove_node(self._id) + self._id = "cluster_" + self._id self.dot.name = self._id - for nodeid, node in self._nodes.items(): - self.dot.node(nodeid, label=node['label'], **node['attrs']) - - for dot in self._subgraphs: - self.dot.subgraph(dot) - - if self._cluster: - self._cluster.subgraph(self.dot) - else: - self._diagram.subgraph(self.dot) - setcluster(self._cluster) + super().__exit__(exc_type, exc_value, traceback) - def _validate_direction(self, direction: str): - direction = direction.upper() - for v in self.__directions: - if v == direction: - return True - return False - - def node(self, nodeid: str, label: str, **attrs) -> None: - """Create a new node in the cluster.""" - self._nodes[nodeid] = {'label': label, 'attrs': attrs} - - def remove_node(self, nodeid: str) -> None: - del self._nodes[nodeid] - - def subgraph(self, dot: Digraph) -> None: - self._subgraphs.append(dot) def __repr__(self): _name = self.__class__.__name__ @@ -562,7 +438,7 @@ class Node: if not isinstance(node, Edge): ValueError(f"{node} is not a valid Edge") # An edge must be added on the global diagrams, not a cluster. - self._diagram.connect(self, node, edge) + getdiagram().connect(self, node, edge) return node @staticmethod @@ -570,8 +446,36 @@ class Node: return uuid.uuid4().hex def _load_icon(self): - basedir = Path(os.path.abspath(os.path.dirname(__file__))) - return os.path.join(basedir.parent, self._icon_dir, self._icon) + if self._icon and self._icon_dir: + basedir = Path(os.path.abspath(os.path.dirname(__file__))) + return os.path.join(basedir.parent, self._icon_dir, self._icon) + return None + + +class Cluster(Node): + def __init__( + self, + label: str = "", + direction: str = "LR", + icon: object = None, + icon_size: int = 30, + **attrs: Dict + ): + """Cluster represents a cluster context. + + :param label: Cluster label. + :param direction: Data flow direction. Default is "LR" (left to right). + :param icon: Custom icon for tihs cluster. Must be a node class or reference. + :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) class Edge: 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 6/9] 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 df6a8f17..c307cb2f 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: From 91b57985db4041075b08647bf3a00b47d677ca8b Mon Sep 17 00:00:00 2001 From: Bruno Meneguello <1322552+bkmeneguello@users.noreply.github.com> Date: Tue, 29 Dec 2020 17:39:18 -0300 Subject: [PATCH 7/9] Fixed issue with nested Clusters and Cluster >> Node edges --- diagrams/__init__.py | 78 +++++++++++++++++++++++--------------------- 1 file changed, 41 insertions(+), 37 deletions(-) diff --git a/diagrams/__init__.py b/diagrams/__init__.py index c307cb2f..1f953e6d 100644 --- a/diagrams/__init__.py +++ b/diagrams/__init__.py @@ -60,28 +60,40 @@ class _Cluster: setcluster(self) return self - def __exit__(self, exc_type, exc_value, traceback): + def __exit__(self, *args): setcluster(self._parent) - for nodeid, node in self.nodes.items(): - self.dot.node(nodeid, label=node['label'], **node['attrs']) + if not (self.nodes or self.subgraphs): + return - for dot in self.subgraphs: - self.dot.subgraph(dot) + for node in self.nodes.values(): + self.dot.node(node.nodeid, label=node.label, **node._attrs) + + for subgraph in self.subgraphs: + self.dot.subgraph(subgraph.dot) if self._parent: - self._parent.subgraph(self.dot) + self._parent.remove_node(self.nodeid) + self._parent.subgraph(self) - def node(self, nodeid: str, label: str, **attrs) -> None: + def node(self, node: "Node") -> None: """Create a new node.""" - self.nodes[nodeid] = {'label': label, 'attrs': attrs} + self.nodes[node.nodeid] = node def remove_node(self, nodeid: str) -> None: del self.nodes[nodeid] - def subgraph(self, dot: Digraph) -> None: + def subgraph(self, subgraph: "_Cluster") -> None: """Create a subgraph for clustering""" - self.subgraphs.append(dot) + self.subgraphs.append(subgraph) + + @property + def nodes_iter(self): + if self.nodes: + yield from self.nodes.values() + if self.subgraphs: + for subgraph in self.subgraphs: + yield from subgraph.nodes_iter def _validate_direction(self, direction: str): direction = direction.upper() @@ -202,21 +214,21 @@ class Diagram(_Cluster): setdiagram(self) super().__enter__() return self - - def __exit__(self, exc_type, exc_value, traceback): - super().__exit__(exc_type, exc_value, traceback) + + def __exit__(self, *args): + super().__exit__(*args) setdiagram(None) - for nodes, edge in self.edges.items(): - node1, node2 = nodes - nodeid1, nodeid2 = node1.nodeid, node2.nodeid - if node1.nodes: - edge._attrs['ltail'] = nodeid1 - nodeid1 = next(iter(node1.nodes.keys())) - if node2.nodes: - edge._attrs['lhead'] = nodeid2 - nodeid2 = next(iter(node2.nodes.keys())) - self.dot.edge(nodeid1, nodeid2, **edge.attrs) + for (node1, node2), edge in self.edges.items(): + cluster_node1 = next(node1.nodes_iter, None) + if cluster_node1: + edge._attrs['ltail'] = node1.nodeid + node1 = cluster_node1 + cluster_node2 = next(node2.nodes_iter, None) + if cluster_node2: + edge._attrs['lhead'] = node2.nodeid + node2 = cluster_node2 + self.dot.edge(node1.nodeid, node2.nodeid, **edge.attrs) self.render() # Remove the graphviz file leaving only the image. @@ -322,11 +334,10 @@ class Node(_Cluster): self._attrs.update(attrs) # If a node is in the cluster context, add it to cluster. - self._parent.node(self._id, self.label, **self._attrs) + self._parent.node(self) def __enter__(self): super().__enter__() - setcluster(self) # Set attributes. for k, v in self._default_graph_attrs.items(): @@ -350,17 +361,10 @@ class Node(_Cluster): return self - def __exit__(self, exc_type, exc_value, traceback): - if not (self.nodes or self.subgraphs): - return - - self._parent.remove_node(self._id) - - self._id = "cluster_" + self._id - self.dot.name = self._id - - super().__exit__(exc_type, exc_value, traceback) - + def __exit__(self, *args): + super().__exit__(*args) + self._id = "cluster_" + self.nodeid + self.dot.name = self.nodeid def __repr__(self): _name = self.__class__.__name__ @@ -435,7 +439,7 @@ class Node(_Cluster): @property def nodeid(self): return self._id - + # TODO: option for adding flow description to the connection edge def connect(self, node: "Node", edge: "Edge"): """Connect to other node. From 417047dae79c902c91b48940bb36214f1fd9dff0 Mon Sep 17 00:00:00 2001 From: Bruno Meneguello <1322552+bkmeneguello@users.noreply.github.com> Date: Tue, 29 Dec 2020 18:19:27 -0300 Subject: [PATCH 8/9] Escape HTML entities of Cluster label --- diagrams/__init__.py | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/diagrams/__init__.py b/diagrams/__init__.py index 1f953e6d..19b23152 100644 --- a/diagrams/__init__.py +++ b/diagrams/__init__.py @@ -1,4 +1,5 @@ import contextvars +import html import os import uuid from pathlib import Path @@ -348,7 +349,7 @@ class Node(_Cluster): self.dot.graph_attr["label"] = '<'\ ''\ - '
'\ '' + self.label + '
>' + '' + html.escape(self.label) + '>' if not self._validate_direction(self._direction): raise ValueError(f'"{self._direction}" is not a valid direction') From 070308c25427ab0a288aa0222716f761e6cd7420 Mon Sep 17 00:00:00 2001 From: Bruno Meneguello <1322552+bkmeneguello@users.noreply.github.com> Date: Tue, 29 Dec 2020 21:24:53 -0300 Subject: [PATCH 9/9] Better handling of multiline Cluster labels --- diagrams/__init__.py | 10 ++++++---- 1 file changed, 6 insertions(+), 4 deletions(-) diff --git a/diagrams/__init__.py b/diagrams/__init__.py index 19b23152..ca84e75c 100644 --- a/diagrams/__init__.py +++ b/diagrams/__init__.py @@ -346,10 +346,12 @@ class Node(_Cluster): icon = self._load_icon() if icon: - self.dot.graph_attr["label"] = '<'\ - ''\ - '
'\ - '' + html.escape(self.label) + '
>' + lines = iter(html.escape(self.label).split("\n")) + self.dot.graph_attr["label"] = '<' +\ + f'' +\ + f'' +\ + ''.join(f'' for line in lines) +\ + '
{next(lines)}
{line}
>' if not self._validate_direction(self._direction): raise ValueError(f'"{self._direction}" is not a valid direction')