Coverage for src / graphable / parsers / graphml.py: 81%

1from logging import getLogger 

2from pathlib import Path 

3from typing import Any 

4 

5from defusedxml import ElementTree as ET 

6 

7from ..graph import Graph 

8from ..registry import register_parser 

9from .utils import build_graph_from_data, is_path 

10 

11logger = getLogger(__name__) 

12 

13 

14@register_parser(".graphml") 

15def load_graph_graphml(source: str | Path, reference_type: type = str) -> Graph[Any]: 

16 """ 

17 Load a Graph from a GraphML XML source. 

18 

19 Args: 

20 source: GraphML XML string or path to a GraphML file. 

21 reference_type: The type to cast the reference string to. 

22 

23 Returns: 

24 Graph: The loaded Graph instance. 

25 """ 

26 if is_path(source): 

27 tree = ET.parse(source) 

28 root = tree.getroot() 

29 else: 

30 root = ET.fromstring(str(source)) 

31 

32 # GraphML uses namespaces 

33 ns = {"g": "http://graphml.graphdrawing.org/xmlns"} 

34 

35 graph_elem = root.find("g:graph", ns) 

36 if graph_elem is None: 

37 # Fallback for no namespace 

38 graph_elem = root.find("graph") 

39 if graph_elem is None: 

40 return Graph() 

41 ns = {} 

42 

43 nodes_data = [] 

44 for node_elem in graph_elem.findall("g:node" if ns else "node", ns): 

45 node_id = node_elem.get("id") 

46 node_entry = {"id": node_id} 

47 

48 # Handle data fields (tags, etc) 

49 for data_elem in node_elem.findall("g:data" if ns else "data", ns): 

50 key = data_elem.get("key") 

51 if key == "tags" and data_elem.text: 

52 node_entry["tags"] = data_elem.text.split(",") 

53 elif key == "duration" and data_elem.text: 

54 node_entry["duration"] = data_elem.text 

55 elif key == "status" and data_elem.text: 

56 node_entry["status"] = data_elem.text 

57 

58 nodes_data.append(node_entry) 

59 

60 edges_data = [] 

61 for edge_elem in graph_elem.findall("g:edge" if ns else "edge", ns): 

62 u_id = edge_elem.get("source") 

63 v_id = edge_elem.get("target") 

64 edge_entry = {"source": u_id, "target": v_id} 

65 

66 # Handle edge attributes 

67 for data_elem in edge_elem.findall("g:data" if ns else "data", ns): 

68 key = data_elem.get("key") 

69 if key and data_elem.text: 

70 edge_entry[key] = data_elem.text 

71 

72 edges_data.append(edge_entry) 

73 

74 g = build_graph_from_data(nodes_data, edges_data, reference_type) 

75 logger.info(f"Loaded graph with {len(g)} nodes from GraphML.") 

76 return g