msb-public
/
diagrams
mirror of https://github.com/mingrammer/diagrams
1
0
Code Issues Projects Releases Wiki Activity

Merge pull request #2 from bkmeneguello/node-as-cluster

Browse Source 
Node as cluster
pull/407/head
dan-ash 5 years ago committed by GitHub
parent d6f73ec1c3 070308c254
commit 88fb1f4d51
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
2 changed files with 182 additions and 136 deletions
Show Stats Download Patch File Download Diff File

311
diagrams/__init__.py
Unescape View File

@ -1,4 +1,5 @@
import contextvars
import html
import os
import uuid
from pathlib import Path
@ -19,7 +20,7 @@ def getdiagram():
try:
return __diagram.get()
except LookupError:
return None
raise EnvironmentError("Global diagrams context not set up")
def setdiagram(diagram):
@ -41,8 +42,72 @@ 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, *args):
setcluster(self._parent)
if not (self.nodes or self.subgraphs):
return
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.remove_node(self.nodeid)
self._parent.subgraph(self)
def node(self, node: "Node") -> None:
"""Create a new node."""
self.nodes[node.nodeid] = node
def remove_node(self, nodeid: str) -> None:
del self.nodes[nodeid]
def subgraph(self, subgraph: "_Cluster") -> None:
"""Create a subgraph for clustering"""
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()
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,15 +171,19 @@ 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"
elif not filename:
filename = "_".join(self.name.split()).lower()
self.filename = filename
self.dot = Digraph(self.name, filename=self.filename)
super().__init__(self.name, filename=self.filename)
self.edges = {}
# 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
@ -142,29 +211,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, *args):
super().__exit__(*args)
setdiagram(None)
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)
def __exit__(self, exc_type, exc_value, traceback):
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:
@ -179,24 +252,16 @@ class Diagram:
return True
return False
def node(self, nodeid: str, label: str, **attrs) -> None:
"""Create a new node."""
self.dot.node(nodeid, label=label, **attrs)
def connect(self, node: "Node", node2: "Node", edge: "Edge") -> None:
"""Connect the two Nodes."""
self.dot.edge(node.nodeid, node2.nodeid, **edge.attrs)
def subgraph(self, dot: Digraph) -> None:
"""Create a subgraph for clustering"""
self.dot.subgraph(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")
class Node(_Cluster):
"""Node represents a node for a specific backend service."""
__bgcolors = ("#E5F5FD", "#EBF3E7", "#ECE8F6", "#FDF7E3")
# fmt: off
@ -209,106 +274,17 @@ class Cluster:
"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)
# 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 = '<<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
# 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):
if self._parent:
self._parent.subgraph(self.dot)
else:
self._diagram.subgraph(self.dot)
setcluster(self._parent)
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.dot.node(nodeid, label=label, **attrs)
def subgraph(self, dot: Digraph) -> None:
self.dot.subgraph(dot)
class Node:
"""Node represents a node for a specific backend service."""
_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,40 +293,81 @@ class Node:
cls.__init__ = new_init(cls, cls.__init__)
return instance
def __init__(self, label: str = "", **attrs: Dict):
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()
self.label = label
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.
# 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)
def __enter__(self):
super().__enter__()
# Set attributes.
for k, v in self._default_graph_attrs.items():
self.dot.graph_attr[k] = v
icon = self._load_icon()
if icon:
lines = iter(html.escape(self.label).split("\n"))
self.dot.graph_attr["label"] = '<<TABLE border="0"><TR>' +\
f'<TD fixedsize="true" width="{self._icon_size}" height="{self._icon_size}"><IMG SRC="{icon}"></IMG></TD>' +\
f'<TD align="left">{next(lines)}</TD></TR>' +\
''.join(f'<TR><TD colspan="2" align="left">{line}</TD></TR>' for line in lines) +\
'</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._parent.depth + 1
coloridx = self.depth % len(self.__bgcolors)
self.dot.graph_attr["bgcolor"] = self.__bgcolors[coloridx]
return self
def __exit__(self, *args):
super().__exit__(*args)
self._id = "cluster_" + self.nodeid
self.dot.name = self.nodeid
def __repr__(self):
_name = self.__class__.__name__
@ -425,7 +442,7 @@ class Node:
@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.
@ -439,7 +456,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
@ -447,8 +464,30 @@ 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 = None,
**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
super().__init__(label, icon, icon_size, **attrs)
class Edge:

7
docs/guides/cluster.md
Unescape View File

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

Write Preview
Loading…
Cancel
Save