Coverage for src / graphable / views / yaml.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 YamlStylingConfig: 

15 """ 

16 Configuration for YAML graph serialization. 

17 

18 Attributes: 

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

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

21 indent: YAML indentation level. 

22 """ 

23 

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

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

26 indent: int = 2 

27 

28 

29def create_topology_yaml(graph: Graph, config: YamlStylingConfig | None = None) -> str: 

30 """ 

31 Generate a YAML representation of the graph. 

32 Requires 'PyYAML' to be installed. 

33 

34 Args: 

35 graph (Graph): The graph to convert. 

36 config (YamlStylingConfig | None): Serialization configuration. 

37 

38 Returns: 

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

40 """ 

41 try: 

42 from yaml import dump 

43 except ImportError: 

44 logger.error("PyYAML not found. Please install it with 'pip install PyYAML'.") 

45 raise ImportError( 

46 "PyYAML is required for YAML export. Install it with 'pip install PyYAML'." 

47 ) 

48 

49 logger.debug("Creating YAML representation of the graph.") 

50 config = config or YamlStylingConfig() 

51 

52 nodes = [] 

53 edges = [] 

54 

55 for node in graph.topological_order(): 

56 node_id = config.reference_fnc(node) 

57 node_entry = { 

58 "id": node_id, 

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

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

61 } 

62 if config.node_data_fnc: 

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

64 nodes.append(node_entry) 

65 

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

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

68 

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

70 

71 return dump(data, indent=config.indent, sort_keys=False) 

72 

73 

74@register_view([".yaml", ".yml"], creator_fnc=create_topology_yaml) 

75def export_topology_yaml( 

76 graph: Graph, 

77 output: Path, 

78 config: YamlStylingConfig | None = None, 

79) -> None: 

80 """ 

81 Export the graph to a YAML file. 

82 

83 Args: 

84 graph (Graph): The graph to export. 

85 output (Path): The output file path. 

86 config (YamlStylingConfig | None): Serialization configuration. 

87 """ 

88 logger.info(f"Exporting YAML to: {output}") 

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

90 f.write(create_topology_yaml(graph, config))