Coverage for src / graphable / parsers / yaml.py: 89%

1from logging import getLogger 

2from pathlib import Path 

3from typing import Any 

4 

5from ..graph import Graph 

6from ..registry import register_parser 

7from .utils import build_graph_from_data, is_path 

8 

9logger = getLogger(__name__) 

10 

11 

12@register_parser([".yaml", ".yml"]) 

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

14 """ 

15 Load a graph from a YAML string or file. 

16 Requires 'PyYAML' to be installed. 

17 

18 Args: 

19 source: YAML string or path to a YAML file. 

20 reference_type: The type to cast the node reference to (default: str). 

21 

22 Returns: 

23 Graph: A new Graph instance populated from the YAML data. 

24 """ 

25 try: 

26 from yaml import safe_load 

27 except ImportError: 

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

29 raise ImportError( 

30 "PyYAML is required for YAML parsing. Install it with 'pip install PyYAML'." 

31 ) 

32 

33 if is_path(source): 

34 logger.debug(f"Loading YAML from file: {source}") 

35 with open(source, "r") as f: 

36 data = safe_load(f) 

37 else: 

38 logger.debug("Loading YAML from string.") 

39 data = safe_load(str(source)) 

40 

41 # Handle wrapped structure 

42 if "graph" in data and ("nodes" not in data or "edges" not in data): 

43 data = data["graph"] 

44 

45 nodes_data = data.get("nodes", []) 

46 edges_data = data.get("edges", []) 

47 

48 g = build_graph_from_data(nodes_data, edges_data, reference_type) 

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

50 return g