Coverage for phml\utils\transform\transform.py: 18%

60 statements  

« prev     ^ index     » next       coverage.py v6.5.0, created at 2022-11-30 09:38 -0600

1from typing import Callable, Optional 

2 

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

4from phml.utils.misc import heading_rank 

5from phml.utils.travel import walk 

6from phml.utils.validate.test import Test, test 

7 

8__all__ = [ 

9 "filter_nodes", 

10 "remove_nodes", 

11 "map_nodes", 

12 "find_and_replace", 

13 "shift_heading", 

14 "replace_node", 

15] 

16 

17 

18def filter_nodes(tree: Root | Element | AST, condition: Test): 

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

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

21 then all children are removed. 

22 

23 Same as remove_nodes but keeps the nodes that match. 

24 

25 Args: 

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

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

28 

29 Returns: 

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

31 """ 

32 

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

34 tree = tree.tree 

35 

36 def filter_children(node): 

37 node.children = [n for n in node.children if test(n, condition)] 

38 for child in node.children: 

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

40 filter_children(child) 

41 

42 filter_children(tree) 

43 

44 

45def remove_nodes(tree: Root | Element | AST, condition: Test): 

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

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

48 

49 Same as filter_nodes except removes nodes that match. 

50 

51 Args: 

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

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

54 """ 

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

56 tree = tree.tree 

57 

58 def filter_children(node): 

59 node.children = [n for n in node.children if not test(n, condition)] 

60 for child in node.children: 

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

62 filter_children(child) 

63 

64 filter_children(tree) 

65 

66 

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

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

69 

70 Signature for the transform function should be as follows: 

71 

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

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

74 

75 ```python 

76 def to_links(node): 

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

78 if node.type == "element" 

79 else node 

80 ``` 

81 

82 Args: 

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

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

85 to each node. 

86 """ 

87 

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

89 tree = tree.tree 

90 

91 for node in walk(tree): 

92 if not isinstance(node, Root): 

93 node = transform(node) 

94 

95 

96def replace_node( 

97 node: Root | Element, condition: Test, replacement: Optional[All_Nodes | list[All_Nodes]] 

98): 

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

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

101 

102 Args: 

103 node (Root | Element): The starting point. 

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

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

106 """ 

107 for n in walk(node): 

108 if test(n, condition): 

109 if n.parent is not None: 

110 idx = n.parent.children.index(n) 

111 if replacement is not None: 

112 n.parent.children = ( 

113 n.parent.children[:idx] + replacement + n.parent.children[idx + 1 :] 

114 if isinstance(replacement, list) 

115 else n.parent.children[:idx] + [replacement] + n.parent.children[idx + 1 :] 

116 ) 

117 else: 

118 n.parent.children.pop(idx) 

119 break 

120 

121 

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

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

124 nodes with matching replacements. 

125 

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

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

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

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

130 """ 

131 from re import finditer 

132 

133 for n in walk(node): 

134 if n.type == "text": 

135 for replacement in replacements: 

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

137 for match in finditer(replacement[0], n.value): 

138 n.value = n.value[: match.start()] + replacement[1] + n.value[match.end() :] 

139 else: 

140 raise NotImplementedError( 

141 "Callables are not yet supported for find_and_replace operations." 

142 ) 

143 # TODO add ability to inject nodes in place of text replacement 

144 # elif isinstance(replacement[1], Callable): 

145 # for match in finditer(replacement[0], n.value): 

146 # result = replacement[1](match.group()) 

147 # if isinstance(result, str): 

148 # n.value = n.value[:match.start()] + replacement[1] + n.value[match.end():] 

149 # elif isinstance(result, All_Nodes): 

150 # pass 

151 # elif isinstance(result, list): 

152 # pass 

153 

154 

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

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

157 

158 value is clamped between 1 and 6. 

159 """ 

160 

161 rank = heading_rank(node) 

162 rank += amount 

163 

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

165 

166 

167def modify_children(func): 

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

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

170 to each child. This means that whatever is returned 

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

172 

173 The wrapped function will be passed the child node, 

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

175 """ 

176 from phml.utils import visit_children 

177 

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

179 if isinstance(start, AST): 

180 start = start.tree 

181 

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

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

184 

185 return inner