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

1from csv import writer as csv_writer 

2from io import StringIO 

3from logging import getLogger 

4from pathlib import Path 

5from typing import Any, Callable 

6 

7from ..graph import Graph 

8from ..graphable import Graphable 

9from ..registry import register_view 

10 

11logger = getLogger(__name__) 

12 

13 

14def create_topology_csv( 

15 graph: Graph, 

16 node_text_fnc: Callable[[Graphable[Any]], str] = lambda n: n.reference, 

17 include_header: bool = True, 

18) -> str: 

19 """ 

20 Generate a CSV edge list from a Graph. 

21 

22 Args: 

23 graph (Graph): The graph to convert. 

24 node_text_fnc: Function to generate the text for each node. 

25 include_header: Whether to include a header row ('source','target'). 

26 

27 Returns: 

28 str: The CSV edge list as a string. 

29 """ 

30 logger.debug("Creating CSV edge list.") 

31 output = StringIO() 

32 writer = csv_writer(output) 

33 

34 if include_header: 

35 writer.writerow(["source", "target"]) 

36 

37 for node in graph.topological_order(): 

38 source_text = node_text_fnc(node) 

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

40 target_text = node_text_fnc(dependent) 

41 writer.writerow([source_text, target_text]) 

42 

43 return output.getvalue() 

44 

45 

46@register_view(".csv", creator_fnc=create_topology_csv) 

47def export_topology_csv( 

48 graph: Graph, 

49 output: Path, 

50 node_text_fnc: Callable[[Graphable[Any]], str] = lambda n: n.reference, 

51 include_header: bool = True, 

52) -> None: 

53 """ 

54 Export the graph to a CSV file. 

55 

56 Args: 

57 graph (Graph): The graph to export. 

58 output (Path): The output file path. 

59 node_text_fnc: Function to generate the text for each node. 

60 include_header: Whether to include a header row. 

61 """ 

62 logger.info(f"Exporting CSV edge list to: {output}") 

63 with open(output, "w+", newline="") as f: 

64 f.write(create_topology_csv(graph, node_text_fnc, include_header))