unifi-network-maps 1.4.1__py3-none-any.whl → 1.4.3__py3-none-any.whl

unifi_network_maps/render/mermaid.py

@@ -2,6 +2,7 @@
 
 from __future__ import annotations
 
+import json
 from collections.abc import Iterable
 
 from ..model.topology import Edge
@@ -9,7 +10,9 @@ from .mermaid_theme import DEFAULT_THEME, MermaidTheme, class_defs
 
 
 def _escape(label: str) -> str:
-    return label.replace('"', '\\"')
+    normalized = label.replace("\r\n", "\n").replace("\r", "\n")
+    escaped = normalized.replace("\\", "\\\\").replace("\n", "\\n")
+    return escaped.replace('"', '\\"')
 
 
 def _slugify(value: str) -> str:
@@ -54,39 +57,45 @@ def _node_ref(name: str, node_id: str) -> str:
     return f'{node_id}["{_escape(name)}"]'
 
 
-def render_mermaid(
-    edges: Iterable[Edge],
-    direction: str = "LR",
+def _group_nodes(groups: dict[str, list[str]] | None) -> list[str]:
+    if not groups:
+        return []
+    nodes: list[str] = []
+    for members in groups.values():
+        nodes.extend(members)
+    return nodes
+
+
+def _render_group_sections(
+    lines: list[str],
+    groups: dict[str, list[str]],
     *,
-    groups: dict[str, list[str]] | None = None,
-    group_order: list[str] | None = None,
-    node_types: dict[str, str] | None = None,
-    theme: MermaidTheme = DEFAULT_THEME,
-) -> str:
-    edge_list = list(edges)
-    group_nodes: list[str] = []
-    if groups:
-        for members in groups.values():
-            group_nodes.extend(members)
-    id_map = _build_id_map(edge_list, group_nodes)
-    lines = [f"graph {direction}"]
+    group_order: list[str] | None,
+    id_map: dict[str, str],
+) -> None:
+    ordered = group_order or list(groups.keys())
+    for group_name in ordered:
+        members = groups.get(group_name, [])
+        if not members:
+            continue
+        group_id = _slugify(f"group_{group_name}")
+        label = group_name.replace("_", " ").title()
+        lines.append(f'  subgraph {group_id}["{_escape(label)}"];')
+        for member in members:
+            lines.append(f"    {_node_ref(member, id_map[member])};")
+        lines.append("  end")
+
+
+def _render_edge_lines(
+    lines: list[str],
+    edges: list[Edge],
+    *,
+    id_map: dict[str, str],
+    use_node_labels: bool,
+) -> tuple[list[int], list[int]]:
     poe_links: list[int] = []
     wireless_links: list[int] = []
-    link_index = 0
-    if groups:
-        ordered = group_order or list(groups.keys())
-        for group_name in ordered:
-            members = groups.get(group_name, [])
-            if not members:
-                continue
-            group_id = _slugify(f"group_{group_name}")
-            label = group_name.replace("_", " ").title()
-            lines.append(f'  subgraph {group_id}["{_escape(label)}"];')
-            for member in members:
-                lines.append(f"    {_node_ref(member, id_map[member])};")
-            lines.append("  end")
-    use_node_labels = not groups
-    for edge in edge_list:
+    for index, edge in enumerate(edges):
         if use_node_labels:
             left = _node_ref(edge.left, id_map[edge.left])
             right = _node_ref(edge.right, id_map[edge.right])
@@ -99,25 +108,41 @@ def render_mermaid(
         else:
             lines.append(f"  {left} --- {right};")
         if edge.poe:
-            poe_links.append(link_index)
+            poe_links.append(index)
         if edge.wireless:
-            wireless_links.append(link_index)
-        link_index += 1
-    if node_types:
-        class_map = {
-            "gateway": "node_gateway",
-            "switch": "node_switch",
-            "ap": "node_ap",
-            "client": "node_client",
-            "other": "node_other",
-        }
-        if node_types:
-            for name, node_type in node_types.items():
-                class_name = class_map.get(node_type, "node_other")
-                node_id = id_map.get(name)
-                if node_id:
-                    lines.append(f"  class {node_id} {class_name};")
-        lines.extend(class_defs(theme))
+            wireless_links.append(index)
+    return poe_links, wireless_links
+
+
+def _render_node_classes(
+    lines: list[str],
+    *,
+    node_types: dict[str, str],
+    id_map: dict[str, str],
+    theme: MermaidTheme,
+) -> None:
+    class_map = {
+        "gateway": "node_gateway",
+        "switch": "node_switch",
+        "ap": "node_ap",
+        "client": "node_client",
+        "other": "node_other",
+    }
+    for name, node_type in node_types.items():
+        class_name = class_map.get(node_type, "node_other")
+        node_id = id_map.get(name)
+        if node_id:
+            lines.append(f"  class {node_id} {class_name};")
+    lines.extend(class_defs(theme))
+
+
+def _render_link_styles(
+    lines: list[str],
+    *,
+    poe_links: list[int],
+    wireless_links: list[int],
+    theme: MermaidTheme,
+) -> None:
     for index in poe_links:
         lines.append(
             "  linkStyle "
@@ -126,6 +151,37 @@ def render_mermaid(
         )
     for index in wireless_links:
         lines.append(f"  linkStyle {index} stroke-dasharray: 5 4;")
+
+
+def render_mermaid(
+    edges: Iterable[Edge],
+    direction: str = "LR",
+    *,
+    groups: dict[str, list[str]] | None = None,
+    group_order: list[str] | None = None,
+    node_types: dict[str, str] | None = None,
+    theme: MermaidTheme = DEFAULT_THEME,
+) -> str:
+    edge_list = list(edges)
+    id_map = _build_id_map(edge_list, _group_nodes(groups))
+    theme_vars: dict[str, object] = {}
+    if theme.edge_label_border:
+        theme_vars["edgeLabelBorderColor"] = theme.edge_label_border
+    if theme.edge_label_border_width:
+        theme_vars["edgeLabelBorderWidth"] = theme.edge_label_border_width
+    lines = []
+    if theme_vars:
+        lines.append(f'%%{{init: {{"themeVariables": {json.dumps(theme_vars)}}}}}%%')
+    lines.append(f"graph {direction}")
+    if groups:
+        _render_group_sections(lines, groups, group_order=group_order, id_map=id_map)
+    use_node_labels = not groups
+    poe_links, wireless_links = _render_edge_lines(
+        lines, edge_list, id_map=id_map, use_node_labels=use_node_labels
+    )
+    if node_types:
+        _render_node_classes(lines, node_types=node_types, id_map=id_map, theme=theme)
+    _render_link_styles(lines, poe_links=poe_links, wireless_links=wireless_links, theme=theme)
     return "\n".join(lines) + "\n"
 
 

unifi_network_maps/render/mermaid_theme.py

@@ -18,6 +18,9 @@ class MermaidTheme:
     standard_link: str
     standard_link_width: int
     standard_link_arrow: str
+    node_text: str | None = None
+    edge_label_border: str | None = None
+    edge_label_border_width: int | None = None
 
 
 DEFAULT_THEME = MermaidTheme(
@@ -32,15 +35,22 @@ DEFAULT_THEME = MermaidTheme(
     standard_link="#2ecc71",
     standard_link_width=2,
     standard_link_arrow="none",
+    node_text=None,
+    edge_label_border=None,
+    edge_label_border_width=None,
 )
 
 
 def class_defs(theme: MermaidTheme = DEFAULT_THEME) -> list[str]:
+    def node_def(name: str, fill: str, stroke: str) -> str:
+        color = f",color:{theme.node_text}" if theme.node_text else ""
+        return f"  classDef {name} fill:{fill},stroke:{stroke},stroke-width:1px{color};"
+
     return [
-        f"  classDef node_gateway fill:{theme.node_gateway[0]},stroke:{theme.node_gateway[1]},stroke-width:1px;",
-        f"  classDef node_switch fill:{theme.node_switch[0]},stroke:{theme.node_switch[1]},stroke-width:1px;",
-        f"  classDef node_ap fill:{theme.node_ap[0]},stroke:{theme.node_ap[1]},stroke-width:1px;",
-        f"  classDef node_client fill:{theme.node_client[0]},stroke:{theme.node_client[1]},stroke-width:1px;",
-        f"  classDef node_other fill:{theme.node_other[0]},stroke:{theme.node_other[1]},stroke-width:1px;",
+        node_def("node_gateway", theme.node_gateway[0], theme.node_gateway[1]),
+        node_def("node_switch", theme.node_switch[0], theme.node_switch[1]),
+        node_def("node_ap", theme.node_ap[0], theme.node_ap[1]),
+        node_def("node_client", theme.node_client[0], theme.node_client[1]),
+        node_def("node_other", theme.node_other[0], theme.node_other[1]),
         "  classDef node_legend font-size:10px;",
     ]