Coverage for src/CSET/graph.py: 100%
44 statements
« prev ^ index » next coverage.py v7.5.4, created at 2024-07-01 09:02 +0000
« prev ^ index » next coverage.py v7.5.4, created at 2024-07-01 09:02 +0000
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.
15"""Visualise recipe into a graph."""
17import logging
18import subprocess
19import sys
20import tempfile
21from pathlib import Path
22from typing import Union
23from uuid import uuid4
25import pygraphviz
27from CSET._common import parse_recipe
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.
39 Parameters
40 ----------
41 recipe_file: Path | str
42 The recipe to be graphed.
44 save_path: Path
45 Path where to save the generated image. Defaults to a temporary file.
47 auto_open: bool
48 Whether to automatically open the graph with the default image viewer.
50 detailed: bool
51 Whether to include operator arguments on the graph.
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")
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)
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
83 graph = pygraphviz.AGraph(directed=True)
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
94 graph.draw(save_path, format="svg", prog="dot")
95 print(f"Graph rendered to {save_path}")
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 )