Coverage for /pythoncovmergedfiles/medio/medio/usr/local/lib/python3.8/site-packages/pasta/base/ast_utils.py: 50%

1# coding=utf-8 

2"""Helpers for working with python ASTs.""" 

3# Copyright 2021 Google LLC 

4# 

5# Licensed under the Apache License, Version 2.0 (the "License"); 

6# you may not use this file except in compliance with the License. 

7# You may obtain a copy of the License at 

8# 

9# https://www.apache.org/licenses/LICENSE-2.0 

10# 

11# Unless required by applicable law or agreed to in writing, software 

12# distributed under the License is distributed on an "AS IS" BASIS, 

13# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 

14# See the License for the specific language governing permissions and 

15# limitations under the License. 

16 

17from __future__ import absolute_import 

18from __future__ import division 

19from __future__ import print_function 

20 

21import ast 

22import re 

23 

24from pasta.augment import errors 

25from pasta.base import formatting as fmt 

26 

27# From PEP-0263 -- https://www.python.org/dev/peps/pep-0263/ 

28_CODING_PATTERN = re.compile('^[ \t\v]*#.*?coding[:=][ \t]*([-_.a-zA-Z0-9]+)') 

29 

30 

31_AST_OP_NODES = ( 

32 ast.And, ast.Or, ast.Eq, ast.NotEq, ast.Is, ast.IsNot, ast.In, ast.NotIn, 

33 ast.Lt, ast.LtE, ast.Gt, ast.GtE, ast.Add, ast.Sub, ast.Mult, ast.Div, 

34 ast.Mod, ast.Pow, ast.LShift, ast.RShift, ast.BitAnd, ast.BitOr, ast.BitXor, 

35 ast.FloorDiv, ast.Invert, ast.Not, ast.UAdd, ast.USub 

36) 

37 

38 

39class _TreeNormalizer(ast.NodeTransformer): 

40 """Replaces all op nodes with unique instances.""" 

41 

42 def visit(self, node): 

43 if isinstance(node, _AST_OP_NODES): 

44 return node.__class__() 

45 return super(_TreeNormalizer, self).visit(node) 

46 

47 

48_tree_normalizer = _TreeNormalizer() 

49 

50 

51def parse(src): 

52 """Replaces ast.parse; ensures additional properties on the parsed tree. 

53 

54 This enforces the assumption that each node in the ast is unique. 

55 """ 

56 tree = ast.parse(sanitize_source(src)) 

57 _tree_normalizer.visit(tree) 

58 return tree 

59 

60 

61def sanitize_source(src): 

62 """Strip the 'coding' directive from python source code, if present. 

63 

64 This is a workaround for https://bugs.python.org/issue18960. Also see PEP-0263. 

65 """ 

66 src_lines = src.splitlines(True) 

67 for i, line in enumerate(src_lines[:2]): 

68 if _CODING_PATTERN.match(line): 

69 src_lines[i] = re.sub('#.*$', '# (removed coding)', line) 

70 return ''.join(src_lines) 

71 

72 

73def find_nodes_by_type(node, accept_types): 

74 visitor = FindNodeVisitor(lambda n: isinstance(n, accept_types)) 

75 visitor.visit(node) 

76 return visitor.results 

77 

78 

79class FindNodeVisitor(ast.NodeVisitor): 

80 

81 def __init__(self, condition): 

82 self._condition = condition 

83 self.results = [] 

84 

85 def visit(self, node): 

86 if self._condition(node): 

87 self.results.append(node) 

88 super(FindNodeVisitor, self).visit(node) 

89 

90 

91def get_last_child(node): 

92 """Get the last child node of a block statement. 

93 

94 The input must be a block statement (e.g. ast.For, ast.With, etc). 

95 

96 Examples: 

97 1. with first(): 

98 second() 

99 last() 

100 

101 2. try: 

102 first() 

103 except: 

104 second() 

105 finally: 

106 last() 

107 

108 In both cases, the last child is the node for `last`. 

109 """ 

110 if isinstance(node, ast.Module): 

111 try: 

112 return node.body[-1] 

113 except IndexError: 

114 return None 

115 if isinstance(node, ast.If): 

116 if (len(node.orelse) == 1 and isinstance(node.orelse[0], ast.If) and 

117 fmt.get(node.orelse[0], 'is_elif')): 

118 return get_last_child(node.orelse[0]) 

119 if node.orelse: 

120 return node.orelse[-1] 

121 elif isinstance(node, ast.With): 

122 if (len(node.body) == 1 and isinstance(node.body[0], ast.With) and 

123 fmt.get(node.body[0], 'is_continued')): 

124 return get_last_child(node.body[0]) 

125 elif hasattr(ast, 'Try') and isinstance(node, ast.Try): 

126 if node.finalbody: 

127 return node.finalbody[-1] 

128 if node.orelse: 

129 return node.orelse[-1] 

130 elif hasattr(ast, 'TryFinally') and isinstance(node, ast.TryFinally): 

131 if node.finalbody: 

132 return node.finalbody[-1] 

133 elif hasattr(ast, 'TryExcept') and isinstance(node, ast.TryExcept): 

134 if node.orelse: 

135 return node.orelse[-1] 

136 if node.handlers: 

137 return get_last_child(node.handlers[-1]) 

138 return node.body[-1] 

139 

140 

141def remove_child(parent, child): 

142 for _, field_value in ast.iter_fields(parent): 

143 if isinstance(field_value, list) and child in field_value: 

144 field_value.remove(child) 

145 return 

146 raise errors.InvalidAstError('Unable to find list containing child %r on ' 

147 'parent node %r' % (child, parent)) 

148 

149 

150def replace_child(parent, node, replace_with): 

151 """Replace a node's child with another node while preserving formatting. 

152 

153 Arguments: 

154 parent: (ast.AST) Parent node to replace a child of. 

155 node: (ast.AST) Child node to replace. 

156 replace_with: (ast.AST) New child node. 

157 """ 

158 # TODO(soupytwist): Don't refer to the formatting dict directly 

159 if hasattr(node, fmt.PASTA_DICT): 

160 fmt.set(replace_with, 'prefix', fmt.get(node, 'prefix')) 

161 fmt.set(replace_with, 'suffix', fmt.get(node, 'suffix')) 

162 for field in parent._fields: 

163 field_val = getattr(parent, field, None) 

164 if field_val == node: 

165 setattr(parent, field, replace_with) 

166 return 

167 elif isinstance(field_val, list): 

168 try: 

169 field_val[field_val.index(node)] = replace_with 

170 return 

171 except ValueError: 

172 pass 

173 raise errors.InvalidAstError('Node %r is not a child of %r' % (node, parent)) 

174 

175 

176def has_docstring(node): 

177 return (hasattr(node, 'body') and node.body and 

178 isinstance(node.body[0], ast.Expr) and 

179 isinstance(node.body[0].value, ast.Str))