Coverage for python/lsst/pipe/base/pipeline_graph/visualization/_show.py: 81%
70 statements
« prev ^ index » next coverage.py v7.16.2, created at 2026-10-07 22:57 +0000
1# This file is part of pipe_base.
2#
3# Developed for the LSST Data Management System.
4# This product includes software developed by the LSST Project
5# (http://www.lsst.org).
6# See the COPYRIGHT file at the top-level directory of this distribution
7# for details of code ownership.
8#
9# This software is dual licensed under the GNU General Public License and also
10# under a 3-clause BSD license. Recipients may choose which of these licenses
11# to use; please see the files gpl-3.0.txt and/or bsd_license.txt,
12# respectively. If you choose the GPL option then the following text applies
13# (but note that there is still no warranty even if you opt for BSD instead):
14#
15# This program is free software: you can redistribute it and/or modify
16# it under the terms of the GNU General Public License as published by
17# the Free Software Foundation, either version 3 of the License, or
18# (at your option) any later version.
19#
20# This program is distributed in the hope that it will be useful,
21# but WITHOUT ANY WARRANTY; without even the implied warranty of
22# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
23# GNU General Public License for more details.
24#
25# You should have received a copy of the GNU General Public License
26# along with this program. If not, see <http://www.gnu.org/licenses/>.
27from __future__ import annotations
29__all__ = ("parse_display_args", "show")
31import sys
32from collections.abc import Sequence
33from shutil import get_terminal_size
34from typing import Any, Literal, TextIO
36import networkx
38from .._nodes import NodeKey, NodeType
39from .._pipeline_graph import PipelineGraph
40from .._tasks import TaskInitNode, TaskNode
41from ._formatting import GetNodeText, get_node_symbol
42from ._layout import ColumnSelector, Layout
43from ._merge import (
44 MergedNodeKey,
45 merge_graph_input_trees,
46 merge_graph_intermediates,
47 merge_graph_output_trees,
48)
49from ._options import NodeAttributeOptions
50from ._printer import make_default_printer
51from ._status_annotator import NodeStatusAnnotator, NodeStatusOptions
53DisplayNodeKey = NodeKey | MergedNodeKey
56def parse_display_args(
57 pipeline_graph: PipelineGraph,
58 *,
59 dataset_types: bool = False,
60 init: bool | None = False,
61 dimensions: Literal["full"] | Literal["concise"] | Literal[False] | None = None,
62 task_classes: Literal["full"] | Literal["concise"] | Literal[False] = False,
63 storage_classes: bool | None = None,
64 merge_input_trees: int = 4,
65 merge_output_trees: int = 4,
66 merge_intermediates: bool = True,
67 include_automatic_connections: bool = False,
68 status_annotator: NodeStatusAnnotator | None = None,
69 status_options: NodeStatusOptions | None = None,
70) -> tuple[networkx.DiGraph | networkx.MultiDiGraph, NodeAttributeOptions]:
71 """Print a text-based ~.PipelineGraph` visualization.
73 Parameters
74 ----------
75 pipeline_graph : `PipelineGraph`
76 Graph to display.
77 dataset_types : `bool`, optional
78 Whether to include dataset type nodes (default is `False`).
79 init : `bool`, optional
80 Whether to include task initialization nodes (i.e. the producers of
81 init-input and init-output dataset types). Default is `False`. `None`
82 will show both runtime and init nodes, while `True` will show only
83 init nodes.
84 dimensions : `str` or `False`, optional
85 How to display the dimensions of tasks and dataset types.
87 Possible values include:
89 - ``"full"``: report fully-expanded dimensions.
90 - ``"concise"``: report only dimensions that are not required or
91 implied dependencies of any reported dimension.
92 - `False`: do not report dimensions at all.
93 - `None` (default): report concise dimensions only if
94 `PipelineGraph.is_fully_resolved` is `True`.
96 This also sets whether merge options consider dimensions when merging
97 nodes.
98 task_classes : `str` or `False`, optional
99 How to display task class names (not task labels, which are always
100 shown).
102 Possible values include:
104 - ``"full"``: report the fully-qualified task class name.
105 - ``"concise"``: report the task class name with no module
106 or package.
107 - `False`: (default) do not report task classes at all.
109 This also sets whether merge options consider task classes when merging
110 nodes. It is ignored if ``tasks=False``.
111 storage_classes : `bool`, optional
112 Whether to display storage classes in dataset type nodes. This also
113 sets whether merge options consider storage classes when merging nodes.
114 It is ignored if ``dataset_types=False``.
115 merge_input_trees : `int`, optional
116 If positive, merge input trees of the graph whose nodes have the same
117 outputs and other properties (dimensions, task classes, storage
118 classes), traversing this many nodes deep into the graph from the
119 beginning. Default is ``4``.
120 merge_output_trees : `int`, optional
121 If positive, merge output trees of the graph whose nodes have the same
122 outputs and other properties (dimensions, task classes, storage
123 classes), traversing this many nodes deep into the graph from the end.
124 Default is ``4``.
125 merge_intermediates : `bool`, optional
126 If `True` (default) merge interior parallel nodes with the same inputs,
127 outputs, and other properties (dimensions, task classes, storage
128 classes).
129 include_automatic_connections : `bool`, optional
130 Whether to include automatically-added connections like the config,
131 log, and metadata dataset types for each task. Default is `False`.
132 status_annotator : `NodeStatusAnnotator`, optional
133 Annotator to add status information to the graph. Default is `None`.
134 status_options : `NodeStatusOptions`, optional
135 Options for displaying execution status. Default is `None`.
136 """
137 if init is None: 137 ↛ 138line 137 didn't jump to line 138 because the condition on line 137 was never true
138 if not dataset_types:
139 raise ValueError("Cannot show init and runtime graphs unless dataset types are shown.")
140 xgraph = pipeline_graph.make_xgraph()
141 if status_annotator is not None:
142 raise ValueError("Cannot show status with both init and runtime graphs.")
143 elif dataset_types:
144 xgraph = pipeline_graph.make_bipartite_xgraph(init)
145 if status_annotator is not None: 145 ↛ 146line 145 didn't jump to line 146 because the condition on line 145 was never true
146 status_annotator(xgraph, dataset_types=True)
147 else:
148 xgraph = pipeline_graph.make_task_xgraph(init)
149 storage_classes = False
150 if status_annotator is not None: 150 ↛ 151line 150 didn't jump to line 151 because the condition on line 150 was never true
151 status_annotator(xgraph, dataset_types=False)
153 options = NodeAttributeOptions(
154 dimensions=dimensions,
155 storage_classes=storage_classes,
156 task_classes=task_classes,
157 status=status_options,
158 )
159 options = options.checked(pipeline_graph.is_fully_resolved, has_status=status_annotator is not None)
161 if dataset_types and not include_automatic_connections:
162 taskish_nodes: list[TaskNode | TaskInitNode] = []
163 for task_node in pipeline_graph.tasks.values():
164 if init is None or init is False: 164 ↛ 166line 164 didn't jump to line 166 because the condition on line 164 was always true
165 taskish_nodes.append(task_node)
166 if init is None or init is True: 166 ↛ 167line 166 didn't jump to line 167 because the condition on line 166 was never true
167 taskish_nodes.append(task_node.init)
168 for t in taskish_nodes:
169 xgraph.remove_nodes_from(
170 edge.dataset_type_key
171 for edge in t.iter_all_outputs()
172 if edge.connection_name not in t.outputs and not xgraph.out_degree(edge.dataset_type_key)
173 )
175 if dataset_types:
176 _annotate_prerequisite_nodes(xgraph)
178 if merge_input_trees:
179 merge_graph_input_trees(xgraph, options, depth=merge_input_trees)
180 if merge_output_trees:
181 merge_graph_output_trees(xgraph, options, depth=merge_output_trees)
182 if merge_intermediates:
183 merge_graph_intermediates(xgraph, options)
185 return xgraph, options
188def _annotate_prerequisite_nodes(xgraph: networkx.DiGraph | networkx.MultiDiGraph) -> None:
189 """Ensure every dataset type node has an ``is_prerequisite`` attribute.
191 Resolved dataset type nodes already carry this attribute, but unresolved
192 ones do not, so it is derived from their outgoing edges instead.
194 Parameters
195 ----------
196 xgraph : `networkx.DiGraph` or `networkx.MultiDiGraph`
197 Graph to be processed; modified in-place.
198 """
199 for node, state in xgraph.nodes.items():
200 if node.node_type is NodeType.DATASET_TYPE and "is_prerequisite" not in state:
201 state["is_prerequisite"] = any(
202 data.get("is_prerequisite", False) for _, _, data in xgraph.out_edges(node, data=True)
203 )
206def show(
207 pipeline_graph: PipelineGraph,
208 stream: TextIO = sys.stdout,
209 *,
210 color: bool | Sequence[str] | None = None,
211 width: int = -1,
212 column_interior_penalty: int = 1,
213 column_crossing_penalty: int = 1,
214 column_insertion_penalty: int = 2,
215 **kwargs: Any,
216) -> None:
217 """Print a text-based ~.PipelineGraph` visualization.
219 Parameters
220 ----------
221 pipeline_graph : `PipelineGraph`
222 Graph to display.
223 stream : `TextIO`, optional
224 Output stream. Defaults to STDOUT.
225 color : `bool` or `~collections.abc.Sequence` [ `str` ]
226 Whether to use tto add color to the graph. Default is to add colors
227 only if the `colorama` package can be imported. `False` disables colors
228 unconditionally, while `True` or a sequence of colors (see
229 `make_colorama_printer`) will result in `ImportError` being propagated
230 up if `colorama` is unavailable.
231 width : `int`, optional
232 Number of columns the full graph should occupy, including text
233 descriptions on the right. If ``0``, there is no limit (and hence no
234 text-wrapping). If negative (default) use the current terminal width.
235 column_interior_penalty : `int`
236 Penalty applied to a prospective column for a node when that column is
237 between two existing columns.
238 column_crossing_penalty : `int`
239 Penalty applied to a prospective column for a node for each ongoing
240 (vertical) edge that node's incoming edges would have to "hop".
241 column_insertion_penalty : `int`
242 Penalty applied to a prospective column for a node when considering a
243 new columns on the sides or between two existing columns.
244 **kwargs
245 Forwarded to parse_display_args.
246 """
247 xgraph, options = parse_display_args(pipeline_graph, **kwargs)
249 column_selector = ColumnSelector(
250 interior_penalty=column_interior_penalty,
251 crossing_penalty=column_crossing_penalty,
252 insertion_penalty=column_insertion_penalty,
253 )
254 layout = Layout[DisplayNodeKey](xgraph, column_selector)
256 if width < 0: 256 ↛ 261line 256 didn't jump to line 261 because the condition on line 256 was always true
257 width, _ = get_terminal_size()
259 # Number of columns used for padding after a symbol.
260 # Must match the padding added by the `Printer` class.
261 symbol_padding = 2
263 printer = make_default_printer(layout.width, color, stream)
264 printer.get_symbol = get_node_symbol
266 get_text = GetNodeText(xgraph, options, (width - printer.width - symbol_padding) if width else 0)
267 printer.get_text = get_text
269 printer.print(stream, layout)
270 for line in get_text.format_deferrals(width): 270 ↛ 271line 270 didn't jump to line 271 because the loop on line 270 never started
271 print(line, file=stream)