Source code for semi_cr.core.lab.station.graphing.view

from typing import Any, Literal

import networkx as nx
import pandas as pd


def _short(value: Any) -> str:
    """Readable representation for graph attributes."""
    if isinstance(value, (str, int, float, bool, type(None))):
        return str(value)

    name = getattr(value, "name", None)
    if name is not None:
        return f"{type(value).__name__}({name})"

    return type(value).__name__


[docs] def graph_view( graph: nx.MultiDiGraph, what: Literal["nodes", "edges", "both"] = "both", include_attrs: bool = True, ) -> pd.DataFrame | tuple[pd.DataFrame, pd.DataFrame]: """ Return a nice tabular view of a NetworkX MultiDiGraph. Parameters ---------- graph: The graph to inspect. what: "nodes", "edges", or "both". include_attrs: Whether to include graph attributes. """ def nodes_df() -> pd.DataFrame: base_columns = [ "node", "in_degree", "out_degree", "degree", ] rows: list[dict[str, Any]] = [] for node, attrs in graph.nodes(data=True): row: dict[str, Any] = { "node": node, "in_degree": graph.in_degree(node), "out_degree": graph.out_degree(node), "degree": graph.degree(node), } if include_attrs: row.update( { key: _short(value) for key, value in attrs.items() } ) rows.append(row) if not rows: return pd.DataFrame(columns=base_columns) return ( pd.DataFrame(rows) .sort_values("node") .reset_index(drop=True) ) def edges_df() -> pd.DataFrame: base_columns = [ "source", "target", "key", ] rows: list[dict[str, Any]] = [] for source, target, key, attrs in graph.edges( keys=True, data=True, ): row: dict[str, Any] = { "source": source, "target": target, "key": key, } if include_attrs: row.update( { attr_name: _short(value) for attr_name, value in attrs.items() } ) rows.append(row) if not rows: return pd.DataFrame(columns=base_columns) return ( pd.DataFrame(rows) .sort_values(["source", "target", "key"]) .reset_index(drop=True) ) if what == "nodes": return nodes_df() if what == "edges": return edges_df() if what == "both": return nodes_df(), edges_df() raise ValueError( "what must be one of: 'nodes', 'edges', 'both'" )