Coverage for src/CSET/graph.py: 100%

1# © Crown copyright, Met Office (2022-2024) and CSET contributors. 

2# 

3# Licensed under the Apache License, Version 2.0 (the "License"); 

4# you may not use this file except in compliance with the License. 

5# You may obtain a copy of the License at 

6# 

7# http://www.apache.org/licenses/LICENSE-2.0 

8# 

9# Unless required by applicable law or agreed to in writing, software 

10# distributed under the License is distributed on an "AS IS" BASIS, 

11# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 

12# See the License for the specific language governing permissions and 

13# limitations under the License. 

14 

15"""Visualise recipe into a graph.""" 

16 

17import logging 

18import subprocess 

19import sys 

20import tempfile 

21from pathlib import Path 

22from uuid import uuid4 

23 

24import pygraphviz 

25 

26from CSET._common import parse_recipe 

27 

28 

29def save_graph( 

30 recipe_file: Path | str, 

31 save_path: Path = None, 

32 auto_open: bool = False, 

33 detailed: bool = False, 

34): 

35 """ 

36 Draws out the graph of a recipe, and saves it to a file. 

37 

38 Parameters 

39 ---------- 

40 recipe_file: Path | str 

41 The recipe to be graphed. 

42 

43 save_path: Path 

44 Path where to save the generated image. Defaults to a temporary file. 

45 

46 auto_open: bool 

47 Whether to automatically open the graph with the default image viewer. 

48 

49 detailed: bool 

50 Whether to include operator arguments on the graph. 

51 

52 Raises 

53 ------ 

54 ValueError 

55 Recipe is invalid. 

56 """ 

57 recipe = parse_recipe(recipe_file) 

58 if save_path is None: 

59 save_path = Path(f"{tempfile.gettempdir()}/{uuid4()}.svg") 

60 

61 def step_parser(step: dict, prev_node: str) -> str: 

62 """Parse recipe to add nodes to graph and link them with edges.""" 

63 logging.debug("Executing step: %s", step) 

64 node = str(uuid4()) 

65 graph.add_node(node, label=step["operator"]) 

66 kwargs = {} 

67 for key in step.keys(): 

68 if isinstance(step[key], dict) and "operator" in step[key]: 

69 logging.debug("Recursing into argument: %s", key) 

70 sub_node = step_parser(step[key], prev_node) 

71 graph.add_edge(sub_node, node) 

72 elif key != "operator": 

73 kwargs[key] = step[key] 

74 graph.add_edge(prev_node, node) 

75 

76 if detailed: 

77 graph.get_node(node).attr["label"] = f"{step['operator']}\n" + "".join( 

78 f"<{key}: {kwargs[key]}>\n" for key in kwargs 

79 ) 

80 return node 

81 

82 graph = pygraphviz.AGraph(directed=True) 

83 

84 prev_node = "START" 

85 graph.add_node(prev_node) 

86 try: 

87 for step in recipe["steps"]: 

88 prev_node = step_parser(step, prev_node) 

89 except KeyError as err: 

90 raise ValueError("Invalid recipe") from err 

91 

92 graph.draw(save_path, format="svg", prog="dot") 

93 print(f"Graph rendered to {save_path}") 

94 

95 if auto_open: 

96 try: 

97 # Stderr is redirected here to suppress gvfs-open deprecation warning. 

98 # See https://bugs.python.org/issue30219 for an example. 

99 subprocess.run( 

100 ("xdg-open", str(save_path)), check=True, stderr=subprocess.DEVNULL 

101 ) 

102 except (subprocess.CalledProcessError, FileNotFoundError): 

103 # Using print rather than logging as this is run interactively. 

104 print( 

105 "Cannot automatically display graph. Specify an output with -o instead.", 

106 file=sys.stderr, 

107 )