Coverage for python/lsst/pipe/base/pipeline_graph/visualization/_printer.py: 22%

101 statements  

« prev     ^ index     » next       coverage.py v7.3.2, created at 2023-10-11 09:32 +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 

28 

29__all__ = ("Printer", "make_default_printer", "make_colorama_printer", "make_simple_printer") 

30 

31import sys 

32from collections.abc import Callable, Sequence 

33from typing import Generic, TextIO 

34 

35from ._layout import _K, Layout, LayoutRow 

36 

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} 

58 

59_CHAR_COMPOSITION = {v: k for k, v in _CHAR_DECOMPOSITION.items()} 

60 

61 

62class PrintRow: 

63 """Base class and default implementation for drawing a single row of a text 

64 DAG visualization. 

65 

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 """ 

74 

75 def __init__(self, width: int, pad: str): 

76 self._cells = [pad] * width 

77 

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. 

80 

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 

91 

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. 

95 

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. 

102 

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) 

113 

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. 

117 

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) 

128 

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. 

138 

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. 

149 

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) 

167 

168 def finish(self) -> str: 

169 """Return the line printer's internal state into a string.""" 

170 return "".join(self._cells) 

171 

172 

173def _default_get_text(node: _K, x: int, style: tuple[str, str]) -> str: 

174 """Return the default text to associate with a node. 

175 

176 This function is the default value for the ``get_text`` argument to 

177 `Printer`; see that for details. 

178 """ 

179 return str(node) 

180 

181 

182def _default_get_symbol(node: _K, x: int) -> str: 

183 """Return the default symbol for a node. 

184 

185 This function is the default value for the ``get_symbol`` argument to 

186 `Printer`; see that for details. 

187 """ 

188 return "⬤" 

189 

190 

191def _default_get_style(node: _K, x: int) -> tuple[str, str]: 

192 """Get the default styling suffix/prefix strings. 

193 

194 This function is the default value for the ``get_style`` argument to 

195 `Printer`; see that for details. 

196 """ 

197 return "", "" 

198 

199 

200class Printer(Generic[_K]): 

201 """High-level tool for drawing a text-based DAG visualization. 

202 

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 """ 

227 

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 

244 

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. 

251 

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") 

282 

283 def print(self, stream: TextIO, layout: Layout) -> None: 

284 """Print the DAG visualization to a file-like object. 

285 

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) 

296 

297 

298class TerminalPrintRow(PrintRow): 

299 """Specialization of `PrintRow` for interactive terminals. 

300 

301 This adds support for styling (mostly colors) using terminal escape 

302 sequences. 

303 

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 """ 

312 

313 def __init__(self, width: int, pad: str): 

314 super().__init__(width, pad) 

315 self._styles = [("", "")] * width 

316 

317 def set(self, x: int, char: str, style: tuple[str, str] = ("", "")) -> None: 

318 super().set(x, char) 

319 self._styles[x] = style 

320 

321 def finish(self) -> str: 

322 return "".join(f"{prefix}{char}{suffix}" for char, (prefix, suffix) in zip(self._cells, self._styles)) 

323 

324 

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. 

328 

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. 

341 

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 ) 

394 

395 

396def make_simple_printer(layout_width: int) -> Printer: 

397 """Return a simple `Printer` instance with no styling. 

398 

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) 

407 

408 

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. 

414 

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)