Coverage for /home/deng/Projects/ete4/hackathon/ete4/ete4/core/text_viz.py: 12%

56 statements  

« prev     ^ index     » next       coverage.py v7.2.7, created at 2024-08-07 10:27 +0200

1""" 

2Tree text visualization. 

3 

4Functions to show a string drawing of a tree suitable for printing on 

5a console. 

6""" 

7 

8# These functions would not normally be used direclty. Instead, they 

9# are used when doing things like: 

10# print(t) 

11# or like: 

12# t.to_str(...) 

13 

14 

15def to_str(tree, show_internal=True, compact=False, props=None, 

16 px=None, py=None, px0=0, cascade=False): 

17 """Return a string containing an ascii drawing of the tree. 

18 

19 :param show_internal: If True, show the internal nodes too. 

20 :param compact: If True, use exactly one line per tip. 

21 :param props: List of node properties to show. If None, show all. 

22 :param px, py, px0: Paddings (x, y, x for leaves). Overrides `compact`. 

23 :param cascade: Use a cascade representation. Overrides 

24 `show_internal`, `compact`, `px`, `py`, `px0`. 

25 """ 

26 if not cascade: 

27 px = px if px is not None else (0 if show_internal else 1) 

28 py = py if py is not None else (0 if compact else 1) 

29 

30 lines, _ = ascii_art(tree, show_internal, props, px, py, px0) 

31 return '\n'.join(lines) 

32 else: 

33 px = px if px is not None else 1 

34 return to_cascade(tree, props, px) 

35 

36 

37# For representations like: 

38# ╭╴a 

39# ╴h╶┤ ╭╴b 

40# ╰╴g╶┼╴e╶┬╴c 

41# │ ╰╴d 

42# ╰╴f 

43 

44def ascii_art(tree, show_internal=True, props=None, px=0, py=0, px0=0): 

45 """Return list of strings representing the tree, and their middle point. 

46 

47 :param tree: Tree to represent as ascii art. 

48 :param show_internal: If True, show the internal node names too. 

49 :param props: List of properties to show for each node. If None, show all. 

50 :param px, py: Padding in x and y. 

51 :param px0: Padding in x for leaves. 

52 """ 

53 # Node description (including all the requested properties). 

54 descr = ','.join( 

55 (f'{k}={v}' for k, v in tree.props.items()) if props is None else 

56 (str(tree.get_prop(p, '')) or '⊗' for p in props)) 

57 

58 if tree.is_leaf: 

59 return (['─' * px0 + '╴' + descr], 0) 

60 

61 lines = [] 

62 padding = ((px0 + 1 + len(descr) + 1) if show_internal else 0) + px 

63 for child in tree.children: 

64 lines_child, mid = ascii_art(child, show_internal, props, px, py, px0) 

65 

66 if len(tree.children) == 1: # only one child 

67 lines += add_prefix(lines_child, padding, mid, ' ', 

68 '─', 

69 ' ') 

70 pos_first = mid 

71 pos_last = len(lines) - mid 

72 elif child == tree.children[0]: # first child 

73 lines += add_prefix(lines_child, padding, mid, ' ', 

74 '╭', 

75 '│') 

76 lines.extend([' ' * padding + '│'] * py) # y padding 

77 pos_first = mid 

78 elif child != tree.children[-1]: # a child in the middle 

79 lines += add_prefix(lines_child, padding, mid, '│', 

80 '├', 

81 '│') 

82 lines.extend([' ' * padding + '│'] * py) # y padding 

83 else: # last child 

84 lines += add_prefix(lines_child, padding, mid, '│', 

85 '╰', 

86 ' ') 

87 pos_last = len(lines_child) - mid 

88 

89 mid = (pos_first + len(lines) - pos_last) // 2 # middle point 

90 

91 lines[mid] = add_base(lines[mid], px, px0, descr, show_internal) 

92 

93 return lines, mid 

94 

95 

96def add_prefix(lines, px, mid, c1, c2, c3): 

97 """Return the given lines adding a prefix. 

98 

99 :param lines: List of strings, to return with prefixes. 

100 :param int px: Padding in x. 

101 :param int mid: Middle point (index of the row where the node would hang). 

102 :param c1, c2, c3: Character to use as prefix before, at, and after mid. 

103 """ 

104 prefix = lambda i: ' ' * px + (c1 if i < mid else (c2 if i == mid else c3)) 

105 

106 return [prefix(i) + line for i, line in enumerate(lines)] 

107 

108 

109def add_base(line, px, px0, txt, show_internal): 

110 """Return the same line but adding a base line.""" 

111 # Example of change at the beginning of line: ' │' -> '─┤' 

112 replacements = { 

113 '│': '┤', 

114 '─': '╌', 

115 '├': '┼', 

116 '╭': '┬'} 

117 

118 padding = ((px0 + 1 + len(txt) + 1) if show_internal else 0) + px 

119 

120 prefix_txt = '─' * px0 + (f'╴{txt}╶' if txt else '──') 

121 

122 return ((prefix_txt if show_internal else '') + 

123 '─' * px + replacements[line[padding]] + line[padding+1:]) 

124 

125 

126# For representations like: 

127# h 

128# ├─╴a 

129# └─┐g 

130# ├─╴b 

131# ├─┐e 

132# │ ├─╴c 

133# │ └─╴d 

134# └─╴f 

135 

136def to_cascade(tree, props=None, px=1, are_last=None): 

137 """Return string with a visual representation of the tree as a cascade.""" 

138 are_last = are_last or [] 

139 

140 # Node description (including all the requested properties). 

141 descr = ','.join( 

142 (f'{k}={v}' for k, v in tree.props.items()) if props is None else 

143 (str(tree.get_prop(p, '')) or '⊗' for p in props)) 

144 

145 branches = get_branches_repr(are_last, tree.is_leaf, px) 

146 

147 wf = lambda n, lasts: to_cascade(n, props, px, lasts) # shortcut 

148 

149 return '\n'.join([branches + descr] + 

150 [wf(n, are_last + [False]) for n in tree.children[:-1]] + 

151 [wf(n, are_last + [True] ) for n in tree.children[-1:]]) 

152 

153 

154def get_branches_repr(are_last, is_leaf, px): 

155 """Return a text line representing the open branches according to are_last. 

156 

157 :param are_last: List of bools that say per level if we are the last node. 

158 :param is_leaf: says if the node to represent in this line has no children. 

159 :param px: Padding in x. 

160 

161 Example (for is_leaf=True, px=6):: 

162 

163 [True , False, True , True , True ] -> 

164 '│ │ │ ├──────╴' 

165 """ 

166 if len(are_last) == 0: 

167 return '' 

168 

169 prefix = ''.join((' ' if is_last else '│') + ' ' * px 

170 for is_last in are_last[:-1]) 

171 

172 return (prefix + ('└' if are_last[-1] else '├') + 

173 '─' * px + ('╴' if is_leaf else '┐')) 

174 

175 

176def to_repr(tree, depth=4, nchildren=3): 

177 """Return a text representation that exactly recreates the tree. 

178 

179 If depth and nchildren are None, return the full representation. 

180 """ 

181 children = tree.children[:nchildren] 

182 depth_1 = depth if depth is None else depth - 1 

183 children_repr = '...' if depth == 0 else ( 

184 ', '.join(to_repr(node, depth_1, nchildren) for node in children) + 

185 ('' if nchildren is None or len(tree.children) <= nchildren 

186 else ', ...')) 

187 

188 return 'Tree(%r, [%s])' % (tree.props, children_repr)