|
1
|
|
|
from functools import lru_cache |
|
2
|
|
|
from typing import Any |
|
3
|
|
|
|
|
4
|
|
|
import networkx as nx |
|
5
|
|
|
from mappingtools.collectors import AutoMapper |
|
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
|
|
|
@lru_cache(maxsize=1024) |
|
22
|
|
|
def _contrast_color(color: str) -> str: |
|
23
|
|
|
""" |
|
24
|
|
|
Return black or white by choosing the best contrast to input color. |
|
25
|
|
|
|
|
26
|
|
|
Args: |
|
27
|
|
|
color: str - hex color code |
|
28
|
|
|
|
|
29
|
|
|
Returns: |
|
30
|
|
|
color: str - hex color code |
|
31
|
|
|
""" |
|
32
|
|
|
if not (isinstance(color, str) and color.startswith("#") and len(color) == 7): |
|
33
|
|
|
raise ValueError(f"Invalid color format: {color}. Expected a 6-digit hex code.") |
|
34
|
|
|
|
|
35
|
|
|
# Parsing hex string as single integer with bitwise operations is faster than string slicing |
|
36
|
|
|
rgb = int(color[1:], 16) |
|
37
|
|
|
r = (rgb >> 16) & 0xFF |
|
38
|
|
|
g = (rgb >> 8) & 0xFF |
|
39
|
|
|
b = rgb & 0xFF |
|
40
|
|
|
|
|
41
|
|
|
# Using scaled integer arithmetic (x1000) instead of floating point math |
|
42
|
|
|
# Threshold: 186 * 1000 = 186000 |
|
43
|
|
|
return "#000000" if (r * 299 + g * 587 + b * 114) > 186000 else "#ffffff" |
|
44
|
|
|
|
|
45
|
|
|
|
|
46
|
|
|
def _node_style(node_id: str, data: dict[str, Any]) -> str: |
|
47
|
|
|
"""Generate a node style string.""" |
|
48
|
|
|
color = data.get("color") |
|
49
|
|
|
if color: |
|
50
|
|
|
return f"\nstyle {node_id} fill:{color}, color:{_contrast_color(color)}" |
|
51
|
|
|
return "" |
|
52
|
|
|
|
|
53
|
|
|
|
|
54
|
|
|
def _graph_title(graph: nx.Graph, title: str | None = None) -> str: |
|
55
|
|
|
"""Generate a graph title string.""" |
|
56
|
|
|
title = title if title is not None else graph.name |
|
57
|
|
|
return f"title: {title}\n" if title else "" |
|
58
|
|
|
|
|
59
|
|
|
|
|
60
|
|
|
class DiagramBuilder: |
|
61
|
|
|
""" |
|
62
|
|
|
A class to generate Mermaid diagrams from NetworkX graphs. |
|
63
|
|
|
""" |
|
64
|
|
|
|
|
65
|
|
|
def __init__( |
|
66
|
|
|
self, |
|
67
|
|
|
orientation: DiagramOrientation = DiagramOrientation.LEFT_RIGHT, |
|
68
|
|
|
node_shape: DiagramNodeShape = DiagramNodeShape.DEFAULT, |
|
69
|
|
|
layout: str = DEFAULT_LAYOUT, |
|
70
|
|
|
look: str = DEFAULT_LOOK, |
|
71
|
|
|
theme: str = DEFAULT_THEME, |
|
72
|
|
|
): |
|
73
|
|
|
""" |
|
74
|
|
|
Initialize the DiagramBuilder. |
|
75
|
|
|
|
|
76
|
|
|
Args: |
|
77
|
|
|
orientation: DiagramOrientation - The orientation of the graph (default: LEFT_RIGHT). |
|
78
|
|
|
node_shape: DiagramNodeShape - The shape of the nodes (default: DiagramNodeShape.DEFAULT). |
|
79
|
|
|
layout: str - the layout to use (default: 'dagre') |
|
80
|
|
|
look: str - the look to use (default: 'neo') |
|
81
|
|
|
theme: str - the theme to use (default: 'neutral') |
|
82
|
|
|
""" |
|
83
|
|
|
self.orientation = orientation |
|
84
|
|
|
self.node_shape = node_shape |
|
85
|
|
|
self.layout = layout |
|
86
|
|
|
self.look = look |
|
87
|
|
|
self.theme = theme |
|
88
|
|
|
|
|
89
|
|
|
if not isinstance(orientation, DiagramOrientation): |
|
90
|
|
|
raise TypeError("orientation must be a valid Orientation enum") |
|
91
|
|
|
if not isinstance(node_shape, DiagramNodeShape): |
|
92
|
|
|
raise TypeError("node_shape must be a valid NodeShape enum") |
|
93
|
|
|
|
|
94
|
|
|
def _diagram_config(self, graph, title: str | None = None) -> str: |
|
95
|
|
|
return ( |
|
96
|
|
|
f"---\n" |
|
97
|
|
|
f"{_graph_title(graph, title)}" |
|
98
|
|
|
f"config:\n" |
|
99
|
|
|
f" layout: {self.layout}\n" |
|
100
|
|
|
f" look: {self.look}\n" |
|
101
|
|
|
f" theme: {self.theme}\n" |
|
102
|
|
|
f"---\n" |
|
103
|
|
|
) |
|
104
|
|
|
|
|
105
|
|
|
def build(self, graph: nx.Graph, title: str | None = None, with_edge_labels: bool = True) -> MermaidDiagram: |
|
106
|
|
|
""" |
|
107
|
|
|
Materialize a graph as a Mermaid flowchart. |
|
108
|
|
|
|
|
109
|
|
|
Args: |
|
110
|
|
|
graph: nx.Graph - The NetworkX graph to convert. |
|
111
|
|
|
title: str - The title of the graph (default: None). |
|
112
|
|
|
If None, the graph name will be used if available. |
|
113
|
|
|
Supplying and empty string will remove the title. |
|
114
|
|
|
with_edge_labels: bool - Whether to include edge labels (default: True). |
|
115
|
|
|
|
|
116
|
|
|
Returns: |
|
117
|
|
|
A string representation of the graph as a Mermaid graph. |
|
118
|
|
|
""" |
|
119
|
|
|
config = self._diagram_config(graph, title) |
|
120
|
|
|
|
|
121
|
|
|
bra, ket = self.node_shape.value |
|
122
|
|
|
|
|
123
|
|
|
minifier = AutoMapper() |
|
124
|
|
|
|
|
125
|
|
|
# Pre-calculate node IDs to avoid repeated function calls |
|
126
|
|
|
node_map = {u: minifier.get(u) for u in graph.nodes()} |
|
127
|
|
|
|
|
128
|
|
|
nodes = "\n".join( |
|
129
|
|
|
f"{node_map[u]}{bra}{d.get('label', u)}{ket}{_node_style(node_map[u], d)}" for u, d in |
|
130
|
|
|
graph.nodes.data()) |
|
131
|
|
|
|
|
132
|
|
|
_edges = ((node_map[u], node_map[v], d) for u, v, d in graph.edges.data()) |
|
|
|
|
|
|
133
|
|
|
edges = "\n".join(f"{u} -->{_edge_label(d) if with_edge_labels else ''} {v}" for u, v, d in _edges) |
|
134
|
|
|
|
|
135
|
|
|
return ( |
|
136
|
|
|
f"{config}" |
|
137
|
|
|
f"graph {self.orientation.value}\n" |
|
138
|
|
|
f"{nodes}\n" |
|
139
|
|
|
f"{edges}" |
|
140
|
|
|
) |
|
141
|
|
|
|
erivlis /
networkx-mermaid
Loading history...