Coverage for src / graphable / parsers / toml.py: 78%

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(".toml") 

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

14 """ 

15 Load a graph from a TOML string or file. 

16 Uses 'tomllib' (Python 3.11+) or 'tomli'. 

17 

18 Args: 

19 source: TOML string or path to a TOML 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 TOML data. 

24 """ 

25 try: 

26 from tomllib import load as toml_load 

27 from tomllib import loads as toml_loads 

28 except ImportError: 

29 try: 

30 from tomli import load as toml_load # type: ignore 

31 from tomli import loads as toml_loads # type: ignore 

32 except ImportError: 

33 logger.error("No TOML parser found. On Python < 3.11, install 'tomli'.") 

34 raise ImportError( 

35 "TOML parsing requires 'tomllib' (Python 3.11+) or 'tomli'." 

36 ) 

37 

38 if is_path(source): 

39 logger.debug(f"Loading TOML from file: {source}") 

40 with open(source, "rb") as f: 

41 data = toml_load(f) 

42 else: 

43 logger.debug("Loading TOML from string.") 

44 data = toml_loads(str(source)) 

45 

46 # Handle wrapped structure 

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

48 data = data["graph"] 

49 

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

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

52 

53 g = build_graph_from_data(nodes_data, edges_data, reference_type) 

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

55 return g