OLD | NEW |
(Empty) | |
| 1 """Utility functions for test code that uses astroid ASTs as input.""" |
| 2 import functools |
| 3 import sys |
| 4 import textwrap |
| 5 |
| 6 from astroid import nodes |
| 7 from astroid import builder |
| 8 # The name of the transient function that is used to |
| 9 # wrap expressions to be extracted when calling |
| 10 # extract_node. |
| 11 _TRANSIENT_FUNCTION = '__' |
| 12 |
| 13 # The comment used to select a statement to be extracted |
| 14 # when calling extract_node. |
| 15 _STATEMENT_SELECTOR = '#@' |
| 16 |
| 17 |
| 18 def _extract_expressions(node): |
| 19 """Find expressions in a call to _TRANSIENT_FUNCTION and extract them. |
| 20 |
| 21 The function walks the AST recursively to search for expressions that |
| 22 are wrapped into a call to _TRANSIENT_FUNCTION. If it finds such an |
| 23 expression, it completely removes the function call node from the tree, |
| 24 replacing it by the wrapped expression inside the parent. |
| 25 |
| 26 :param node: An astroid node. |
| 27 :type node: astroid.bases.NodeNG |
| 28 :yields: The sequence of wrapped expressions on the modified tree |
| 29 expression can be found. |
| 30 """ |
| 31 if (isinstance(node, nodes.CallFunc) |
| 32 and isinstance(node.func, nodes.Name) |
| 33 and node.func.name == _TRANSIENT_FUNCTION): |
| 34 real_expr = node.args[0] |
| 35 real_expr.parent = node.parent |
| 36 # Search for node in all _astng_fields (the fields checked when |
| 37 # get_children is called) of its parent. Some of those fields may |
| 38 # be lists or tuples, in which case the elements need to be checked. |
| 39 # When we find it, replace it by real_expr, so that the AST looks |
| 40 # like no call to _TRANSIENT_FUNCTION ever took place. |
| 41 for name in node.parent._astroid_fields: |
| 42 child = getattr(node.parent, name) |
| 43 if isinstance(child, (list, tuple)): |
| 44 for idx, compound_child in enumerate(child): |
| 45 if compound_child is node: |
| 46 child[idx] = real_expr |
| 47 elif child is node: |
| 48 setattr(node.parent, name, real_expr) |
| 49 yield real_expr |
| 50 else: |
| 51 for child in node.get_children(): |
| 52 for result in _extract_expressions(child): |
| 53 yield result |
| 54 |
| 55 |
| 56 def _find_statement_by_line(node, line): |
| 57 """Extracts the statement on a specific line from an AST. |
| 58 |
| 59 If the line number of node matches line, it will be returned; |
| 60 otherwise its children are iterated and the function is called |
| 61 recursively. |
| 62 |
| 63 :param node: An astroid node. |
| 64 :type node: astroid.bases.NodeNG |
| 65 :param line: The line number of the statement to extract. |
| 66 :type line: int |
| 67 :returns: The statement on the line, or None if no statement for the line |
| 68 can be found. |
| 69 :rtype: astroid.bases.NodeNG or None |
| 70 """ |
| 71 if isinstance(node, (nodes.Class, nodes.Function)): |
| 72 # This is an inaccuracy in the AST: the nodes that can be |
| 73 # decorated do not carry explicit information on which line |
| 74 # the actual definition (class/def), but .fromline seems to |
| 75 # be close enough. |
| 76 node_line = node.fromlineno |
| 77 else: |
| 78 node_line = node.lineno |
| 79 |
| 80 if node_line == line: |
| 81 return node |
| 82 |
| 83 for child in node.get_children(): |
| 84 result = _find_statement_by_line(child, line) |
| 85 if result: |
| 86 return result |
| 87 |
| 88 return None |
| 89 |
| 90 def extract_node(code, module_name=''): |
| 91 """Parses some Python code as a module and extracts a designated AST node. |
| 92 |
| 93 Statements: |
| 94 To extract one or more statement nodes, append #@ to the end of the line |
| 95 |
| 96 Examples: |
| 97 >>> def x(): |
| 98 >>> def y(): |
| 99 >>> return 1 #@ |
| 100 |
| 101 The return statement will be extracted. |
| 102 |
| 103 >>> class X(object): |
| 104 >>> def meth(self): #@ |
| 105 >>> pass |
| 106 |
| 107 The funcion object 'meth' will be extracted. |
| 108 |
| 109 Expressions: |
| 110 To extract arbitrary expressions, surround them with the fake |
| 111 function call __(...). After parsing, the surrounded expression |
| 112 will be returned and the whole AST (accessible via the returned |
| 113 node's parent attribute) will look like the function call was |
| 114 never there in the first place. |
| 115 |
| 116 Examples: |
| 117 >>> a = __(1) |
| 118 |
| 119 The const node will be extracted. |
| 120 |
| 121 >>> def x(d=__(foo.bar)): pass |
| 122 |
| 123 The node containing the default argument will be extracted. |
| 124 |
| 125 >>> def foo(a, b): |
| 126 >>> return 0 < __(len(a)) < b |
| 127 |
| 128 The node containing the function call 'len' will be extracted. |
| 129 |
| 130 If no statements or expressions are selected, the last toplevel |
| 131 statement will be returned. |
| 132 |
| 133 If the selected statement is a discard statement, (i.e. an expression |
| 134 turned into a statement), the wrapped expression is returned instead. |
| 135 |
| 136 For convenience, singleton lists are unpacked. |
| 137 |
| 138 :param str code: A piece of Python code that is parsed as |
| 139 a module. Will be passed through textwrap.dedent first. |
| 140 :param str module_name: The name of the module. |
| 141 :returns: The designated node from the parse tree, or a list of nodes. |
| 142 :rtype: astroid.bases.NodeNG, or a list of nodes. |
| 143 """ |
| 144 def _extract(node): |
| 145 if isinstance(node, nodes.Discard): |
| 146 return node.value |
| 147 else: |
| 148 return node |
| 149 |
| 150 requested_lines = [] |
| 151 for idx, line in enumerate(code.splitlines()): |
| 152 if line.strip().endswith(_STATEMENT_SELECTOR): |
| 153 requested_lines.append(idx + 1) |
| 154 |
| 155 tree = build_module(code, module_name=module_name) |
| 156 extracted = [] |
| 157 if requested_lines: |
| 158 for line in requested_lines: |
| 159 extracted.append(_find_statement_by_line(tree, line)) |
| 160 |
| 161 # Modifies the tree. |
| 162 extracted.extend(_extract_expressions(tree)) |
| 163 |
| 164 if not extracted: |
| 165 extracted.append(tree.body[-1]) |
| 166 |
| 167 extracted = [_extract(node) for node in extracted] |
| 168 if len(extracted) == 1: |
| 169 return extracted[0] |
| 170 else: |
| 171 return extracted |
| 172 |
| 173 |
| 174 def build_module(code, module_name='', path=None): |
| 175 """Parses a string module with a builder. |
| 176 :param code: The code for the module. |
| 177 :type code: str |
| 178 :param module_name: The name for the module |
| 179 :type module_name: str |
| 180 :param path: The path for the module |
| 181 :type module_name: str |
| 182 :returns: The module AST. |
| 183 :rtype: astroid.bases.NodeNG |
| 184 """ |
| 185 code = textwrap.dedent(code) |
| 186 return builder.AstroidBuilder(None).string_build(code, modname=module_name,
path=path) |
| 187 |
| 188 |
| 189 def require_version(minver=None, maxver=None): |
| 190 """ Compare version of python interpreter to the given one. Skip the test |
| 191 if older. |
| 192 """ |
| 193 def parse(string, default=None): |
| 194 string = string or default |
| 195 try: |
| 196 return tuple(int(v) for v in string.split('.')) |
| 197 except ValueError: |
| 198 raise ValueError('%s is not a correct version : should be X.Y[.Z].'
% version) |
| 199 |
| 200 def check_require_version(f): |
| 201 current = sys.version_info[:3] |
| 202 if parse(minver, "0") < current <= parse(maxver, "4"): |
| 203 return f |
| 204 else: |
| 205 str_version = '.'.join(str(v) for v in sys.version_info) |
| 206 @functools.wraps(f) |
| 207 def new_f(self, *args, **kwargs): |
| 208 if minver is not None: |
| 209 self.skipTest('Needs Python > %s. Current version is %s.' %
(minver, str_version)) |
| 210 elif maxver is not None: |
| 211 self.skipTest('Needs Python <= %s. Current version is %s.' %
(maxver, str_version)) |
| 212 return new_f |
| 213 |
| 214 |
| 215 return check_require_version |
| 216 |
| 217 def get_name_node(start_from, name, index=0): |
| 218 return [n for n in start_from.nodes_of_class(nodes.Name) if n.name == name][
index] |
OLD | NEW |