Coverage for python/lsst/pipe/base/pipeline_graph/visualization/_printer.py: 22%
101 statements
« prev ^ index » next coverage.py v7.4.4, created at 2024-03-26 02:50 -0700
« prev ^ index » next coverage.py v7.4.4, created at 2024-03-26 02:50 -0700
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__ = ("Printer", "make_default_printer", "make_colorama_printer", "make_simple_printer")
31import sys
32from collections.abc import Callable, Sequence
33from typing import Generic, TextIO
35from ._layout import _K, Layout, LayoutRow
37_CHAR_DECOMPOSITION = {
38 # This mapping provides the "logic" for how to decompose the relevant
39 # box-drawing symbols into symbols with just a single line segment,
40 # allowing us to do set-operations on just the single-segment characters.
41 " ": frozenset(),
42 "╴": frozenset({"╴"}),
43 "╵": frozenset({"╵"}),
44 "╶": frozenset({"╶"}),
45 "╷": frozenset({"╷"}),
46 "╯": frozenset({"╴", "╵"}),
47 "─": frozenset({"╴", "╶"}),
48 "╮": frozenset({"╴", "╷"}),
49 "╰": frozenset({"╵", "╶"}),
50 "│": frozenset({"╵", "╷"}),
51 "╭": frozenset({"╶", "╷"}),
52 "┴": frozenset({"╴", "╵", "╶"}),
53 "┤": frozenset({"╴", "╵", "╷"}),
54 "┬": frozenset({"╴", "╶", "╷"}),
55 "├": frozenset({"╵", "╶", "╷"}),
56 "┼": frozenset({"╴", "╵", "╶", "╷"}),
57}
59_CHAR_COMPOSITION = {v: k for k, v in _CHAR_DECOMPOSITION.items()}
62class PrintRow:
63 """Base class and default implementation for drawing a single row of a text
64 DAG visualization.
66 Parameters
67 ----------
68 width : `int`
69 Number of columns the graph will occupy (not including text
70 descriptions on the right).
71 pad : `str`
72 Character to use for empty cells.
73 """
75 def __init__(self, width: int, pad: str):
76 self._cells = [pad] * width
78 def set(self, x: int, char: str, style: tuple[str, str] = ("", "")) -> None:
79 """Set the character in a single cell, overriding what is there now.
81 Parameters
82 ----------
83 x : `int`
84 Column index.
85 char : `str`
86 Single character value to set.
87 style : `tuple` [ `str`, `str` ], optional
88 Styling prefix and suffix strings. Ignored by the base class.
89 """
90 self._cells[x] = char
92 def vert(self, x: int, style: tuple[str, str] = ("", "")) -> None:
93 """Add a vertical line, taking into account the current contents of the
94 cell.
96 Parameters
97 ----------
98 x : `int`
99 Column index.
100 style : `tuple` [ `str`, `str` ], optional
101 Styling prefix and suffix strings. Ignored by the base class.
103 Notes
104 -----
105 This merges the vertical line with any other character other than a
106 complete complete horizontal segment, which it replaces to represent a
107 non-intersection as a visual "hop", i.e. ``───`` becomes ``─│─``.
108 """
109 if self._cells[x] in (" ", "─"):
110 self.set(x, "│", style)
111 else:
112 self.update(x, "│", style)
114 def update(self, x: int, char: str, style: tuple[str, str] = ("", "")) -> None:
115 """Combine a new line-drawing character with the line segments already
116 present in a cell.
118 Parameters
119 ----------
120 x : `int`
121 Column index.
122 char : `str`
123 Single character value to add.
124 style : `tuple` [ `str`, `str` ], optional
125 Styling prefix and suffix strings. Ignored by the base class.
126 """
127 self.set(x, _CHAR_COMPOSITION[_CHAR_DECOMPOSITION[char] | _CHAR_DECOMPOSITION[self._cells[x]]], style)
129 def bend(
130 self,
131 start: int,
132 stop: int,
133 start_style: tuple[str, str] = ("", ""),
134 stop_style: tuple[str, str] = ("", ""),
135 ) -> None:
136 """Draw a sideways-S bend representing a connection in one column from
137 above to a different column below.
139 Parameters
140 ----------
141 start : `int`
142 Incoming column index (connecting from the previous line, above).
143 stop : `int`
144 Outgoing column index (connecting to the next line, below).
145 start_style : `tuple` [ `str`, `str` ], optional
146 Styling prefix and suffix strings. Ignored by the base class.
147 stop_style : `tuple` [ `str`, `str` ], optional
148 Styling prefix and suffix strings. Ignored by the base class.
150 Notes
151 -----
152 If the incoming and outgoing columns are the same, this yields a
153 vertical line.
154 """
155 if start < stop:
156 self.update(start, "╰", start_style)
157 self.update(stop, "╮", stop_style)
158 for x in range(start + 1, stop):
159 self.update(x, "─", stop_style)
160 elif start > stop:
161 self.update(start, "╯", start_style)
162 self.update(stop, "╭", stop_style)
163 for x in range(stop + 1, start):
164 self.update(x, "─", stop_style)
165 else:
166 self.update(start, "│", start_style)
168 def finish(self) -> str:
169 """Return the line printer's internal state into a string."""
170 return "".join(self._cells)
173def _default_get_text(node: _K, x: int, style: tuple[str, str]) -> str:
174 """Return the default text to associate with a node.
176 This function is the default value for the ``get_text`` argument to
177 `Printer`; see that for details.
178 """
179 return str(node)
182def _default_get_symbol(node: _K, x: int) -> str:
183 """Return the default symbol for a node.
185 This function is the default value for the ``get_symbol`` argument to
186 `Printer`; see that for details.
187 """
188 return "⬤"
191def _default_get_style(node: _K, x: int) -> tuple[str, str]:
192 """Get the default styling suffix/prefix strings.
194 This function is the default value for the ``get_style`` argument to
195 `Printer`; see that for details.
196 """
197 return "", ""
200class Printer(Generic[_K]):
201 """High-level tool for drawing a text-based DAG visualization.
203 Parameters
204 ----------
205 layout_width : `int`
206 Logical width of the layout. Actual width of the graph is
207 ``layout_width * 2 + 1`` to space out the nodes and make line
208 intersections and non-intersections comprehensible.
209 pad : `str`, optional
210 Character to use for unpopulated cells.
211 make_blank_row : `~collections.abc.Callable`
212 Callback that returns a new `PrintRow` instance with no new cells
213 populated. Callback arguments are ``(self.width, self.pad)``. This
214 can return a specialization of `PrintRow` to add support for styling.
215 get_text : `~collections.abc.Callable`
216 Callback that returns the text description for a node. Arguments
217 are the node key, the column in which the node appears, and a
218 style prefix/suffix 2-tuple.
219 get_symbol : `~collections.abc.Callable`
220 Callback that returns the symbol for a node. Arguments are the node
221 key, the column in which the node appears.
222 get_style : `~collections.abc.Callable`
223 Callback that returns a prefix/suffix style 2-tuple for a node. Prefix
224 and suffix values are strings, but are otherwise subclass-dependent;
225 styles are ignored by the base class.
226 """
228 def __init__(
229 self,
230 layout_width: int,
231 *,
232 pad: str = " ",
233 make_blank_row: Callable[[int, str], PrintRow] = PrintRow,
234 get_text: Callable[[_K, int, tuple[str, str]], str] = _default_get_text,
235 get_symbol: Callable[[_K, int], str] = _default_get_symbol,
236 get_style: Callable[[_K, int], tuple[str, str]] = _default_get_style,
237 ):
238 self.width = layout_width * 2 + 1
239 self.pad = pad
240 self.make_blank_row = make_blank_row
241 self.get_text = get_text
242 self.get_symbol = get_symbol
243 self.get_style = get_style
245 def print_row(
246 self,
247 stream: TextIO,
248 layout_row: LayoutRow[_K],
249 ) -> None:
250 """Print a single row of the DAG visualization to a file-like object.
252 Parameters
253 ----------
254 stream : `TextIO`
255 Output stream to write to.
256 layout_row : `LayoutRow`
257 Struct that indicates the columns in which nodes an edges should
258 be drawn.
259 """
260 node_style = self.get_style(layout_row.node, layout_row.x)
261 if layout_row.continuing or layout_row.connecting:
262 print_row = self.make_blank_row(self.width, self.pad)
263 for x, source in layout_row.connecting:
264 print_row.bend(
265 2 * x,
266 2 * layout_row.x,
267 start_style=self.get_style(source, x),
268 stop_style=node_style,
269 )
270 for x, source, _ in layout_row.continuing:
271 print_row.vert(2 * x, self.get_style(source, x))
272 stream.write(print_row.finish())
273 stream.write("\n")
274 print_row = self.make_blank_row(self.width, self.pad)
275 for x, source, _ in layout_row.continuing:
276 print_row.vert(2 * x, self.get_style(source, x))
277 print_row.set(2 * layout_row.x, self.get_symbol(layout_row.node, layout_row.x), node_style)
278 stream.write(print_row.finish())
279 stream.write(self.pad * 2)
280 stream.write(self.get_text(layout_row.node, layout_row.x, node_style))
281 stream.write("\n")
283 def print(self, stream: TextIO, layout: Layout) -> None:
284 """Print the DAG visualization to a file-like object.
286 Parameters
287 ----------
288 stream : `TextIO`
289 Output stream to write to.
290 layout : `Layout`
291 Struct that determines the rows and columns in which nodes and
292 edges are drawn.
293 """
294 for layout_row in layout:
295 self.print_row(stream, layout_row)
298class TerminalPrintRow(PrintRow):
299 """Specialization of `PrintRow` for interactive terminals.
301 This adds support for styling (mostly colors) using terminal escape
302 sequences.
304 Parameters
305 ----------
306 width : `int`
307 Number of columns the graph will occupy (not including text
308 descriptions on the right).
309 pad : `str`
310 Character to use for empty cells.
311 """
313 def __init__(self, width: int, pad: str):
314 super().__init__(width, pad)
315 self._styles = [("", "")] * width
317 def set(self, x: int, char: str, style: tuple[str, str] = ("", "")) -> None:
318 super().set(x, char)
319 self._styles[x] = style
321 def finish(self) -> str:
322 return "".join(f"{prefix}{char}{suffix}" for char, (prefix, suffix) in zip(self._cells, self._styles))
325def make_colorama_printer(layout_width: int, palette: Sequence[str] = ()) -> Printer | None:
326 """Return a `Printer` instance that uses terminal escape codes provided
327 by the `colorama` package.
329 Parameters
330 ----------
331 layout_width : `int`
332 Logical width of the layout. Actual width of the graph is
333 ``layout_width * 2 + 1`` to space out the nodes and make line
334 intersections and non-intersections comprehensible.
335 palette : `~collections.abc.Sequence` [ `str` ], optional
336 Sequence of colors, in which each is a single character (any of
337 ``rgbcym``), a color name (any of ``red``, ``green``, ``blue``,
338 ``cyan``, ``yellow``, or ``magenta``), or one of those color names
339 preceded by ``light`` (with no space). Case is ignored. If empty,
340 a predefined sequence using both light and dark colors is used.
342 Returns
343 -------
344 printer : `Printer` or `None`
345 A `Printer` instance, or `None` if the `colorama` package could not
346 be imported.
347 """
348 try:
349 import colorama
350 except ImportError:
351 return None
352 if not palette:
353 palette = [
354 colorama.Fore.RED,
355 colorama.Fore.LIGHTBLUE_EX,
356 colorama.Fore.GREEN,
357 colorama.Fore.LIGHTMAGENTA_EX,
358 colorama.Fore.YELLOW,
359 colorama.Fore.LIGHTCYAN_EX,
360 colorama.Fore.LIGHTRED_EX,
361 colorama.Fore.BLUE,
362 colorama.Fore.LIGHTGREEN_EX,
363 colorama.Fore.MAGENTA,
364 colorama.Fore.LIGHTYELLOW_EX,
365 colorama.Fore.CYAN,
366 ]
367 else:
368 translate_color = {
369 "R": colorama.Fore.RED,
370 "RED": colorama.Fore.RED,
371 "LIGHTRED": colorama.Fore.LIGHTRED_EX,
372 "G": colorama.Fore.GREEN,
373 "GREEN": colorama.Fore.GREEN,
374 "LIGHTGREEN": colorama.Fore.LIGHTGREEN_EX,
375 "B": colorama.Fore.BLUE,
376 "BLUE": colorama.Fore.BLUE,
377 "LIGHTBLUE": colorama.Fore.LIGHTBLUE_EX,
378 "C": colorama.Fore.CYAN,
379 "CYAN": colorama.Fore.CYAN,
380 "LIGHTCYAN": colorama.Fore.LIGHTCYAN_EX,
381 "Y": colorama.Fore.YELLOW,
382 "YELLOW": colorama.Fore.YELLOW,
383 "LIGHTYELLOW": colorama.Fore.LIGHTYELLOW_EX,
384 "M": colorama.Fore.MAGENTA,
385 "MAGENTA": colorama.Fore.MAGENTA,
386 "LIGHTMAGENTA": colorama.Fore.LIGHTMAGENTA_EX,
387 }
388 palette = [translate_color.get(c.upper(), c) for c in palette]
389 return Printer(
390 layout_width,
391 make_blank_row=lambda width, pad: TerminalPrintRow(width, pad),
392 get_style=lambda node, x: (palette[x % len(palette)], colorama.Style.RESET_ALL),
393 )
396def make_simple_printer(layout_width: int) -> Printer:
397 """Return a simple `Printer` instance with no styling.
399 Parameters
400 ----------
401 layout_width : `int`
402 Logical width of the layout. Actual width of the graph is
403 ``layout_width * 2 + 1`` to space out the nodes and make line
404 intersections and non-intersections comprehensible.
405 """
406 return Printer(layout_width)
409def make_default_printer(
410 layout_width: int, color: bool | Sequence[str] | None = None, stream: TextIO = sys.stdout
411) -> Printer:
412 """Return a `Printer` instance with terminal escape-code coloring if
413 possible (and requested), or a simple one otherwise.
415 Parameters
416 ----------
417 layout_width : `int`
418 Logical width of the layout. Actual width of the graph is
419 ``layout_width * 2 + 1`` to space out the nodes and make line
420 intersections and non-intersections comprehensible.
421 color : `bool` or `~collections.abc.Sequence` [ `str` ]
422 Whether to use terminal escape codes to add color to the graph. Default
423 is to add colors only if the `colorama` package can be imported.
424 `False` disables colors unconditionally, while `True` or a sequence of
425 colors (see `make_colorama_printer`) will result in `ImportError` being
426 propagated up if `colorama` is unavailable.
427 stream : `io.TextIO`, optional
428 Output stream the printer will write to.
429 """
430 if color is None:
431 if stream.isatty():
432 if printer := make_colorama_printer(layout_width):
433 return printer
434 elif color:
435 palette = color if color is not True else ()
436 printer = make_colorama_printer(layout_width, palette)
437 if printer is None:
438 raise ImportError("Cannot use color unless the 'colorama' module is available.")
439 return printer
440 return make_simple_printer(layout_width)