Passed
Push — main ( 2bcd5c...05fcff )
by Eran
01:27
created

networkx_mermaid.title()   A

Complexity

Conditions 2

Size

Total Lines 2
Code Lines 2

Duplication

Lines 0
Ratio 0 %

Importance

Changes 0
Metric Value
eloc 2
dl 0
loc 2
rs 10
c 0
b 0
f 0
cc 2
nop 1
1
"""
2
The `networkx_mermaid.py` script defines a function to materialize a NetworkX graph as a Mermaid flowchart.
3
 It includes enumerations for graph orientation, node shapes, and output formats.
4
  The script also provides an example of generating a Mermaid diagram and serving it via a temporary HTTP server.
5
"""
6
7
from enum import Enum
8
from tempfile import TemporaryDirectory
9
from typing import TypeVar
10
11
import networkx as nx
12
13
MermaidGraph = TypeVar('MermaidGraph', bound='str')
14
15
class Orientation(Enum):
16
    """
17
    Orientation of a Mermaid graph.
18
    """
19
    TOP_DOWN = "TD"
20
    BOTTOM_UP = "BT"
21
    LEFT_RIGHT = "LR"
22
    RIGHT_LEFT = "RL"
23
24
25
class NodeShape(Enum):
26
    """
27
    Shapes of a Mermaid graph node.
28
    """
29
    DEFAULT = ('(', ')')
30
    RECTANGLE = ('[', ']')
31
    ROUND_RECTANGLE = ('([', '])')
32
    SUBROUTINE = ('[[', ']]')
33
    DATABASE = ('[(', ')]')
34
    CIRCLE = ('((', '))')
35
    DOUBLE_CIRCLE = ('(((', ')))')
36
    FLAG = ('>', ']')
37
    DIAMOND = ('{', '}')
38
    HEXAGON = ('{{', '}}')
39
    PARALLELOGRAM = ('[/', '/]')
40
    PARALLELOGRAM_ALT = ('[\\', '\\]')
41
    TRAPEZOID = ('[/', '\\]')
42
    TRAPEZOID_ALT = ('[\\', '/]')
43
44
45
HTML_TEMPLATE = """<!doctype html>
46
<html lang="en">
47
  <head>
48
    <link rel="icon" type="image/x-icon" href="https://mermaid.js.org/favicon.ico">
49
    <meta charset="utf-8">
50
    <title>Mermaid Diagram</title>
51
    <style>
52
    pre.mermaid {
53
      font-family: "Fira Mono", "Roboto Mono", "Source Code Pro", monospace;
54
    }
55
    </style>
56
  </head>
57
  <body>
58
    <pre class="mermaid">
59
    </pre>
60
    <script type="module">
61
      import mermaid from 'https://cdn.jsdelivr.net/npm/mermaid@11/dist/mermaid.esm.min.mjs';
62
      let config = { startOnLoad: true, flowchart: { useMaxWidth: false, htmlLabels: true } };
63
      mermaid.initialize(config);
64
    </script>
65
  </body>
66
</html>
67
"""
68
69
70
def edge_label(data: dict) -> str:
71
    label = data.get('label')
72
    return f"|{label}|" if label else ""
73
74
75
def contrast_color(color: str) -> str:
76
    """
77
    Return black or white by choosing the best contrast to input color
78
79
    Args:
80
        color: str - hex color code
81
82
    Returns:
83
        color: str - hex color code
84
    """
85
    r, g, b = int(color[1:3], 16), int(color[3:5], 16), int(color[5:7], 16)
86
    return '#000000' if (r * 0.299 + g * 0.587 + b * 0.114) > 186 else '#ffffff'
87
88
89
def style(node_id: str, data: dict) -> str:
90
    color = data.get('color')
91
    return f"\nstyle {node_id} fill:{color}, color:{contrast_color(color)}" if color else ""
92
93
94
def title(graph: nx.Graph) -> str:
95
    return f"title: {graph.name}\n" if graph.name else ""
96
97
98
def mermaid(graph: nx.Graph,
99
            orientation: Orientation = Orientation.LEFT_RIGHT,
100
            node_shape: NodeShape = NodeShape.DEFAULT) -> MermaidGraph:
101
    """
102
    Materialize a graph as a Mermaid flowchart.
103
104
    Parameters
105
    ----------
106
    graph : nx.Graph
107
        A NetworkX graph.
108
    orientation : Orientation, optional
109
        The orientation of the graph, by default Orientation.LEFT_RIGHT.
110
    node_shape : NodeShape, optional
111
        The shape of the nodes, by default NodeShape.DEFAULT.
112
113
    Returns
114
    -------
115
    str
116
        A string representation of the graph as a Mermaid graph.
117
    """
118
    layout = 'dagre'
119
    look = 'neo'
120
    theme = 'neutral'
121
    config = (f"---\n"
122
              f"{title(graph)}"
123
              f"config:\n"
124
              f"  layout: {layout}\n"
125
              f"  look: {look}\n"
126
              f"  theme: {theme}\n"
127
              f"---")
128
129
    bra, ket = node_shape.value
130
    nodes = '\n'.join(
131
        f"{u}{bra}{v.get('label', u)}{ket}{style(u, v)}" for u, v in
132
        graph.nodes.data()
133
    )
134
    edges = '\n'.join(
135
        f"{u} -->{edge_label(d)} {v}" for u, v, d in graph.edges.data()
136
    )
137
138
    return (f"{config}\n"
139
            f"graph {orientation.value}\n"
140
            f"{nodes}\n"
141
            f"{edges}")
142
143
144
def markdown(diagram: str, title: str | None = None) -> str:
145
    output = f"```mermaid\n{diagram}\n```"
146
    if title:
147
        output = f"## {title}\n\n{output}"
148
    return output
149
150
151
def html(diagram: str, title: str | None = None) -> str:
152
    output = HTML_TEMPLATE.replace('<pre class="mermaid">', f'<pre class="mermaid">\n{diagram}\n')
153
    if title:
154
        output = output.replace('<title>Mermaid Diagram</title>', f'<title>{title}</title>')
155
    return output
156
157
158
if __name__ == '__main__':
159
    # colors = ['#FF0000', '#00FF00', '#0000FF', '#FFFF00', '#00FFFF', '#FF00FF']
160
    pastel_colors = ['#FFCCCC', '#CCFFCC', '#CCCCFF', '#FFFFCC', '#CCFFFF', '#FFCCFF']
161
    graphs: list[nx.Graph] = [nx.tetrahedral_graph(), nx.dodecahedral_graph()]
162
163
    for i, g in enumerate(graphs):
164
        nx.set_node_attributes(g, {n: {'color': pastel_colors[i]} for n in g.nodes})
165
166
    graph: nx.Graph = nx.disjoint_union_all(graphs)
167
168
    graph.name = ' + '.join(g.name for g in graphs)
169
170
    mermaid_diagram = mermaid(graph,
171
                              orientation=Orientation.LEFT_RIGHT,
172
                              node_shape=NodeShape.ROUND_RECTANGLE)
173
174
    with TemporaryDirectory() as temp_dir:
175
        with open(f"{temp_dir}/index.html", 'w') as f:
176
            rendered = html(mermaid_diagram, title=graph.name)
177
            f.write(rendered)
178
179
        import http.server
180
        import socketserver
181
182
        PORT = 8073
183
184
185
        class Handler(http.server.SimpleHTTPRequestHandler):
186
            def __init__(self, *args, **kwargs):
187
                super().__init__(*args, directory=temp_dir, **kwargs)
0 ignored issues
show
introduced by
The variable temp_dir does not seem to be defined in case __name__ == '__main__' on line 158 is False. Are you sure this can never be the case?
Loading history...
188
189
190
        with socketserver.TCPServer(("", PORT), Handler) as httpd:
191
            print("serving at port", PORT)
192
            httpd.serve_forever()
193