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 

28logger = logging.getLogger(__name__) 

29 

30 

31def save_graph( 

32 recipe_file: Path | str, 

33 save_path: Path | None = None, 

34 auto_open: bool = False, 

35 detailed: bool = False, 

36): 

37 """ 

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

39 

40 Parameters 

41 ---------- 

42 recipe_file: Path | str 

43 The recipe to be graphed. 

44 

45 save_path: Path 

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

47 

48 auto_open: bool 

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

50 

51 detailed: bool 

52 Whether to include operator arguments on the graph. 

53 

54 Raises 

55 ------ 

56 ValueError 

57 Recipe is invalid. 

58 """ 

59 recipe = parse_recipe(recipe_file) 

60 if save_path is None: 

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

62 

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

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

65 logger.debug("Executing step: %s", step) 

66 node = str(uuid4()) 

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

68 kwargs = {} 

69 for key, value in step.items(): 

70 if isinstance(value, dict) and "operator" in value: 

71 logger.debug("Recursing into argument: %s", key) 

72 sub_node = step_parser(value, prev_node) 

73 graph.add_edge(sub_node, node) 

74 elif key != "operator": 

75 kwargs[key] = value 

76 graph.add_edge(prev_node, node) 

77 

78 if detailed: 

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

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

81 ) 

82 return node 

83 

84 graph = pygraphviz.AGraph(directed=True) 

85 

86 prev_node = "START" 

87 graph.add_node(prev_node) 

88 try: 

89 for step in recipe["steps"]: 

90 prev_node = step_parser(step, prev_node) 

91 except KeyError as err: 

92 raise ValueError("Invalid recipe") from err 

93 

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

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

96 

97 if auto_open: 

98 try: 

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

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

101 subprocess.run( 

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

103 ) 

104 except (subprocess.CalledProcessError, FileNotFoundError): 

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

106 print( 

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

108 file=sys.stderr, 

109 )