Martin v. Löwis | ef04c44 | 2008-03-19 05:04:44 +0000 | [diff] [blame] | 1 | """Utility functions, node construction macros, etc.""" |
| 2 | # Author: Collin Winter |
| 3 | |
Benjamin Peterson | 4eb5fa5 | 2010-08-08 19:01:25 +0000 | [diff] [blame] | 4 | from itertools import islice |
| 5 | |
Martin v. Löwis | ef04c44 | 2008-03-19 05:04:44 +0000 | [diff] [blame] | 6 | # Local imports |
Benjamin Peterson | df6dc8f | 2008-06-15 02:57:40 +0000 | [diff] [blame] | 7 | from .pgen2 import token |
| 8 | from .pytree import Leaf, Node |
| 9 | from .pygram import python_symbols as syms |
| 10 | from . import patcomp |
Martin v. Löwis | ef04c44 | 2008-03-19 05:04:44 +0000 | [diff] [blame] | 11 | |
| 12 | |
| 13 | ########################################################### |
| 14 | ### Common node-construction "macros" |
| 15 | ########################################################### |
| 16 | |
| 17 | def KeywordArg(keyword, value): |
| 18 | return Node(syms.argument, |
Benjamin Peterson | 4eb5fa5 | 2010-08-08 19:01:25 +0000 | [diff] [blame] | 19 | [keyword, Leaf(token.EQUAL, "="), value]) |
Martin v. Löwis | ef04c44 | 2008-03-19 05:04:44 +0000 | [diff] [blame] | 20 | |
| 21 | def LParen(): |
| 22 | return Leaf(token.LPAR, "(") |
| 23 | |
| 24 | def RParen(): |
| 25 | return Leaf(token.RPAR, ")") |
| 26 | |
| 27 | def Assign(target, source): |
| 28 | """Build an assignment statement""" |
| 29 | if not isinstance(target, list): |
| 30 | target = [target] |
| 31 | if not isinstance(source, list): |
Benjamin Peterson | 2c3ac6b | 2009-06-11 23:47:38 +0000 | [diff] [blame] | 32 | source.prefix = " " |
Martin v. Löwis | ef04c44 | 2008-03-19 05:04:44 +0000 | [diff] [blame] | 33 | source = [source] |
| 34 | |
| 35 | return Node(syms.atom, |
| 36 | target + [Leaf(token.EQUAL, "=", prefix=" ")] + source) |
| 37 | |
| 38 | def Name(name, prefix=None): |
| 39 | """Return a NAME leaf""" |
| 40 | return Leaf(token.NAME, name, prefix=prefix) |
| 41 | |
| 42 | def Attr(obj, attr): |
| 43 | """A node tuple for obj.attr""" |
| 44 | return [obj, Node(syms.trailer, [Dot(), attr])] |
| 45 | |
| 46 | def Comma(): |
| 47 | """A comma leaf""" |
| 48 | return Leaf(token.COMMA, ",") |
| 49 | |
| 50 | def Dot(): |
| 51 | """A period (.) leaf""" |
| 52 | return Leaf(token.DOT, ".") |
| 53 | |
| 54 | def ArgList(args, lparen=LParen(), rparen=RParen()): |
| 55 | """A parenthesised argument list, used by Call()""" |
Benjamin Peterson | cf60382 | 2008-09-01 19:56:06 +0000 | [diff] [blame] | 56 | node = Node(syms.trailer, [lparen.clone(), rparen.clone()]) |
| 57 | if args: |
| 58 | node.insert_child(1, Node(syms.arglist, args)) |
| 59 | return node |
Martin v. Löwis | ef04c44 | 2008-03-19 05:04:44 +0000 | [diff] [blame] | 60 | |
Benjamin Peterson | cf60382 | 2008-09-01 19:56:06 +0000 | [diff] [blame] | 61 | def Call(func_name, args=None, prefix=None): |
Martin v. Löwis | ef04c44 | 2008-03-19 05:04:44 +0000 | [diff] [blame] | 62 | """A function call""" |
| 63 | node = Node(syms.power, [func_name, ArgList(args)]) |
| 64 | if prefix is not None: |
Benjamin Peterson | 2c3ac6b | 2009-06-11 23:47:38 +0000 | [diff] [blame] | 65 | node.prefix = prefix |
Martin v. Löwis | ef04c44 | 2008-03-19 05:04:44 +0000 | [diff] [blame] | 66 | return node |
| 67 | |
| 68 | def Newline(): |
| 69 | """A newline literal""" |
| 70 | return Leaf(token.NEWLINE, "\n") |
| 71 | |
| 72 | def BlankLine(): |
| 73 | """A blank line""" |
| 74 | return Leaf(token.NEWLINE, "") |
| 75 | |
| 76 | def Number(n, prefix=None): |
| 77 | return Leaf(token.NUMBER, n, prefix=prefix) |
| 78 | |
| 79 | def Subscript(index_node): |
| 80 | """A numeric or string subscript""" |
Benjamin Peterson | 4eb5fa5 | 2010-08-08 19:01:25 +0000 | [diff] [blame] | 81 | return Node(syms.trailer, [Leaf(token.LBRACE, "["), |
Martin v. Löwis | ef04c44 | 2008-03-19 05:04:44 +0000 | [diff] [blame] | 82 | index_node, |
Benjamin Peterson | 4eb5fa5 | 2010-08-08 19:01:25 +0000 | [diff] [blame] | 83 | Leaf(token.RBRACE, "]")]) |
Martin v. Löwis | ef04c44 | 2008-03-19 05:04:44 +0000 | [diff] [blame] | 84 | |
| 85 | def String(string, prefix=None): |
| 86 | """A string leaf""" |
| 87 | return Leaf(token.STRING, string, prefix=prefix) |
| 88 | |
| 89 | def ListComp(xp, fp, it, test=None): |
| 90 | """A list comprehension of the form [xp for fp in it if test]. |
| 91 | |
| 92 | If test is None, the "if test" part is omitted. |
| 93 | """ |
Benjamin Peterson | 2c3ac6b | 2009-06-11 23:47:38 +0000 | [diff] [blame] | 94 | xp.prefix = "" |
| 95 | fp.prefix = " " |
| 96 | it.prefix = " " |
Martin v. Löwis | ef04c44 | 2008-03-19 05:04:44 +0000 | [diff] [blame] | 97 | for_leaf = Leaf(token.NAME, "for") |
Benjamin Peterson | 2c3ac6b | 2009-06-11 23:47:38 +0000 | [diff] [blame] | 98 | for_leaf.prefix = " " |
Martin v. Löwis | ef04c44 | 2008-03-19 05:04:44 +0000 | [diff] [blame] | 99 | in_leaf = Leaf(token.NAME, "in") |
Benjamin Peterson | 2c3ac6b | 2009-06-11 23:47:38 +0000 | [diff] [blame] | 100 | in_leaf.prefix = " " |
Martin v. Löwis | ef04c44 | 2008-03-19 05:04:44 +0000 | [diff] [blame] | 101 | inner_args = [for_leaf, fp, in_leaf, it] |
| 102 | if test: |
Benjamin Peterson | 2c3ac6b | 2009-06-11 23:47:38 +0000 | [diff] [blame] | 103 | test.prefix = " " |
Martin v. Löwis | ef04c44 | 2008-03-19 05:04:44 +0000 | [diff] [blame] | 104 | if_leaf = Leaf(token.NAME, "if") |
Benjamin Peterson | 2c3ac6b | 2009-06-11 23:47:38 +0000 | [diff] [blame] | 105 | if_leaf.prefix = " " |
Martin v. Löwis | ef04c44 | 2008-03-19 05:04:44 +0000 | [diff] [blame] | 106 | inner_args.append(Node(syms.comp_if, [if_leaf, test])) |
| 107 | inner = Node(syms.listmaker, [xp, Node(syms.comp_for, inner_args)]) |
| 108 | return Node(syms.atom, |
| 109 | [Leaf(token.LBRACE, "["), |
| 110 | inner, |
| 111 | Leaf(token.RBRACE, "]")]) |
| 112 | |
Martin v. Löwis | a675ef1 | 2008-03-24 00:50:58 +0000 | [diff] [blame] | 113 | def FromImport(package_name, name_leafs): |
| 114 | """ Return an import statement in the form: |
| 115 | from package import name_leafs""" |
| 116 | # XXX: May not handle dotted imports properly (eg, package_name='foo.bar') |
Benjamin Peterson | d613bb4 | 2008-07-16 18:44:47 +0000 | [diff] [blame] | 117 | #assert package_name == '.' or '.' not in package_name, "FromImport has "\ |
| 118 | # "not been tested with dotted package names -- use at your own "\ |
| 119 | # "peril!" |
Martin v. Löwis | a675ef1 | 2008-03-24 00:50:58 +0000 | [diff] [blame] | 120 | |
| 121 | for leaf in name_leafs: |
| 122 | # Pull the leaves out of their old tree |
| 123 | leaf.remove() |
| 124 | |
Benjamin Peterson | 4eb5fa5 | 2010-08-08 19:01:25 +0000 | [diff] [blame] | 125 | children = [Leaf(token.NAME, "from"), |
Martin v. Löwis | a675ef1 | 2008-03-24 00:50:58 +0000 | [diff] [blame] | 126 | Leaf(token.NAME, package_name, prefix=" "), |
Benjamin Peterson | 4eb5fa5 | 2010-08-08 19:01:25 +0000 | [diff] [blame] | 127 | Leaf(token.NAME, "import", prefix=" "), |
Martin v. Löwis | a675ef1 | 2008-03-24 00:50:58 +0000 | [diff] [blame] | 128 | Node(syms.import_as_names, name_leafs)] |
| 129 | imp = Node(syms.import_from, children) |
| 130 | return imp |
| 131 | |
Benjamin Peterson | 448e81b | 2012-12-07 22:44:10 -0500 | [diff] [blame] | 132 | def ImportAndCall(node, results, names): |
| 133 | """Returns an import statement and calls a method |
| 134 | of the module: |
| 135 | |
| 136 | import module |
| 137 | module.name()""" |
| 138 | obj = results["obj"].clone() |
| 139 | if obj.type == syms.arglist: |
| 140 | newarglist = obj.clone() |
| 141 | else: |
| 142 | newarglist = Node(syms.arglist, [obj.clone()]) |
| 143 | after = results["after"] |
| 144 | if after: |
| 145 | after = [n.clone() for n in after] |
| 146 | new = Node(syms.power, |
| 147 | Attr(Name(names[0]), Name(names[1])) + |
| 148 | [Node(syms.trailer, |
| 149 | [results["lpar"].clone(), |
| 150 | newarglist, |
| 151 | results["rpar"].clone()])] + after) |
| 152 | new.prefix = node.prefix |
| 153 | return new |
| 154 | |
Martin v. Löwis | a675ef1 | 2008-03-24 00:50:58 +0000 | [diff] [blame] | 155 | |
Martin v. Löwis | ef04c44 | 2008-03-19 05:04:44 +0000 | [diff] [blame] | 156 | ########################################################### |
| 157 | ### Determine whether a node represents a given literal |
| 158 | ########################################################### |
| 159 | |
| 160 | def is_tuple(node): |
| 161 | """Does the node represent a tuple literal?""" |
| 162 | if isinstance(node, Node) and node.children == [LParen(), RParen()]: |
| 163 | return True |
| 164 | return (isinstance(node, Node) |
| 165 | and len(node.children) == 3 |
| 166 | and isinstance(node.children[0], Leaf) |
| 167 | and isinstance(node.children[1], Node) |
| 168 | and isinstance(node.children[2], Leaf) |
| 169 | and node.children[0].value == "(" |
| 170 | and node.children[2].value == ")") |
| 171 | |
| 172 | def is_list(node): |
| 173 | """Does the node represent a list literal?""" |
| 174 | return (isinstance(node, Node) |
| 175 | and len(node.children) > 1 |
| 176 | and isinstance(node.children[0], Leaf) |
| 177 | and isinstance(node.children[-1], Leaf) |
| 178 | and node.children[0].value == "[" |
| 179 | and node.children[-1].value == "]") |
| 180 | |
Martin v. Löwis | ef04c44 | 2008-03-19 05:04:44 +0000 | [diff] [blame] | 181 | |
| 182 | ########################################################### |
| 183 | ### Misc |
| 184 | ########################################################### |
| 185 | |
Benjamin Peterson | 0b24b3d | 2008-12-16 03:57:54 +0000 | [diff] [blame] | 186 | def parenthesize(node): |
| 187 | return Node(syms.atom, [LParen(), node, RParen()]) |
| 188 | |
Martin v. Löwis | 3de92bf | 2008-04-10 02:50:50 +0000 | [diff] [blame] | 189 | |
Serhiy Storchaka | db9b65d | 2014-12-13 21:50:49 +0200 | [diff] [blame] | 190 | consuming_calls = {"sorted", "list", "set", "any", "all", "tuple", "sum", |
| 191 | "min", "max", "enumerate"} |
Martin v. Löwis | 3de92bf | 2008-04-10 02:50:50 +0000 | [diff] [blame] | 192 | |
Martin v. Löwis | ef04c44 | 2008-03-19 05:04:44 +0000 | [diff] [blame] | 193 | def attr_chain(obj, attr): |
| 194 | """Follow an attribute chain. |
Martin v. Löwis | f733c60 | 2008-03-19 05:26:18 +0000 | [diff] [blame] | 195 | |
Martin v. Löwis | ef04c44 | 2008-03-19 05:04:44 +0000 | [diff] [blame] | 196 | If you have a chain of objects where a.foo -> b, b.foo-> c, etc, |
| 197 | use this to iterate over all objects in the chain. Iteration is |
| 198 | terminated by getattr(x, attr) is None. |
Martin v. Löwis | f733c60 | 2008-03-19 05:26:18 +0000 | [diff] [blame] | 199 | |
Martin v. Löwis | ef04c44 | 2008-03-19 05:04:44 +0000 | [diff] [blame] | 200 | Args: |
| 201 | obj: the starting object |
| 202 | attr: the name of the chaining attribute |
Martin v. Löwis | f733c60 | 2008-03-19 05:26:18 +0000 | [diff] [blame] | 203 | |
Martin v. Löwis | ef04c44 | 2008-03-19 05:04:44 +0000 | [diff] [blame] | 204 | Yields: |
| 205 | Each successive object in the chain. |
| 206 | """ |
| 207 | next = getattr(obj, attr) |
| 208 | while next: |
| 209 | yield next |
| 210 | next = getattr(next, attr) |
| 211 | |
Martin v. Löwis | f733c60 | 2008-03-19 05:26:18 +0000 | [diff] [blame] | 212 | p0 = """for_stmt< 'for' any 'in' node=any ':' any* > |
| 213 | | comp_for< 'for' any 'in' node=any any* > |
| 214 | """ |
| 215 | p1 = """ |
| 216 | power< |
| 217 | ( 'iter' | 'list' | 'tuple' | 'sorted' | 'set' | 'sum' | |
Benjamin Peterson | 57af387 | 2012-11-29 10:55:22 -0500 | [diff] [blame] | 218 | 'any' | 'all' | 'enumerate' | (any* trailer< '.' 'join' >) ) |
Martin v. Löwis | f733c60 | 2008-03-19 05:26:18 +0000 | [diff] [blame] | 219 | trailer< '(' node=any ')' > |
| 220 | any* |
| 221 | > |
| 222 | """ |
| 223 | p2 = """ |
| 224 | power< |
Benjamin Peterson | 57af387 | 2012-11-29 10:55:22 -0500 | [diff] [blame] | 225 | ( 'sorted' | 'enumerate' ) |
Martin v. Löwis | f733c60 | 2008-03-19 05:26:18 +0000 | [diff] [blame] | 226 | trailer< '(' arglist<node=any any*> ')' > |
| 227 | any* |
| 228 | > |
| 229 | """ |
| 230 | pats_built = False |
| 231 | def in_special_context(node): |
| 232 | """ Returns true if node is in an environment where all that is required |
Benjamin Peterson | 57af387 | 2012-11-29 10:55:22 -0500 | [diff] [blame] | 233 | of it is being iterable (ie, it doesn't matter if it returns a list |
| 234 | or an iterator). |
Martin v. Löwis | f733c60 | 2008-03-19 05:26:18 +0000 | [diff] [blame] | 235 | See test_map_nochange in test_fixers.py for some examples and tests. |
| 236 | """ |
| 237 | global p0, p1, p2, pats_built |
| 238 | if not pats_built: |
Martin v. Löwis | f733c60 | 2008-03-19 05:26:18 +0000 | [diff] [blame] | 239 | p0 = patcomp.compile_pattern(p0) |
Benjamin Peterson | 57af387 | 2012-11-29 10:55:22 -0500 | [diff] [blame] | 240 | p1 = patcomp.compile_pattern(p1) |
Martin v. Löwis | f733c60 | 2008-03-19 05:26:18 +0000 | [diff] [blame] | 241 | p2 = patcomp.compile_pattern(p2) |
| 242 | pats_built = True |
| 243 | patterns = [p0, p1, p2] |
| 244 | for pattern, parent in zip(patterns, attr_chain(node, "parent")): |
| 245 | results = {} |
| 246 | if pattern.match(parent, results) and results["node"] is node: |
| 247 | return True |
| 248 | return False |
| 249 | |
Benjamin Peterson | 8bcddca | 2009-01-03 16:53:14 +0000 | [diff] [blame] | 250 | def is_probably_builtin(node): |
| 251 | """ |
| 252 | Check that something isn't an attribute or function name etc. |
| 253 | """ |
Benjamin Peterson | 608d8bc | 2009-05-05 23:23:31 +0000 | [diff] [blame] | 254 | prev = node.prev_sibling |
Benjamin Peterson | 8bcddca | 2009-01-03 16:53:14 +0000 | [diff] [blame] | 255 | if prev is not None and prev.type == token.DOT: |
| 256 | # Attribute lookup. |
| 257 | return False |
| 258 | parent = node.parent |
| 259 | if parent.type in (syms.funcdef, syms.classdef): |
| 260 | return False |
| 261 | if parent.type == syms.expr_stmt and parent.children[0] is node: |
| 262 | # Assignment. |
| 263 | return False |
| 264 | if parent.type == syms.parameters or \ |
| 265 | (parent.type == syms.typedargslist and ( |
| 266 | (prev is not None and prev.type == token.COMMA) or |
| 267 | parent.children[0] is node |
| 268 | )): |
| 269 | # The name of an argument. |
| 270 | return False |
| 271 | return True |
| 272 | |
Benjamin Peterson | 4eb5fa5 | 2010-08-08 19:01:25 +0000 | [diff] [blame] | 273 | def find_indentation(node): |
| 274 | """Find the indentation of *node*.""" |
| 275 | while node is not None: |
| 276 | if node.type == syms.suite and len(node.children) > 2: |
| 277 | indent = node.children[1] |
| 278 | if indent.type == token.INDENT: |
| 279 | return indent.value |
| 280 | node = node.parent |
| 281 | return "" |
| 282 | |
Martin v. Löwis | ef04c44 | 2008-03-19 05:04:44 +0000 | [diff] [blame] | 283 | ########################################################### |
| 284 | ### The following functions are to find bindings in a suite |
| 285 | ########################################################### |
| 286 | |
| 287 | def make_suite(node): |
| 288 | if node.type == syms.suite: |
| 289 | return node |
| 290 | node = node.clone() |
| 291 | parent, node.parent = node.parent, None |
| 292 | suite = Node(syms.suite, [node]) |
| 293 | suite.parent = parent |
| 294 | return suite |
| 295 | |
Benjamin Peterson | 0b24b3d | 2008-12-16 03:57:54 +0000 | [diff] [blame] | 296 | def find_root(node): |
| 297 | """Find the top level namespace.""" |
Martin v. Löwis | ef04c44 | 2008-03-19 05:04:44 +0000 | [diff] [blame] | 298 | # Scamper up to the top level namespace |
| 299 | while node.type != syms.file_input: |
Martin v. Löwis | ef04c44 | 2008-03-19 05:04:44 +0000 | [diff] [blame] | 300 | node = node.parent |
Benjamin Peterson | 1654d74 | 2012-09-25 11:48:50 -0400 | [diff] [blame] | 301 | if not node: |
| 302 | raise ValueError("root found before file_input node was found.") |
Benjamin Peterson | 0b24b3d | 2008-12-16 03:57:54 +0000 | [diff] [blame] | 303 | return node |
Martin v. Löwis | ef04c44 | 2008-03-19 05:04:44 +0000 | [diff] [blame] | 304 | |
Benjamin Peterson | 0b24b3d | 2008-12-16 03:57:54 +0000 | [diff] [blame] | 305 | def does_tree_import(package, name, node): |
| 306 | """ Returns true if name is imported from package at the |
| 307 | top level of the tree which node belongs to. |
| 308 | To cover the case of an import like 'import foo', use |
| 309 | None for the package and 'foo' for the name. """ |
| 310 | binding = find_binding(name, find_root(node), package) |
Martin v. Löwis | ef04c44 | 2008-03-19 05:04:44 +0000 | [diff] [blame] | 311 | return bool(binding) |
| 312 | |
Benjamin Peterson | 0b24b3d | 2008-12-16 03:57:54 +0000 | [diff] [blame] | 313 | def is_import(node): |
| 314 | """Returns true if the node is an import statement.""" |
| 315 | return node.type in (syms.import_name, syms.import_from) |
| 316 | |
| 317 | def touch_import(package, name, node): |
| 318 | """ Works like `does_tree_import` but adds an import statement |
| 319 | if it was not imported. """ |
| 320 | def is_import_stmt(node): |
Benjamin Peterson | f37eb3a | 2010-10-14 23:00:04 +0000 | [diff] [blame] | 321 | return (node.type == syms.simple_stmt and node.children and |
| 322 | is_import(node.children[0])) |
Benjamin Peterson | 0b24b3d | 2008-12-16 03:57:54 +0000 | [diff] [blame] | 323 | |
| 324 | root = find_root(node) |
| 325 | |
| 326 | if does_tree_import(package, name, root): |
| 327 | return |
| 328 | |
Benjamin Peterson | 0b24b3d | 2008-12-16 03:57:54 +0000 | [diff] [blame] | 329 | # figure out where to insert the new import. First try to find |
| 330 | # the first import and then skip to the last one. |
| 331 | insert_pos = offset = 0 |
| 332 | for idx, node in enumerate(root.children): |
| 333 | if not is_import_stmt(node): |
| 334 | continue |
| 335 | for offset, node2 in enumerate(root.children[idx:]): |
| 336 | if not is_import_stmt(node2): |
| 337 | break |
| 338 | insert_pos = idx + offset |
| 339 | break |
| 340 | |
| 341 | # if there are no imports where we can insert, find the docstring. |
| 342 | # if that also fails, we stick to the beginning of the file |
| 343 | if insert_pos == 0: |
| 344 | for idx, node in enumerate(root.children): |
Benjamin Peterson | f37eb3a | 2010-10-14 23:00:04 +0000 | [diff] [blame] | 345 | if (node.type == syms.simple_stmt and node.children and |
| 346 | node.children[0].type == token.STRING): |
Benjamin Peterson | 0b24b3d | 2008-12-16 03:57:54 +0000 | [diff] [blame] | 347 | insert_pos = idx + 1 |
Benjamin Peterson | 0b24b3d | 2008-12-16 03:57:54 +0000 | [diff] [blame] | 348 | break |
| 349 | |
| 350 | if package is None: |
| 351 | import_ = Node(syms.import_name, [ |
Benjamin Peterson | 4eb5fa5 | 2010-08-08 19:01:25 +0000 | [diff] [blame] | 352 | Leaf(token.NAME, "import"), |
| 353 | Leaf(token.NAME, name, prefix=" ") |
Benjamin Peterson | 0b24b3d | 2008-12-16 03:57:54 +0000 | [diff] [blame] | 354 | ]) |
| 355 | else: |
Benjamin Peterson | 4eb5fa5 | 2010-08-08 19:01:25 +0000 | [diff] [blame] | 356 | import_ = FromImport(package, [Leaf(token.NAME, name, prefix=" ")]) |
Benjamin Peterson | 0b24b3d | 2008-12-16 03:57:54 +0000 | [diff] [blame] | 357 | |
| 358 | children = [import_, Newline()] |
Benjamin Peterson | 0b24b3d | 2008-12-16 03:57:54 +0000 | [diff] [blame] | 359 | root.insert_child(insert_pos, Node(syms.simple_stmt, children)) |
| 360 | |
| 361 | |
Serhiy Storchaka | db9b65d | 2014-12-13 21:50:49 +0200 | [diff] [blame] | 362 | _def_syms = {syms.classdef, syms.funcdef} |
Martin v. Löwis | ef04c44 | 2008-03-19 05:04:44 +0000 | [diff] [blame] | 363 | def find_binding(name, node, package=None): |
| 364 | """ Returns the node which binds variable name, otherwise None. |
| 365 | If optional argument package is supplied, only imports will |
| 366 | be returned. |
| 367 | See test cases for examples.""" |
| 368 | for child in node.children: |
| 369 | ret = None |
| 370 | if child.type == syms.for_stmt: |
| 371 | if _find(name, child.children[1]): |
| 372 | return child |
| 373 | n = find_binding(name, make_suite(child.children[-1]), package) |
| 374 | if n: ret = n |
| 375 | elif child.type in (syms.if_stmt, syms.while_stmt): |
| 376 | n = find_binding(name, make_suite(child.children[-1]), package) |
| 377 | if n: ret = n |
| 378 | elif child.type == syms.try_stmt: |
| 379 | n = find_binding(name, make_suite(child.children[2]), package) |
| 380 | if n: |
| 381 | ret = n |
| 382 | else: |
| 383 | for i, kid in enumerate(child.children[3:]): |
| 384 | if kid.type == token.COLON and kid.value == ":": |
| 385 | # i+3 is the colon, i+4 is the suite |
| 386 | n = find_binding(name, make_suite(child.children[i+4]), package) |
| 387 | if n: ret = n |
| 388 | elif child.type in _def_syms and child.children[1].value == name: |
| 389 | ret = child |
| 390 | elif _is_import_binding(child, name, package): |
| 391 | ret = child |
| 392 | elif child.type == syms.simple_stmt: |
| 393 | ret = find_binding(name, child, package) |
| 394 | elif child.type == syms.expr_stmt: |
Martin v. Löwis | f733c60 | 2008-03-19 05:26:18 +0000 | [diff] [blame] | 395 | if _find(name, child.children[0]): |
| 396 | ret = child |
Martin v. Löwis | ef04c44 | 2008-03-19 05:04:44 +0000 | [diff] [blame] | 397 | |
| 398 | if ret: |
| 399 | if not package: |
| 400 | return ret |
Benjamin Peterson | 0b24b3d | 2008-12-16 03:57:54 +0000 | [diff] [blame] | 401 | if is_import(ret): |
Martin v. Löwis | ef04c44 | 2008-03-19 05:04:44 +0000 | [diff] [blame] | 402 | return ret |
| 403 | return None |
| 404 | |
Serhiy Storchaka | db9b65d | 2014-12-13 21:50:49 +0200 | [diff] [blame] | 405 | _block_syms = {syms.funcdef, syms.classdef, syms.trailer} |
Martin v. Löwis | ef04c44 | 2008-03-19 05:04:44 +0000 | [diff] [blame] | 406 | def _find(name, node): |
| 407 | nodes = [node] |
| 408 | while nodes: |
| 409 | node = nodes.pop() |
| 410 | if node.type > 256 and node.type not in _block_syms: |
| 411 | nodes.extend(node.children) |
| 412 | elif node.type == token.NAME and node.value == name: |
| 413 | return node |
| 414 | return None |
| 415 | |
| 416 | def _is_import_binding(node, name, package=None): |
| 417 | """ Will reuturn node if node will import name, or node |
| 418 | will import * from package. None is returned otherwise. |
| 419 | See test cases for examples. """ |
| 420 | |
| 421 | if node.type == syms.import_name and not package: |
| 422 | imp = node.children[1] |
| 423 | if imp.type == syms.dotted_as_names: |
| 424 | for child in imp.children: |
| 425 | if child.type == syms.dotted_as_name: |
| 426 | if child.children[2].value == name: |
| 427 | return node |
| 428 | elif child.type == token.NAME and child.value == name: |
| 429 | return node |
| 430 | elif imp.type == syms.dotted_as_name: |
| 431 | last = imp.children[-1] |
| 432 | if last.type == token.NAME and last.value == name: |
| 433 | return node |
| 434 | elif imp.type == token.NAME and imp.value == name: |
| 435 | return node |
| 436 | elif node.type == syms.import_from: |
Benjamin Peterson | df6dc8f | 2008-06-15 02:57:40 +0000 | [diff] [blame] | 437 | # str(...) is used to make life easier here, because |
Martin v. Löwis | ef04c44 | 2008-03-19 05:04:44 +0000 | [diff] [blame] | 438 | # from a.b import parses to ['import', ['a', '.', 'b'], ...] |
Martin v. Löwis | 8a5f8ca | 2008-03-19 05:33:36 +0000 | [diff] [blame] | 439 | if package and str(node.children[1]).strip() != package: |
Martin v. Löwis | ef04c44 | 2008-03-19 05:04:44 +0000 | [diff] [blame] | 440 | return None |
| 441 | n = node.children[3] |
Benjamin Peterson | 4eb5fa5 | 2010-08-08 19:01:25 +0000 | [diff] [blame] | 442 | if package and _find("as", n): |
Martin v. Löwis | ef04c44 | 2008-03-19 05:04:44 +0000 | [diff] [blame] | 443 | # See test_from_import_as for explanation |
| 444 | return None |
| 445 | elif n.type == syms.import_as_names and _find(name, n): |
| 446 | return node |
| 447 | elif n.type == syms.import_as_name: |
| 448 | child = n.children[2] |
| 449 | if child.type == token.NAME and child.value == name: |
| 450 | return node |
| 451 | elif n.type == token.NAME and n.value == name: |
| 452 | return node |
| 453 | elif package and n.type == token.STAR: |
| 454 | return node |
| 455 | return None |