Coverage for phml\transform\transform.py: 100%

79 statements  

« prev     ^ index     » next       coverage.py v6.5.0, created at 2022-12-08 16:33 -0600

1"""phml.utils.transform.transform 

2 

3Utility methods that revolve around transforming or manipulating the ast. 

4""" 

5 

6from typing import Callable, Optional 

7 

8from phml.misc import heading_rank 

9from phml.nodes import AST, All_Nodes, Element, Root 

10from phml.travel.travel import walk 

11from phml.validate.check import Test, check 

12 

13__all__ = [ 

14 "filter_nodes", 

15 "remove_nodes", 

16 "map_nodes", 

17 "find_and_replace", 

18 "shift_heading", 

19 "replace_node", 

20 "modify_children", 

21] 

22 

23 

24def filter_nodes( 

25 tree: Root | Element | AST, 

26 condition: Test, 

27 strict: bool = True, 

28): 

29 """Take a given tree and filter the nodes with the condition. 

30 Only nodes passing the condition stay. If the parent node fails, 

31 all children are moved up in scope. Depth first 

32 

33 Same as remove_nodes but keeps the nodes that match. 

34 

35 Args: 

36 tree (Root | Element): The tree node to filter. 

37 condition (Test): The condition to apply to each node. 

38 

39 Returns: 

40 Root | Element: The given tree after being filtered. 

41 """ 

42 

43 if tree.__class__.__name__ == "AST": 

44 tree = tree.tree 

45 

46 def filter_children(node): 

47 children = [] 

48 for i, child in enumerate(node.children): 

49 if child.type in ["root", "element"]: 

50 node.children[i] = filter_children(node.children[i]) 

51 if not check(child, condition, strict=strict): 

52 for idx, _ in enumerate(child.children): 

53 child.children[idx].parent = node 

54 children.extend(node.children[i].children) 

55 else: 

56 children.append(node.children[i]) 

57 elif check(child, condition, strict=strict): 

58 children.append(node.children[i]) 

59 

60 node.children = children 

61 if len(node.children) == 0 and isinstance(node, Element): 

62 node.startend = True 

63 return node 

64 

65 filter_children(tree) 

66 

67 

68def remove_nodes( 

69 tree: Root | Element | AST, 

70 condition: Test, 

71 strict: bool = True, 

72): 

73 """Take a given tree and remove the nodes that match the condition. 

74 If a parent node is removed so is all the children. 

75 

76 Same as filter_nodes except removes nodes that match. 

77 

78 Args: 

79 tree (Root | Element): The parent node to start recursively removing from. 

80 condition (Test): The condition to apply to each node. 

81 """ 

82 if tree.__class__.__name__ == "AST": 

83 tree = tree.tree 

84 

85 def filter_children(node): 

86 node.children = [n for n in node.children if not check(n, condition, strict=strict)] 

87 for child in node.children: 

88 if child.type in ["root", "element"]: 

89 filter_children(child) 

90 

91 if len(node.children) == 0 and isinstance(node, Element): 

92 node.startend = True 

93 

94 filter_children(tree) 

95 

96 

97def map_nodes(tree: Root | Element | AST, transform: Callable): 

98 """Takes a tree and a callable that returns a node and maps each node. 

99 

100 Signature for the transform function should be as follows: 

101 

102 1. Takes a single argument that is the node. 

103 2. Returns any type of node that is assigned to the original node. 

104 

105 ```python 

106 def to_links(node): 

107 return Element("a", {}, node.parent, children=node.children) 

108 if node.type == "element" 

109 else node 

110 ``` 

111 

112 Args: 

113 tree (Root | Element): Tree to transform. 

114 transform (Callable): The Callable that returns a node that is assigned 

115 to each node. 

116 """ 

117 

118 if tree.__class__.__name__ == "AST": 

119 tree = tree.tree 

120 

121 def recursive_map(node): 

122 for i, child in enumerate(node.children): 

123 if isinstance(child, Element): 

124 recursive_map(node.children[i]) 

125 node.children[i] = transform(child) 

126 else: 

127 node.children[i] = transform(child) 

128 

129 recursive_map(tree) 

130 

131 

132def replace_node( 

133 start: Root | Element, 

134 condition: Test, 

135 replacement: Optional[All_Nodes | list[All_Nodes]], 

136 strict: bool = True, 

137): 

138 """Search for a specific node in the tree and replace it with either 

139 a node or list of nodes. If replacement is None the found node is just removed. 

140 

141 Args: 

142 start (Root | Element): The starting point. 

143 condition (test): Test condition to find the correct node. 

144 replacement (All_Nodes | list[All_Nodes] | None): What to replace the node with. 

145 """ 

146 for node in walk(start): 

147 if check(node, condition, strict=strict): 

148 if node.parent is not None: 

149 idx = node.parent.children.index(node) 

150 if replacement is not None: 

151 parent = node.parent 

152 parent.children = ( 

153 node.parent.children[:idx] + replacement + node.parent.children[idx + 1 :] 

154 if isinstance(replacement, list) 

155 else node.parent.children[:idx] 

156 + [replacement] 

157 + node.parent.children[idx + 1 :] 

158 ) 

159 else: 

160 parent = node.parent 

161 parent.children.pop(idx) 

162 if len(parent.children) == 0 and isinstance(parent, Element): 

163 parent.startend = True 

164 

165 

166def find_and_replace(start: Root | Element, *replacements: tuple[str, str | Callable]) -> int: 

167 """Takes a ast, root, or any node and replaces text in `text` 

168 nodes with matching replacements. 

169 

170 First value in each replacement tuple is the regex to match and 

171 the second value is what to replace it with. This can either be 

172 a string or a callable that returns a string or a new node. If 

173 a new node is returned then the text element will be split. 

174 """ 

175 from re import finditer # pylint: disable=import-outside-toplevel 

176 

177 for node in walk(start): 

178 if node.type == "text": 

179 for replacement in replacements: 

180 if isinstance(replacement[1], str): 

181 for match in finditer(replacement[0], node.value): 

182 node.value = ( 

183 node.value[: match.start()] + replacement[1] + node.value[match.end() :] 

184 ) 

185 

186 

187def shift_heading(node: Element, amount: int): 

188 """Shift the heading by the amount specified. 

189 

190 value is clamped between 1 and 6. 

191 """ 

192 

193 rank = heading_rank(node) 

194 rank += amount 

195 

196 node.tag = f"h{min(6, max(1, rank))}" 

197 

198 

199def modify_children(func): 

200 """Function wrapper that when called and passed an 

201 AST, Root, or Element will apply the wrapped function 

202 to each child. This means that whatever is returned 

203 from the wrapped function will be assigned to the child. 

204 

205 The wrapped function will be passed the child node, 

206 the index in the parents children, and the parent node 

207 """ 

208 from phml import visit_children # pylint: disable=import-outside-toplevel 

209 

210 def inner(start: AST | Element | Root): 

211 if isinstance(start, AST): 

212 start = start.tree 

213 

214 for idx, child in enumerate(visit_children(start)): 

215 start.children[idx] = func(child, idx, child.parent) 

216 

217 return inner