unifi-network-maps 1.2.1__py3-none-any.whl → 1.3.0__py3-none-any.whl

{unifi_mermaid → unifi_network_maps/model}/topology.py

@@ -6,10 +6,11 @@ import logging
 from collections import deque
 from collections.abc import Iterable
 from dataclasses import dataclass, field
-from typing import Protocol
+from typing import Protocol, TypeAlias
 
-from .labels import compose_port_label
+from .labels import compose_port_label, order_edge_names
 from .lldp import LLDPEntry, coerce_lldp, local_port_label
+from .ports import extract_port_number
 
 logger = logging.getLogger(__name__)
 
@@ -75,6 +76,10 @@ class PortInfo:
     poe_power: float | None
 
 
+PortMap: TypeAlias = dict[tuple[str, str], str]
+PoeMap: TypeAlias = dict[tuple[str, str], bool]
+
+
 def _get_attr(obj: object, name: str) -> object | None:
     return getattr(obj, name, None)
 
@@ -114,50 +119,77 @@ def _as_int(value: object | None) -> int | None:
     return None
 
 
+def _resolve_port_idx_from_lldp(lldp_entry: LLDPEntry, port_table: list[PortInfo]) -> int | None:
+    if lldp_entry.local_port_idx is not None:
+        return lldp_entry.local_port_idx
+    candidates = []
+    if lldp_entry.local_port_name:
+        candidates.append(lldp_entry.local_port_name)
+    if lldp_entry.port_id:
+        candidates.append(lldp_entry.port_id)
+    for candidate in candidates:
+        normalized = candidate.strip().lower()
+        for port in port_table:
+            if port.ifname and port.ifname.strip().lower() == normalized:
+                return port.port_idx
+            if port.name and port.name.strip().lower() == normalized:
+                return port.port_idx
+    for candidate in candidates:
+        number = extract_port_number(candidate)
+        if number is None:
+            continue
+        for port in port_table:
+            if port.port_idx == number:
+                return port.port_idx
+    return None
+
+
+def _port_info_from_entry(port_entry: object) -> PortInfo:
+    if isinstance(port_entry, dict):
+        port_idx = port_entry.get("port_idx") or port_entry.get("portIdx")
+        name = port_entry.get("name")
+        ifname = port_entry.get("ifname")
+        port_poe = _as_bool(port_entry.get("port_poe"))
+        poe_enable = _as_bool(port_entry.get("poe_enable"))
+        poe_good = _as_bool(port_entry.get("poe_good"))
+        poe_power = _as_float(port_entry.get("poe_power"))
+    else:
+        port_idx = _get_attr(port_entry, "port_idx") or _get_attr(port_entry, "portIdx")
+        name = _get_attr(port_entry, "name")
+        ifname = _get_attr(port_entry, "ifname")
+        port_poe = _as_bool(_get_attr(port_entry, "port_poe"))
+        poe_enable = _as_bool(_get_attr(port_entry, "poe_enable"))
+        poe_good = _as_bool(_get_attr(port_entry, "poe_good"))
+        poe_power = _as_float(_get_attr(port_entry, "poe_power"))
+    return PortInfo(
+        port_idx=_as_int(port_idx),
+        name=str(name) if isinstance(name, str) and name.strip() else None,
+        ifname=str(ifname) if isinstance(ifname, str) and ifname.strip() else None,
+        port_poe=port_poe,
+        poe_enable=poe_enable,
+        poe_good=poe_good,
+        poe_power=poe_power,
+    )
+
+
 def _coerce_port_table(device: DeviceLike) -> list[PortInfo]:
     port_table = _get_attr(device, "port_table") or []
-    result: list[PortInfo] = []
-    for entry in port_table:
-        if isinstance(entry, dict):
-            port_idx = entry.get("port_idx") or entry.get("portIdx")
-            name = entry.get("name")
-            ifname = entry.get("ifname")
-            port_poe = _as_bool(entry.get("port_poe"))
-            poe_enable = _as_bool(entry.get("poe_enable"))
-            poe_good = _as_bool(entry.get("poe_good"))
-            poe_power = _as_float(entry.get("poe_power"))
-        else:
-            port_idx = _get_attr(entry, "port_idx") or _get_attr(entry, "portIdx")
-            name = _get_attr(entry, "name")
-            ifname = _get_attr(entry, "ifname")
-            port_poe = _as_bool(_get_attr(entry, "port_poe"))
-            poe_enable = _as_bool(_get_attr(entry, "poe_enable"))
-            poe_good = _as_bool(_get_attr(entry, "poe_good"))
-            poe_power = _as_float(_get_attr(entry, "poe_power"))
-        result.append(
-            PortInfo(
-                port_idx=_as_int(port_idx),
-                name=str(name) if isinstance(name, str) and name.strip() else None,
-                ifname=str(ifname) if isinstance(ifname, str) and ifname.strip() else None,
-                port_poe=port_poe,
-                poe_enable=poe_enable,
-                poe_good=poe_good,
-                poe_power=poe_power,
-            )
-        )
-    return result
+    return [_port_info_from_entry(port_entry) for port_entry in port_table]
 
 
 def _poe_ports_from_device(device: DeviceLike) -> dict[int, bool]:
     port_table = _coerce_port_table(device)
     poe_ports: dict[int, bool] = {}
-    for entry in port_table:
-        if entry.port_idx is None:
+    for port_entry in port_table:
+        if port_entry.port_idx is None:
             continue
         active = (
-            entry.poe_enable or entry.port_poe or entry.poe_good or _as_float(entry.poe_power) > 0.0
+            port_entry.poe_enable
+            or port_entry.port_poe
+            or port_entry.poe_good
+            or _as_float(port_entry.poe_power) > 0.0
         )
-        poe_ports[int(entry.port_idx)] = active
+        poe_ports[int(port_entry.port_idx)] = active
     return poe_ports
 
 
@@ -224,7 +256,7 @@ def coerce_device(device: DeviceLike) -> Device:
         else:
             raise ValueError(f"Device {name} missing LLDP info")
 
-    coerced_lldp = [coerce_lldp(entry) for entry in lldp_info]
+    coerced_lldp = [coerce_lldp(lldp_entry) for lldp_entry in lldp_info]
     port_table = _coerce_port_table(device)
     poe_ports = _poe_ports_from_device(device)
 
@@ -431,15 +463,59 @@ def build_edges(
     ordered_devices = sorted(devices, key=lambda item: (item.name.lower(), item.mac.lower()))
     index = build_device_index(ordered_devices)
     device_by_name = {device.name: device for device in ordered_devices}
-    pairs: list[tuple[str, str]] = []
+    raw_links: list[tuple[str, str]] = []
     seen: set[frozenset[str]] = set()
-    port_map: dict[tuple[str, str], str] = {}
-    poe_map: dict[tuple[str, str], bool] = {}
-    devices_with_lldp_edges: set[str] = set()
+    port_map: PortMap = {}
+    poe_map: PoeMap = {}
+
+    devices_with_lldp_edges = _collect_lldp_links(
+        ordered_devices,
+        index,
+        port_map,
+        poe_map,
+        raw_links,
+        seen,
+        only_unifi=only_unifi,
+    )
+    _collect_uplink_links(
+        ordered_devices,
+        devices_with_lldp_edges,
+        index,
+        device_by_name,
+        port_map,
+        poe_map,
+        raw_links,
+        seen,
+        include_ports=include_ports,
+        only_unifi=only_unifi,
+    )
+    edges = _build_ordered_edges(
+        raw_links,
+        port_map,
+        poe_map,
+        device_by_name,
+        include_ports=include_ports,
+    )
+
+    poe_edges = sum(1 for edge in edges if edge.poe)
+    logger.info("Built %d unique edges (%d PoE)", len(edges), poe_edges)
+    return edges
 
-    for device in ordered_devices:
+
+def _collect_lldp_links(
+    devices: list[Device],
+    index: dict[str, str],
+    port_map: PortMap,
+    poe_map: PoeMap,
+    raw_links: list[tuple[str, str]],
+    seen: set[frozenset[str]],
+    *,
+    only_unifi: bool,
+) -> set[str]:
+    devices_with_lldp_edges: set[str] = set()
+    for device in devices:
         poe_ports = device.poe_ports
-        for entry in sorted(
+        for lldp_entry in sorted(
             device.lldp_info,
             key=lambda item: (
                 _normalize_mac(item.chassis_id),
@@ -447,78 +523,117 @@ def build_edges(
                 str(item.port_desc or ""),
             ),
         ):
-            neighbor_mac = _normalize_mac(entry.chassis_id)
-            neighbor_name = index.get(neighbor_mac)
-            if neighbor_name is None:
+            peer_mac = _normalize_mac(lldp_entry.chassis_id)
+            peer_name = index.get(peer_mac)
+            if peer_name is None:
                 if only_unifi:
                     continue
-                neighbor_name = entry.chassis_id
-
-            label = local_port_label(entry)
+                peer_name = lldp_entry.chassis_id
+
+            resolved_port_idx = _resolve_port_idx_from_lldp(lldp_entry, device.port_table)
+            entry_for_label = (
+                LLDPEntry(
+                    chassis_id=lldp_entry.chassis_id,
+                    port_id=lldp_entry.port_id,
+                    port_desc=lldp_entry.port_desc,
+                    local_port_name=lldp_entry.local_port_name,
+                    local_port_idx=resolved_port_idx,
+                )
+                if resolved_port_idx is not None
+                else lldp_entry
+            )
+            label = local_port_label(entry_for_label)
             if label:
-                port_map[(device.name, neighbor_name)] = label
-            if entry.local_port_idx is not None and entry.local_port_idx in poe_ports:
-                poe_map[(device.name, neighbor_name)] = poe_ports[entry.local_port_idx]
+                port_map[(device.name, peer_name)] = label
+            if resolved_port_idx is not None and resolved_port_idx in poe_ports:
+                poe_map[(device.name, peer_name)] = poe_ports[resolved_port_idx]
 
-            key = frozenset({device.name, neighbor_name})
+            key = frozenset({device.name, peer_name})
             if key in seen:
                 continue
 
-            pairs.append((device.name, neighbor_name))
+            raw_links.append((device.name, peer_name))
             seen.add(key)
             devices_with_lldp_edges.add(device.name)
-
-    for device in ordered_devices:
+    return devices_with_lldp_edges
+
+
+def _collect_uplink_links(
+    devices: list[Device],
+    devices_with_lldp_edges: set[str],
+    index: dict[str, str],
+    device_by_name: dict[str, Device],
+    port_map: PortMap,
+    poe_map: PoeMap,
+    raw_links: list[tuple[str, str]],
+    seen: set[frozenset[str]],
+    *,
+    include_ports: bool,
+    only_unifi: bool,
+) -> None:
+    for device in devices:
         if device.name in devices_with_lldp_edges:
             continue
         uplink = device.uplink or device.last_uplink
-        uplink_name = None
+        upstream_name = None
         if uplink and uplink.mac:
-            uplink_name = index.get(_normalize_mac(uplink.mac))
-        if not uplink_name and uplink and uplink.name:
-            uplink_name = uplink.name
-        if not uplink_name and not only_unifi and uplink and uplink.mac:
-            uplink_name = uplink.mac
-        if not uplink_name:
+            upstream_name = index.get(_normalize_mac(uplink.mac))
+        if not upstream_name and uplink and uplink.name:
+            upstream_name = uplink.name
+        if not upstream_name and not only_unifi and uplink and uplink.mac:
+            upstream_name = uplink.mac
+        if not upstream_name:
             continue
-        if only_unifi and uplink_name not in device_by_name:
+        if only_unifi and upstream_name not in device_by_name:
             continue
-        key = frozenset({device.name, uplink_name})
+        key = frozenset({device.name, upstream_name})
         if key in seen:
             continue
         poe = False
         if uplink and uplink.port is not None:
             if include_ports:
-                port_map[(uplink_name, device.name)] = f"Port {uplink.port}"
-            uplink_device = device_by_name.get(uplink_name)
+                port_map[(upstream_name, device.name)] = f"Port {uplink.port}"
+            uplink_device = device_by_name.get(upstream_name)
             if uplink_device and uplink.port in uplink_device.poe_ports:
                 poe = uplink_device.poe_ports[uplink.port]
-        pairs.append((uplink_name, device.name))
+        raw_links.append((upstream_name, device.name))
         seen.add(key)
         if poe:
-            poe_map[(uplink_name, device.name)] = poe
+            poe_map[(upstream_name, device.name)] = poe
+
+
+def _build_ordered_edges(
+    raw_links: list[tuple[str, str]],
+    port_map: PortMap,
+    poe_map: PoeMap,
+    device_by_name: dict[str, Device],
+    *,
+    include_ports: bool,
+) -> list[Edge]:
+    type_rank = {"gateway": 0, "switch": 1, "ap": 2, "other": 3}
+
+    def _rank_for_name(name: str) -> int:
+        device = device_by_name.get(name)
+        if not device:
+            return 3
+        return type_rank.get(classify_device_type(device), 3)
 
     edges: list[Edge] = []
-    for left, right in pairs:
+    for source_name, target_name in raw_links:
+        left_name = source_name
+        right_name = target_name
         if include_ports:
-            left_label = port_map.get((left, right))
-            right_label = port_map.get((right, left))
-            if left_label is None and right_label is not None:
-                left, right = right, left
-            elif left_label and right_label:
-                left_device = device_by_name.get(left)
-                right_device = device_by_name.get(right)
-                if left_device and right_device:
-                    type_rank = {"gateway": 0, "switch": 1, "ap": 2, "other": 3}
-                    left_rank = type_rank.get(classify_device_type(left_device), 3)
-                    right_rank = type_rank.get(classify_device_type(right_device), 3)
-                    if (left_rank, left.lower()) > (right_rank, right.lower()):
-                        left, right = right, left
-        poe = poe_map.get((left, right), False) or poe_map.get((right, left), False)
-        label = compose_port_label(left, right, port_map) if include_ports else None
-        edges.append(Edge(left=left, right=right, label=label, poe=poe))
-
-    logger.info("Built %d unique edges", len(edges))
+            left_name, right_name = order_edge_names(
+                left_name,
+                right_name,
+                port_map,
+                _rank_for_name,
+            )
+        poe = poe_map.get((left_name, right_name), False) or poe_map.get(
+            (right_name, left_name), False
+        )
+        label = compose_port_label(left_name, right_name, port_map) if include_ports else None
+        edges.append(Edge(left=left_name, right=right_name, label=label, poe=poe))
     return edges
 
 
@@ -535,6 +650,18 @@ def build_topology(
     only_unifi: bool,
     gateways: list[str],
 ) -> TopologyResult:
-    raw_edges = build_edges(devices, include_ports=include_ports, only_unifi=only_unifi)
+    normalized_devices = list(devices)
+    lldp_entries = sum(len(device.lldp_info) for device in normalized_devices)
+    logger.info(
+        "Normalized %d devices (%d LLDP entries)",
+        len(normalized_devices),
+        lldp_entries,
+    )
+    raw_edges = build_edges(normalized_devices, include_ports=include_ports, only_unifi=only_unifi)
     tree_edges = build_tree_edges_by_topology(raw_edges, gateways)
+    logger.info(
+        "Built %d hierarchy edges (gateways=%d)",
+        len(tree_edges),
+        len(gateways),
+    )
     return TopologyResult(raw_edges=raw_edges, tree_edges=tree_edges)

unifi_network_maps/render/__init__.py

@@ -0,0 +1 @@
+"""Package module."""

{unifi_mermaid → unifi_network_maps/render}/mermaid.py

@@ -4,7 +4,8 @@ from __future__ import annotations
 
 from collections.abc import Iterable
 
-from .topology import Edge
+from ..model.topology import Edge
+from .mermaid_theme import DEFAULT_THEME, MermaidTheme, class_defs
 
 
 def _escape(label: str) -> str:
@@ -60,6 +61,7 @@ def render_mermaid(
     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] = []
@@ -112,17 +114,17 @@ def render_mermaid(
                 node_id = id_map.get(name)
                 if node_id:
                     lines.append(f"  class {node_id} {class_name};")
-        lines.append("  classDef node_gateway fill:#ffe3b3,stroke:#d98300,stroke-width:1px;")
-        lines.append("  classDef node_switch fill:#d6ecff,stroke:#3a7bd5,stroke-width:1px;")
-        lines.append("  classDef node_ap fill:#d7f5e7,stroke:#27ae60,stroke-width:1px;")
-        lines.append("  classDef node_client fill:#f2e5ff,stroke:#7f3fbf,stroke-width:1px;")
-        lines.append("  classDef node_other fill:#eeeeee,stroke:#8f8f8f,stroke-width:1px;")
+        lines.extend(class_defs(theme))
     for index in poe_links:
-        lines.append(f"  linkStyle {index} stroke:#1e88e5,stroke-width:2px;")
+        lines.append(
+            "  linkStyle "
+            f"{index} stroke:{theme.poe_link},stroke-width:{theme.poe_link_width}px,"
+            f"arrowhead:{theme.poe_link_arrow};"
+        )
     return "\n".join(lines) + "\n"
 
 
-def render_legend() -> str:
+def render_legend(theme: MermaidTheme = DEFAULT_THEME) -> str:
     lines = [
         "graph TB",
         '  subgraph legend["Legend"];',
@@ -149,13 +151,16 @@ def render_legend() -> str:
         "  class legend_poe_b node_legend;",
         "  class legend_no_poe_a node_legend;",
         "  class legend_no_poe_b node_legend;",
-        "  classDef node_gateway fill:#ffe3b3,stroke:#d98300,stroke-width:1px;",
-        "  classDef node_switch fill:#d6ecff,stroke:#3a7bd5,stroke-width:1px;",
-        "  classDef node_ap fill:#d7f5e7,stroke:#27ae60,stroke-width:1px;",
-        "  classDef node_client fill:#f2e5ff,stroke:#7f3fbf,stroke-width:1px;",
-        "  classDef node_other fill:#eeeeee,stroke:#8f8f8f,stroke-width:1px;",
-        "  classDef node_legend font-size:10px;",
     ]
-    lines.append("  linkStyle 0 stroke:#1e88e5,stroke-width:2px,arrowhead:none;")
-    lines.append("  linkStyle 1 stroke:#2ecc71,stroke-width:2px,arrowhead:none;")
+    lines.extend(class_defs(theme))
+    lines.append(
+        "  linkStyle 0 "
+        f"stroke:{theme.poe_link},stroke-width:{theme.poe_link_width}px,"
+        f"arrowhead:{theme.poe_link_arrow};"
+    )
+    lines.append(
+        "  linkStyle 1 "
+        f"stroke:{theme.standard_link},stroke-width:{theme.standard_link_width}px,"
+        f"arrowhead:{theme.standard_link_arrow};"
+    )
     return "\n".join(lines) + "\n"

unifi_network_maps/render/mermaid_theme.py

@@ -0,0 +1,46 @@
+"""Mermaid theming helpers."""
+
+from __future__ import annotations
+
+from dataclasses import dataclass
+
+
+@dataclass(frozen=True)
+class MermaidTheme:
+    node_gateway: tuple[str, str]
+    node_switch: tuple[str, str]
+    node_ap: tuple[str, str]
+    node_client: tuple[str, str]
+    node_other: tuple[str, str]
+    poe_link: str
+    poe_link_width: int
+    poe_link_arrow: str
+    standard_link: str
+    standard_link_width: int
+    standard_link_arrow: str
+
+
+DEFAULT_THEME = MermaidTheme(
+    node_gateway=("#ffe3b3", "#d98300"),
+    node_switch=("#d6ecff", "#3a7bd5"),
+    node_ap=("#d7f5e7", "#27ae60"),
+    node_client=("#f2e5ff", "#7f3fbf"),
+    node_other=("#eeeeee", "#8f8f8f"),
+    poe_link="#1e88e5",
+    poe_link_width=2,
+    poe_link_arrow="none",
+    standard_link="#2ecc71",
+    standard_link_width=2,
+    standard_link_arrow="none",
+)
+
+
+def class_defs(theme: MermaidTheme = DEFAULT_THEME) -> list[str]:
+    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;",
+        "  classDef node_legend font-size:10px;",
+    ]