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

1import xml.etree.ElementTree as ET 

2from dataclasses import dataclass 

3from logging import getLogger 

4from pathlib import Path 

5from typing import Any, Callable 

6from xml.dom import minidom 

7 

8from ..graph import Graph 

9from ..graphable import Graphable 

10from ..registry import register_view 

11 

12logger = getLogger(__name__) 

13 

14 

15@dataclass 

16class GraphmlStylingConfig: 

17 """ 

18 Configuration for GraphML serialization. 

19 

20 Attributes: 

21 node_ref_fnc: Function to generate the unique ID for each node. 

22 attr_mapping: Dictionary mapping Graphable attributes to GraphML data keys. 

23 """ 

24 

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

26 # We could extend this to support more complex data mapping if needed. 

27 

28 

29def create_topology_graphml( 

30 graph: Graph, config: GraphmlStylingConfig | None = None 

31) -> str: 

32 """ 

33 Generate a GraphML (XML) representation of the graph. 

34 

35 Args: 

36 graph (Graph): The graph to convert. 

37 config (GraphmlStylingConfig | None): Export configuration. 

38 

39 Returns: 

40 str: The GraphML XML string. 

41 """ 

42 logger.debug("Creating GraphML representation.") 

43 config = config or GraphmlStylingConfig() 

44 

45 # Create root element 

46 root = ET.Element( 

47 "graphml", 

48 { 

49 "xmlns": "http://graphml.graphdrawing.org/xmlns", 

50 "xmlns:xsi": "http://www.w3.org/2001/XMLSchema-instance", 

51 "xsi:schemaLocation": "http://graphml.graphdrawing.org/xmlns http://graphml.graphdrawing.org/xmlns/1.0/graphml.xsd", 

52 }, 

53 ) 

54 

55 # Define attributes (keys) 

56 # 1. Tags 

57 ET.SubElement( 

58 root, 

59 "key", 

60 { 

61 "id": "tags", 

62 "for": "node", 

63 "attr.name": "tags", 

64 "attr.type": "string", 

65 }, 

66 ) 

67 

68 # Create graph element 

69 graph_elem = ET.SubElement(root, "graph", {"id": "G", "edgedefault": "directed"}) 

70 

71 # Nodes 

72 for node in graph.topological_order(): 

73 node_id = config.node_ref_fnc(node) 

74 node_elem = ET.SubElement(graph_elem, "node", {"id": node_id}) 

75 

76 # Add tags as data 

77 if node.tags: 

78 data_elem = ET.SubElement(node_elem, "data", {"key": "tags"}) 

79 data_elem.text = ",".join(node.tags) 

80 

81 # Edges 

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

83 dep_id = config.node_ref_fnc(dependent) 

84 ET.SubElement( 

85 graph_elem, 

86 "edge", 

87 { 

88 "id": f"e_{node_id}_{dep_id}", 

89 "source": node_id, 

90 "target": dep_id, 

91 }, 

92 ) 

93 

94 # Pretty print XML 

95 xml_str = ET.tostring(root, encoding="utf-8") 

96 parsed_xml = minidom.parseString(xml_str) 

97 return parsed_xml.toprettyxml(indent=" ") 

98 

99 

100@register_view(".graphml", creator_fnc=create_topology_graphml) 

101def export_topology_graphml( 

102 graph: Graph, output: Path, config: GraphmlStylingConfig | None = None 

103) -> None: 

104 """ 

105 Export the graph to a GraphML (.graphml) file. 

106 

107 Args: 

108 graph (Graph): The graph to export. 

109 output (Path): The output file path. 

110 config (GraphmlStylingConfig | None): Export configuration. 

111 """ 

112 logger.info(f"Exporting GraphML to: {output}") 

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

114 f.write(create_topology_graphml(graph, config))