Coverage for src / graphable / views / toml.py: 92%

1from dataclasses import dataclass 

2from logging import getLogger 

3from pathlib import Path 

4from typing import Any, Callable 

5 

6from ..graph import Graph 

7from ..graphable import Graphable 

8from ..registry import register_view 

9 

10logger = getLogger(__name__) 

11 

12 

13@dataclass 

14class TomlStylingConfig: 

15 """ 

16 Configuration for TOML graph serialization. 

17 

18 Attributes: 

19 node_data_fnc: Optional function to add extra data to each node's TOML object. 

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

21 """ 

22 

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

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

25 

26 

27def create_topology_toml(graph: Graph, config: TomlStylingConfig | None = None) -> str: 

28 """ 

29 Generate a TOML representation of the graph. 

30 Requires 'tomli-w' to be installed. 

31 

32 Args: 

33 graph (Graph): The graph to convert. 

34 config (TomlStylingConfig | None): Serialization configuration. 

35 

36 Returns: 

37 str: A TOML string containing 'nodes' and 'edges'. 

38 """ 

39 try: 

40 import tomli_w 

41 except ImportError: 

42 logger.error("tomli-w not found. Please install it with 'pip install tomli-w'.") 

43 raise ImportError( 

44 "tomli-w is required for TOML export. Install it with 'pip install tomli-w'." 

45 ) 

46 

47 logger.debug("Creating TOML representation of the graph.") 

48 config = config or TomlStylingConfig() 

49 

50 nodes = [] 

51 edges = [] 

52 

53 for node in graph.topological_order(): 

54 node_id = config.reference_fnc(node) 

55 node_entry = { 

56 "id": node_id, 

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

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

59 } 

60 if config.node_data_fnc: 

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

62 nodes.append(node_entry) 

63 

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

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

66 

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

68 

69 return tomli_w.dumps(data) 

70 

71 

72@register_view(".toml", creator_fnc=create_topology_toml) 

73def export_topology_toml( 

74 graph: Graph, 

75 output: Path, 

76 config: TomlStylingConfig | None = None, 

77) -> None: 

78 """ 

79 Export the graph to a TOML file. 

80 

81 Args: 

82 graph (Graph): The graph to export. 

83 output (Path): The output file path. 

84 config (TomlStylingConfig | None): Serialization configuration. 

85 """ 

86 logger.info(f"Exporting TOML to: {output}") 

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

88 f.write(create_topology_toml(graph, config))