turbo_broccoli.custom.networkx

 1"""NetworkX graph serialization and deserialization."""
 2
 3from typing import Any, Callable, Tuple
 4
 5import networkx as nx
 6
 7from ..context import Context
 8from ..exceptions import DeserializationError, TypeNotSupported
 9
10
11def _graph_to_json(obj: nx.Graph, ctx: Context) -> dict:
12    return {
13        "__type__": "networkx.graph",
14        "__version__": 1,
15        "data": nx.adjacency_data(obj),
16    }
17
18
19def _json_to_graph(dct: dict, ctx: Context) -> nx.Graph:
20    decoders = {1: _json_to_graph_v1}
21    return decoders[dct["__version__"]](dct, ctx)
22
23
24def _json_to_graph_v1(dct: dict, ctx: Context) -> nx.Graph:
25    return nx.adjacency_graph(dct["data"])
26
27
28def from_json(dct: dict, ctx: Context) -> nx.Graph:
29    decoders = {
30        "networkx.graph": _json_to_graph,
31    }
32    try:
33        type_name = dct["__type__"]
34        return decoders[type_name](dct, ctx)
35    except KeyError as exc:
36        raise DeserializationError() from exc
37
38
39def to_json(obj: nx.Graph, ctx: Context) -> dict:
40    """
41    Serializes a graph into JSON by cases. The return dict has the following
42    structure
43
44    ```py
45    {
46        "__type__": "networkx.graph",
47        "__version__": 1,
48        "data": {...},
49    }
50    ```
51
52    where the `{...}` is produced by
53    [`networkx.adjacency_data`](https://networkx.org/documentation/stable/reference/readwrite/generated/networkx.readwrite.json_graph.adjacency_data.html#adjacency-data).
54    """
55    encoders: list[Tuple[type, Callable[[Any, Context], dict]]] = [
56        (nx.Graph, _graph_to_json),
57    ]
58    for t, f in encoders:
59        if isinstance(obj, t):
60            return f(obj, ctx)
61    raise TypeNotSupported()