Coverage for python/lsst/pipe/base/pipeline_graph/visualization/_formatting.py: 15%

82 statements  

« prev     ^ index     » next       coverage.py v7.3.2, created at 2023-12-06 10:56 +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__ = ("get_node_symbol", "GetNodeText") 

30 

31import textwrap 

32from collections.abc import Iterator 

33 

34import networkx 

35import networkx.algorithms.community 

36from lsst.daf.butler import DimensionGroup 

37 

38from .._nodes import NodeKey, NodeType 

39from ._merge import MergedNodeKey 

40from ._options import NodeAttributeOptions 

41 

42DisplayNodeKey = NodeKey | MergedNodeKey 

43"""Type alias for graph keys that may be original task, task init, or dataset 

44type keys, or a merge of several keys for display purposes. 

45""" 

46 

47 

48def get_node_symbol(node: DisplayNodeKey, x: int | None = None) -> str: 

49 """Return a single-character symbol for a particular node type. 

50 

51 Parameters 

52 ---------- 

53 node : `DisplayNodeKey` 

54 Named tuple used as the node key. 

55 x : `str`, optional 

56 Ignored; may be passed for compatibility with the `Printer` class's 

57 ``get_symbol`` callback. 

58 

59 Returns 

60 ------- 

61 symbol : `str` 

62 A single-character symbol. 

63 """ 

64 match node: 

65 case NodeKey(node_type=NodeType.TASK): 

66 return "■" 

67 case NodeKey(node_type=NodeType.DATASET_TYPE): 

68 return "○" 

69 case NodeKey(node_type=NodeType.TASK_INIT): 

70 return "▣" 

71 case MergedNodeKey(node_type=NodeType.TASK): 

72 return "▥" 

73 case MergedNodeKey(node_type=NodeType.DATASET_TYPE): 

74 return "◍" 

75 case MergedNodeKey(node_type=NodeType.TASK_INIT): 

76 return "▤" 

77 raise ValueError(f"Unexpected node key: {node} of type {type(node)}.") 

78 

79 

80class GetNodeText: 

81 """A callback for the `Printer` class's `get_text` callback that 

82 prints detailed information about a node and can defer long entries to 

83 a footnote. 

84 

85 Parameters 

86 ---------- 

87 xgraph : `networkx.DiGraph` or `networkx.MultiDiGraph` 

88 NetworkX export of a `.PipelineGraph` that is being displayed. 

89 options : `NodeAttributeOptions` 

90 Options for how much information to display. 

91 width : `int` or `None` 

92 Number of display columns that node text can occupy. `None` for 

93 unlimited. 

94 """ 

95 

96 def __init__( 

97 self, 

98 xgraph: networkx.DiGraph | networkx.MultiDiGraph, 

99 options: NodeAttributeOptions, 

100 width: int | None, 

101 ): 

102 self.xgraph = xgraph 

103 self.options = options 

104 self.width = width 

105 self.deferred: list[tuple[str, tuple[str, str], list[str]]] = [] 

106 

107 def __call__(self, node: DisplayNodeKey, x: int, style: tuple[str, str]) -> str: 

108 state = self.xgraph.nodes[node] 

109 terms: list[str] = [f"{node}:" if self.options else str(node)] 

110 if self.options.dimensions and node.node_type != NodeType.TASK_INIT: 

111 terms.append(self.format_dimensions(state["dimensions"])) 

112 if ( 

113 self.options.task_classes 

114 and node.node_type is NodeType.TASK 

115 or node.node_type is NodeType.TASK_INIT 

116 ): 

117 terms.append(self.format_task_class(state["task_class_name"])) 

118 if self.options.storage_classes and node.node_type is NodeType.DATASET_TYPE: 

119 terms.append(state["storage_class_name"]) 

120 description = " ".join(terms) 

121 if self.width and len(description) > self.width: 

122 index = f"[{len(self.deferred) + 1}]" 

123 self.deferred.append((index, style, terms)) 

124 return f"{description[:self.width - len(index) - 6]}...{style[0]}{index}{style[1]} " 

125 return description 

126 

127 def format_dimensions(self, dimensions: DimensionGroup) -> str: 

128 """Format the dimensions of a task or dataset type node.""" 

129 match self.options.dimensions: 

130 case "full": 

131 return str(dimensions.names) 

132 case "concise": 

133 kept = set(dimensions.names) 

134 done = False 

135 while not done: 

136 for dimension in kept: 

137 if any( 

138 dimension != other and dimension in dimensions.universe[other].dimensions 

139 for other in kept 

140 ): 

141 kept.remove(dimension) 

142 break 

143 else: 

144 done = True 

145 # We still iterate over dimensions instead of kept to preserve 

146 # order. 

147 return f"{{{', '.join(d for d in dimensions.names if d in kept)}}}" 

148 case False: 

149 return "" 

150 raise ValueError(f"Invalid display option for dimensions: {self.options.dimensions!r}.") 

151 

152 def format_task_class(self, task_class_name: str) -> str: 

153 """Format the type object for a task or task init node.""" 

154 match self.options.task_classes: 

155 case "full": 

156 return task_class_name 

157 case "concise": 

158 return task_class_name.split(".")[-1] 

159 case False: 

160 return "" 

161 raise ValueError(f"Invalid display option for task_classes: {self.options.task_classes!r}.") 

162 

163 def format_deferrals(self, width: int | None) -> Iterator[str]: 

164 """Iterate over all descriptions that were truncated earlier and 

165 replace with footnote placeholders. 

166 

167 Parameters 

168 ---------- 

169 width : `int` or `None` 

170 Number of columns to wrap descriptions at. 

171 

172 Returns 

173 ------- 

174 deferrals : `collections.abc.Iterator` [ `str` ] 

175 Lines of deferred text, already wrapped. 

176 """ 

177 indent = " " 

178 for index, style, terms in self.deferred: 

179 yield f"{style[0]}{index}{style[1]}" 

180 for term in terms: 

181 if width: 

182 yield from textwrap.wrap(term, width, initial_indent=indent, subsequent_indent=indent) 

183 else: 

184 yield term