Allow nodes to be user as cluster

pull/407/head^2
Bruno Meneguello 5 years ago
parent d6f73ec1c3
commit dc330102d5

@ -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"] = '<<TABLE border="0"><TR>'\
'<TD fixedsize="true" width="' + str(self._icon_size) + '" height="' + str(self._icon_size) + '">'\
'<IMG SRC="' + self._load_icon() + '"></IMG></TD>'\
'<TD>' + self.label + '</TD></TR></TABLE>>'
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}>"

@ -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
```

Loading…
Cancel
Save