Coverage for src / graphable / views / json.py: 100%

1from dataclasses import dataclass 

2from json import dumps 

3from logging import getLogger 

4from pathlib import Path 

5from typing import Any, Callable 

6 

7from ..graph import Graph 

8from ..graphable import Graphable 

9from ..registry import register_view 

10 

11logger = getLogger(__name__) 

12 

13 

14@dataclass 

15class JsonStylingConfig: 

16 """ 

17 Configuration for JSON graph serialization. 

18 

19 Attributes: 

20 node_data_fnc: Optional function to add extra data to each node's JSON object. 

21 reference_fnc: Function to generate the string identifier for each node. 

22 indent: JSON indentation level. 

23 """ 

24 

25 node_data_fnc: Callable[[Graphable[Any]], dict[str, Any]] | None = None 

26 reference_fnc: Callable[[Graphable[Any]], str] = lambda n: str(n.reference) 

27 indent: int | str | None = 2 

28 

29 

30def create_topology_json(graph: Graph, config: JsonStylingConfig | None = None) -> str: 

31 """ 

32 Generate a JSON representation of the graph. 

33 

34 Args: 

35 graph (Graph): The graph to convert. 

36 config (JsonStylingConfig | None): Serialization configuration. 

37 

38 Returns: 

39 str: A JSON string containing 'nodes' and 'edges'. 

40 """ 

41 logger.debug("Creating JSON representation of the graph.") 

42 config = config or JsonStylingConfig() 

43 

44 nodes = [] 

45 edges = [] 

46 

47 for node in graph.topological_order(): 

48 node_id = config.reference_fnc(node) 

49 node_entry = { 

50 "id": node_id, 

51 "reference": str(node.reference), 

52 "tags": list(node.tags), 

53 } 

54 if config.node_data_fnc: 

55 node_entry.update(config.node_data_fnc(node)) 

56 nodes.append(node_entry) 

57 

58 for dependent, _ in graph.internal_dependents(node): 

59 edges.append({"source": node_id, "target": config.reference_fnc(dependent)}) 

60 

61 data = {"nodes": nodes, "edges": edges} 

62 

63 return dumps(data, indent=config.indent) 

64 

65 

66@register_view(".json", creator_fnc=create_topology_json) 

67def export_topology_json( 

68 graph: Graph, 

69 output: Path, 

70 config: JsonStylingConfig | None = None, 

71) -> None: 

72 """ 

73 Export the graph to a JSON file. 

74 

75 Args: 

76 graph (Graph): The graph to export. 

77 output (Path): The output file path. 

78 config (JSONStylingConfig | None): Serialization configuration. 

79 """ 

80 logger.info(f"Exporting JSON to: {output}") 

81 with open(output, "w+") as f: 

82 f.write(create_topology_json(graph, config))