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