|
1
|
|
|
import base64 |
|
2
|
|
|
import struct |
|
3
|
|
|
from typing import Any |
|
4
|
|
|
|
|
5
|
|
|
import networkx as nx |
|
6
|
|
|
|
|
7
|
|
|
from .models import DiagramNodeShape, DiagramOrientation |
|
8
|
|
|
from .typing import MermaidDiagram |
|
9
|
|
|
|
|
10
|
|
|
DEFAULT_LAYOUT = "dagre" |
|
11
|
|
|
DEFAULT_LOOK = "neo" |
|
12
|
|
|
DEFAULT_THEME = "neutral" |
|
13
|
|
|
|
|
14
|
|
|
|
|
15
|
|
|
def _edge_label(data: dict[str, Any]) -> str: |
|
16
|
|
|
"""Generate an edge label string.""" |
|
17
|
|
|
label = data.get("label") |
|
18
|
|
|
return f"|{label}|" if label else "" |
|
19
|
|
|
|
|
20
|
|
|
|
|
21
|
|
|
def _contrast_color(color: str) -> str: |
|
22
|
|
|
""" |
|
23
|
|
|
Return black or white by choosing the best contrast to input color. |
|
24
|
|
|
|
|
25
|
|
|
Args: |
|
26
|
|
|
color: str - hex color code |
|
27
|
|
|
|
|
28
|
|
|
Returns: |
|
29
|
|
|
color: str - hex color code |
|
30
|
|
|
""" |
|
31
|
|
|
if not (isinstance(color, str) and color.startswith("#") and len(color) == 7): |
|
32
|
|
|
raise ValueError(f"Invalid color format: {color}. Expected a 6-digit hex code.") |
|
33
|
|
|
|
|
34
|
|
|
r, g, b = int(color[1:3], 16), int(color[3:5], 16), int(color[5:7], 16) |
|
35
|
|
|
return "#000000" if (r * 0.299 + g * 0.587 + b * 0.114) > 186 else "#ffffff" |
|
36
|
|
|
|
|
37
|
|
|
|
|
38
|
|
|
def _node_style(node_id: str, data: dict[str, Any]) -> str: |
|
39
|
|
|
"""Generate a node style string.""" |
|
40
|
|
|
color = data.get("color") |
|
41
|
|
|
if color: |
|
42
|
|
|
return f"\nstyle {node_id} fill:{color}, color:{_contrast_color(color)}" |
|
43
|
|
|
return "" |
|
44
|
|
|
|
|
45
|
|
|
|
|
46
|
|
|
def _graph_title(graph: nx.Graph) -> str: |
|
47
|
|
|
"""Generate a graph title string.""" |
|
48
|
|
|
return f"title: {graph.name}\n" if graph.name else "" |
|
49
|
|
|
|
|
50
|
|
|
|
|
51
|
|
|
def _node_id(node_id) -> str: |
|
52
|
|
|
"""Generate a node id string.""" |
|
53
|
|
|
n_id = base64.b64encode(struct.pack('>h', node_id)).decode().strip('=') |
|
54
|
|
|
return n_id |
|
55
|
|
|
|
|
56
|
|
|
|
|
57
|
|
|
class DiagramBuilder: |
|
58
|
|
|
""" |
|
59
|
|
|
A class to generate Mermaid diagrams from NetworkX graphs. |
|
60
|
|
|
""" |
|
61
|
|
|
|
|
62
|
|
|
def __init__( |
|
63
|
|
|
self, |
|
64
|
|
|
orientation: DiagramOrientation = DiagramOrientation.LEFT_RIGHT, |
|
65
|
|
|
node_shape: DiagramNodeShape = DiagramNodeShape.DEFAULT, |
|
66
|
|
|
layout: str = DEFAULT_LAYOUT, |
|
67
|
|
|
look: str = DEFAULT_LOOK, |
|
68
|
|
|
theme: str = DEFAULT_THEME, |
|
69
|
|
|
): |
|
70
|
|
|
""" |
|
71
|
|
|
Initialize the MermaidDiagramGenerator. |
|
72
|
|
|
|
|
73
|
|
|
Args: |
|
74
|
|
|
orientation: The orientation of the graph (default: LEFT_RIGHT). |
|
75
|
|
|
node_shape: The shape of the nodes (default: DEFAULT). |
|
76
|
|
|
layout: the layout to use (default: 'dagre') |
|
77
|
|
|
look: the look to use (default: 'neo') |
|
78
|
|
|
theme: the theme to use (default: 'neutral') |
|
79
|
|
|
""" |
|
80
|
|
|
self.orientation = orientation |
|
81
|
|
|
self.node_shape = node_shape |
|
82
|
|
|
self.layout = layout |
|
83
|
|
|
self.look = look |
|
84
|
|
|
self.theme = theme |
|
85
|
|
|
|
|
86
|
|
|
if not isinstance(orientation, DiagramOrientation): |
|
87
|
|
|
raise TypeError("orientation must be a valid Orientation enum") |
|
88
|
|
|
if not isinstance(node_shape, DiagramNodeShape): |
|
89
|
|
|
raise TypeError("node_shape must be a valid NodeShape enum") |
|
90
|
|
|
|
|
91
|
|
|
def build(self, graph: nx.Graph) -> MermaidDiagram: |
|
92
|
|
|
""" |
|
93
|
|
|
Materialize a graph as a Mermaid flowchart. |
|
94
|
|
|
|
|
95
|
|
|
Returns: |
|
96
|
|
|
A string representation of the graph as a Mermaid graph. |
|
97
|
|
|
""" |
|
98
|
|
|
config = ( |
|
99
|
|
|
f"---\n" |
|
100
|
|
|
f"{_graph_title(graph)}" |
|
101
|
|
|
f"config:\n" |
|
102
|
|
|
f" layout: {self.layout}\n" |
|
103
|
|
|
f" look: {self.look}\n" |
|
104
|
|
|
f" theme: {self.theme}\n" |
|
105
|
|
|
f"---\n" |
|
106
|
|
|
) |
|
107
|
|
|
|
|
108
|
|
|
bra, ket = self.node_shape.value |
|
109
|
|
|
|
|
110
|
|
|
nodes = "\n".join( |
|
111
|
|
|
f"{_node_id(u)}{bra}{d.get('label', u)}{ket}{_node_style(_node_id(u), d)}" for u, d in graph.nodes.data()) |
|
112
|
|
|
|
|
113
|
|
|
_edges = ((_node_id(u), _node_id(v), d) for u, v, d in graph.edges.data()) |
|
|
|
|
|
|
114
|
|
|
edges = "\n".join(f"{u} -->{_edge_label(d)} {v}" for u, v, d in _edges) |
|
115
|
|
|
|
|
116
|
|
|
return ( |
|
117
|
|
|
f"{config}" |
|
118
|
|
|
f"graph {self.orientation.value}\n" |
|
119
|
|
|
f"{nodes}\n" |
|
120
|
|
|
f"{edges}" |
|
121
|
|
|
) |
|
122
|
|
|
|
erivlis /
networkx-mermaid
Loading history...