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

1# Copyright 2022-2023 Met Office and 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 typing import Union 

23from uuid import uuid4 

24 

25import pygraphviz 

26 

27from CSET._common import parse_recipe 

28 

29 

30def save_graph( 

31 recipe_file: Union[Path, str], 

32 save_path: Path = None, 

33 auto_open: bool = False, 

34 detailed: bool = False, 

35): 

36 """ 

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

38 

39 Parameters 

40 ---------- 

41 recipe_file: Path | str 

42 The recipe to be graphed. 

43 

44 save_path: Path 

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

46 

47 auto_open: bool 

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

49 

50 detailed: bool 

51 Whether to include operator arguments on the graph. 

52 

53 Raises 

54 ------ 

55 ValueError 

56 Recipe is invalid. 

57 """ 

58 recipe = parse_recipe(recipe_file) 

59 if save_path is None: 

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

61 

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

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

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

65 node = str(uuid4()) 

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

67 kwargs = {} 

68 for key in step.keys(): 

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

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

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

72 graph.add_edge(sub_node, node) 

73 elif key != "operator": 

74 kwargs[key] = step[key] 

75 graph.add_edge(prev_node, node) 

76 

77 if detailed: 

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

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

80 ) 

81 return node 

82 

83 graph = pygraphviz.AGraph(directed=True) 

84 

85 prev_node = "START" 

86 graph.add_node(prev_node) 

87 try: 

88 # TODO: Expand to cover collate too. 

89 for step in recipe["parallel"]: 

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 )