1ef04c44e29a8276a484f58d03a75a2dec516302dMartin v. Löwis"""Utility functions, node construction macros, etc.""" 2ef04c44e29a8276a484f58d03a75a2dec516302dMartin v. Löwis# Author: Collin Winter 3ef04c44e29a8276a484f58d03a75a2dec516302dMartin v. Löwis 4ef04c44e29a8276a484f58d03a75a2dec516302dMartin v. Löwis# Local imports 5df6dc8f10770f92db68c69d87abe7c89774d128cBenjamin Petersonfrom .pgen2 import token 6df6dc8f10770f92db68c69d87abe7c89774d128cBenjamin Petersonfrom .pytree import Leaf, Node 7df6dc8f10770f92db68c69d87abe7c89774d128cBenjamin Petersonfrom .pygram import python_symbols as syms 8df6dc8f10770f92db68c69d87abe7c89774d128cBenjamin Petersonfrom . import patcomp 9ef04c44e29a8276a484f58d03a75a2dec516302dMartin v. Löwis 10ef04c44e29a8276a484f58d03a75a2dec516302dMartin v. Löwis 11ef04c44e29a8276a484f58d03a75a2dec516302dMartin v. Löwis########################################################### 12ef04c44e29a8276a484f58d03a75a2dec516302dMartin v. Löwis### Common node-construction "macros" 13ef04c44e29a8276a484f58d03a75a2dec516302dMartin v. Löwis########################################################### 14ef04c44e29a8276a484f58d03a75a2dec516302dMartin v. Löwis 15ef04c44e29a8276a484f58d03a75a2dec516302dMartin v. Löwisdef KeywordArg(keyword, value): 16ef04c44e29a8276a484f58d03a75a2dec516302dMartin v. Löwis return Node(syms.argument, 174eb5fa56e4b84a3a76ddc9487f19764dfab52a7fBenjamin Peterson [keyword, Leaf(token.EQUAL, "="), value]) 18ef04c44e29a8276a484f58d03a75a2dec516302dMartin v. Löwis 19ef04c44e29a8276a484f58d03a75a2dec516302dMartin v. Löwisdef LParen(): 20ef04c44e29a8276a484f58d03a75a2dec516302dMartin v. Löwis return Leaf(token.LPAR, "(") 21ef04c44e29a8276a484f58d03a75a2dec516302dMartin v. Löwis 22ef04c44e29a8276a484f58d03a75a2dec516302dMartin v. Löwisdef RParen(): 23ef04c44e29a8276a484f58d03a75a2dec516302dMartin v. Löwis return Leaf(token.RPAR, ")") 24ef04c44e29a8276a484f58d03a75a2dec516302dMartin v. Löwis 25ef04c44e29a8276a484f58d03a75a2dec516302dMartin v. Löwisdef Assign(target, source): 26ef04c44e29a8276a484f58d03a75a2dec516302dMartin v. Löwis """Build an assignment statement""" 27ef04c44e29a8276a484f58d03a75a2dec516302dMartin v. Löwis if not isinstance(target, list): 28ef04c44e29a8276a484f58d03a75a2dec516302dMartin v. Löwis target = [target] 29ef04c44e29a8276a484f58d03a75a2dec516302dMartin v. Löwis if not isinstance(source, list): 302c3ac6b8757f8fc1dacf6aeaa9a9518d8729d52bBenjamin Peterson source.prefix = " " 31ef04c44e29a8276a484f58d03a75a2dec516302dMartin v. Löwis source = [source] 32ef04c44e29a8276a484f58d03a75a2dec516302dMartin v. Löwis 33ef04c44e29a8276a484f58d03a75a2dec516302dMartin v. Löwis return Node(syms.atom, 34ef04c44e29a8276a484f58d03a75a2dec516302dMartin v. Löwis target + [Leaf(token.EQUAL, "=", prefix=" ")] + source) 35ef04c44e29a8276a484f58d03a75a2dec516302dMartin v. Löwis 36ef04c44e29a8276a484f58d03a75a2dec516302dMartin v. Löwisdef Name(name, prefix=None): 37ef04c44e29a8276a484f58d03a75a2dec516302dMartin v. Löwis """Return a NAME leaf""" 38ef04c44e29a8276a484f58d03a75a2dec516302dMartin v. Löwis return Leaf(token.NAME, name, prefix=prefix) 39ef04c44e29a8276a484f58d03a75a2dec516302dMartin v. Löwis 40ef04c44e29a8276a484f58d03a75a2dec516302dMartin v. Löwisdef Attr(obj, attr): 41ef04c44e29a8276a484f58d03a75a2dec516302dMartin v. Löwis """A node tuple for obj.attr""" 42ef04c44e29a8276a484f58d03a75a2dec516302dMartin v. Löwis return [obj, Node(syms.trailer, [Dot(), attr])] 43ef04c44e29a8276a484f58d03a75a2dec516302dMartin v. Löwis 44ef04c44e29a8276a484f58d03a75a2dec516302dMartin v. Löwisdef Comma(): 45ef04c44e29a8276a484f58d03a75a2dec516302dMartin v. Löwis """A comma leaf""" 46ef04c44e29a8276a484f58d03a75a2dec516302dMartin v. Löwis return Leaf(token.COMMA, ",") 47ef04c44e29a8276a484f58d03a75a2dec516302dMartin v. Löwis 48ef04c44e29a8276a484f58d03a75a2dec516302dMartin v. Löwisdef Dot(): 49ef04c44e29a8276a484f58d03a75a2dec516302dMartin v. Löwis """A period (.) leaf""" 50ef04c44e29a8276a484f58d03a75a2dec516302dMartin v. Löwis return Leaf(token.DOT, ".") 51ef04c44e29a8276a484f58d03a75a2dec516302dMartin v. Löwis 52ef04c44e29a8276a484f58d03a75a2dec516302dMartin v. Löwisdef ArgList(args, lparen=LParen(), rparen=RParen()): 53ef04c44e29a8276a484f58d03a75a2dec516302dMartin v. Löwis """A parenthesised argument list, used by Call()""" 54cf60382420f849881c3e1e9d691d9da2d4e95bc7Benjamin Peterson node = Node(syms.trailer, [lparen.clone(), rparen.clone()]) 55cf60382420f849881c3e1e9d691d9da2d4e95bc7Benjamin Peterson if args: 56cf60382420f849881c3e1e9d691d9da2d4e95bc7Benjamin Peterson node.insert_child(1, Node(syms.arglist, args)) 57cf60382420f849881c3e1e9d691d9da2d4e95bc7Benjamin Peterson return node 58ef04c44e29a8276a484f58d03a75a2dec516302dMartin v. Löwis 59cf60382420f849881c3e1e9d691d9da2d4e95bc7Benjamin Petersondef Call(func_name, args=None, prefix=None): 60ef04c44e29a8276a484f58d03a75a2dec516302dMartin v. Löwis """A function call""" 61ef04c44e29a8276a484f58d03a75a2dec516302dMartin v. Löwis node = Node(syms.power, [func_name, ArgList(args)]) 62ef04c44e29a8276a484f58d03a75a2dec516302dMartin v. Löwis if prefix is not None: 632c3ac6b8757f8fc1dacf6aeaa9a9518d8729d52bBenjamin Peterson node.prefix = prefix 64ef04c44e29a8276a484f58d03a75a2dec516302dMartin v. Löwis return node 65ef04c44e29a8276a484f58d03a75a2dec516302dMartin v. Löwis 66ef04c44e29a8276a484f58d03a75a2dec516302dMartin v. Löwisdef Newline(): 67ef04c44e29a8276a484f58d03a75a2dec516302dMartin v. Löwis """A newline literal""" 68ef04c44e29a8276a484f58d03a75a2dec516302dMartin v. Löwis return Leaf(token.NEWLINE, "\n") 69ef04c44e29a8276a484f58d03a75a2dec516302dMartin v. Löwis 70ef04c44e29a8276a484f58d03a75a2dec516302dMartin v. Löwisdef BlankLine(): 71ef04c44e29a8276a484f58d03a75a2dec516302dMartin v. Löwis """A blank line""" 72ef04c44e29a8276a484f58d03a75a2dec516302dMartin v. Löwis return Leaf(token.NEWLINE, "") 73ef04c44e29a8276a484f58d03a75a2dec516302dMartin v. Löwis 74ef04c44e29a8276a484f58d03a75a2dec516302dMartin v. Löwisdef Number(n, prefix=None): 75ef04c44e29a8276a484f58d03a75a2dec516302dMartin v. Löwis return Leaf(token.NUMBER, n, prefix=prefix) 76ef04c44e29a8276a484f58d03a75a2dec516302dMartin v. Löwis 77ef04c44e29a8276a484f58d03a75a2dec516302dMartin v. Löwisdef Subscript(index_node): 78ef04c44e29a8276a484f58d03a75a2dec516302dMartin v. Löwis """A numeric or string subscript""" 794eb5fa56e4b84a3a76ddc9487f19764dfab52a7fBenjamin Peterson return Node(syms.trailer, [Leaf(token.LBRACE, "["), 80ef04c44e29a8276a484f58d03a75a2dec516302dMartin v. Löwis index_node, 814eb5fa56e4b84a3a76ddc9487f19764dfab52a7fBenjamin Peterson Leaf(token.RBRACE, "]")]) 82ef04c44e29a8276a484f58d03a75a2dec516302dMartin v. Löwis 83ef04c44e29a8276a484f58d03a75a2dec516302dMartin v. Löwisdef String(string, prefix=None): 84ef04c44e29a8276a484f58d03a75a2dec516302dMartin v. Löwis """A string leaf""" 85ef04c44e29a8276a484f58d03a75a2dec516302dMartin v. Löwis return Leaf(token.STRING, string, prefix=prefix) 86ef04c44e29a8276a484f58d03a75a2dec516302dMartin v. Löwis 87ef04c44e29a8276a484f58d03a75a2dec516302dMartin v. Löwisdef ListComp(xp, fp, it, test=None): 88ef04c44e29a8276a484f58d03a75a2dec516302dMartin v. Löwis """A list comprehension of the form [xp for fp in it if test]. 89ef04c44e29a8276a484f58d03a75a2dec516302dMartin v. Löwis 90ef04c44e29a8276a484f58d03a75a2dec516302dMartin v. Löwis If test is None, the "if test" part is omitted. 91ef04c44e29a8276a484f58d03a75a2dec516302dMartin v. Löwis """ 922c3ac6b8757f8fc1dacf6aeaa9a9518d8729d52bBenjamin Peterson xp.prefix = "" 932c3ac6b8757f8fc1dacf6aeaa9a9518d8729d52bBenjamin Peterson fp.prefix = " " 942c3ac6b8757f8fc1dacf6aeaa9a9518d8729d52bBenjamin Peterson it.prefix = " " 95ef04c44e29a8276a484f58d03a75a2dec516302dMartin v. Löwis for_leaf = Leaf(token.NAME, "for") 962c3ac6b8757f8fc1dacf6aeaa9a9518d8729d52bBenjamin Peterson for_leaf.prefix = " " 97ef04c44e29a8276a484f58d03a75a2dec516302dMartin v. Löwis in_leaf = Leaf(token.NAME, "in") 982c3ac6b8757f8fc1dacf6aeaa9a9518d8729d52bBenjamin Peterson in_leaf.prefix = " " 99ef04c44e29a8276a484f58d03a75a2dec516302dMartin v. Löwis inner_args = [for_leaf, fp, in_leaf, it] 100ef04c44e29a8276a484f58d03a75a2dec516302dMartin v. Löwis if test: 1012c3ac6b8757f8fc1dacf6aeaa9a9518d8729d52bBenjamin Peterson test.prefix = " " 102ef04c44e29a8276a484f58d03a75a2dec516302dMartin v. Löwis if_leaf = Leaf(token.NAME, "if") 1032c3ac6b8757f8fc1dacf6aeaa9a9518d8729d52bBenjamin Peterson if_leaf.prefix = " " 104ef04c44e29a8276a484f58d03a75a2dec516302dMartin v. Löwis inner_args.append(Node(syms.comp_if, [if_leaf, test])) 105ef04c44e29a8276a484f58d03a75a2dec516302dMartin v. Löwis inner = Node(syms.listmaker, [xp, Node(syms.comp_for, inner_args)]) 106ef04c44e29a8276a484f58d03a75a2dec516302dMartin v. Löwis return Node(syms.atom, 107ef04c44e29a8276a484f58d03a75a2dec516302dMartin v. Löwis [Leaf(token.LBRACE, "["), 108ef04c44e29a8276a484f58d03a75a2dec516302dMartin v. Löwis inner, 109ef04c44e29a8276a484f58d03a75a2dec516302dMartin v. Löwis Leaf(token.RBRACE, "]")]) 110ef04c44e29a8276a484f58d03a75a2dec516302dMartin v. Löwis 111a675ef1141e2533bd5596222d0a64b6a3f74d2d0Martin v. Löwisdef FromImport(package_name, name_leafs): 112a675ef1141e2533bd5596222d0a64b6a3f74d2d0Martin v. Löwis """ Return an import statement in the form: 113a675ef1141e2533bd5596222d0a64b6a3f74d2d0Martin v. Löwis from package import name_leafs""" 114a675ef1141e2533bd5596222d0a64b6a3f74d2d0Martin v. Löwis # XXX: May not handle dotted imports properly (eg, package_name='foo.bar') 115d613bb4c292635ead23b30af0c9bb83abc81d2beBenjamin Peterson #assert package_name == '.' or '.' not in package_name, "FromImport has "\ 116d613bb4c292635ead23b30af0c9bb83abc81d2beBenjamin Peterson # "not been tested with dotted package names -- use at your own "\ 117d613bb4c292635ead23b30af0c9bb83abc81d2beBenjamin Peterson # "peril!" 118a675ef1141e2533bd5596222d0a64b6a3f74d2d0Martin v. Löwis 119a675ef1141e2533bd5596222d0a64b6a3f74d2d0Martin v. Löwis for leaf in name_leafs: 120a675ef1141e2533bd5596222d0a64b6a3f74d2d0Martin v. Löwis # Pull the leaves out of their old tree 121a675ef1141e2533bd5596222d0a64b6a3f74d2d0Martin v. Löwis leaf.remove() 122a675ef1141e2533bd5596222d0a64b6a3f74d2d0Martin v. Löwis 1234eb5fa56e4b84a3a76ddc9487f19764dfab52a7fBenjamin Peterson children = [Leaf(token.NAME, "from"), 124a675ef1141e2533bd5596222d0a64b6a3f74d2d0Martin v. Löwis Leaf(token.NAME, package_name, prefix=" "), 1254eb5fa56e4b84a3a76ddc9487f19764dfab52a7fBenjamin Peterson Leaf(token.NAME, "import", prefix=" "), 126a675ef1141e2533bd5596222d0a64b6a3f74d2d0Martin v. Löwis Node(syms.import_as_names, name_leafs)] 127a675ef1141e2533bd5596222d0a64b6a3f74d2d0Martin v. Löwis imp = Node(syms.import_from, children) 128a675ef1141e2533bd5596222d0a64b6a3f74d2d0Martin v. Löwis return imp 129a675ef1141e2533bd5596222d0a64b6a3f74d2d0Martin v. Löwis 130448e81b2da41e2ad89302cc8cb384795f9f6306cBenjamin Petersondef ImportAndCall(node, results, names): 131448e81b2da41e2ad89302cc8cb384795f9f6306cBenjamin Peterson """Returns an import statement and calls a method 132448e81b2da41e2ad89302cc8cb384795f9f6306cBenjamin Peterson of the module: 133448e81b2da41e2ad89302cc8cb384795f9f6306cBenjamin Peterson 134448e81b2da41e2ad89302cc8cb384795f9f6306cBenjamin Peterson import module 135448e81b2da41e2ad89302cc8cb384795f9f6306cBenjamin Peterson module.name()""" 136448e81b2da41e2ad89302cc8cb384795f9f6306cBenjamin Peterson obj = results["obj"].clone() 137448e81b2da41e2ad89302cc8cb384795f9f6306cBenjamin Peterson if obj.type == syms.arglist: 138448e81b2da41e2ad89302cc8cb384795f9f6306cBenjamin Peterson newarglist = obj.clone() 139448e81b2da41e2ad89302cc8cb384795f9f6306cBenjamin Peterson else: 140448e81b2da41e2ad89302cc8cb384795f9f6306cBenjamin Peterson newarglist = Node(syms.arglist, [obj.clone()]) 141448e81b2da41e2ad89302cc8cb384795f9f6306cBenjamin Peterson after = results["after"] 142448e81b2da41e2ad89302cc8cb384795f9f6306cBenjamin Peterson if after: 143448e81b2da41e2ad89302cc8cb384795f9f6306cBenjamin Peterson after = [n.clone() for n in after] 144448e81b2da41e2ad89302cc8cb384795f9f6306cBenjamin Peterson new = Node(syms.power, 145448e81b2da41e2ad89302cc8cb384795f9f6306cBenjamin Peterson Attr(Name(names[0]), Name(names[1])) + 146448e81b2da41e2ad89302cc8cb384795f9f6306cBenjamin Peterson [Node(syms.trailer, 147448e81b2da41e2ad89302cc8cb384795f9f6306cBenjamin Peterson [results["lpar"].clone(), 148448e81b2da41e2ad89302cc8cb384795f9f6306cBenjamin Peterson newarglist, 149448e81b2da41e2ad89302cc8cb384795f9f6306cBenjamin Peterson results["rpar"].clone()])] + after) 150448e81b2da41e2ad89302cc8cb384795f9f6306cBenjamin Peterson new.prefix = node.prefix 151448e81b2da41e2ad89302cc8cb384795f9f6306cBenjamin Peterson return new 152448e81b2da41e2ad89302cc8cb384795f9f6306cBenjamin Peterson 153a675ef1141e2533bd5596222d0a64b6a3f74d2d0Martin v. Löwis 154ef04c44e29a8276a484f58d03a75a2dec516302dMartin v. Löwis########################################################### 155ef04c44e29a8276a484f58d03a75a2dec516302dMartin v. Löwis### Determine whether a node represents a given literal 156ef04c44e29a8276a484f58d03a75a2dec516302dMartin v. Löwis########################################################### 157ef04c44e29a8276a484f58d03a75a2dec516302dMartin v. Löwis 158ef04c44e29a8276a484f58d03a75a2dec516302dMartin v. Löwisdef is_tuple(node): 159ef04c44e29a8276a484f58d03a75a2dec516302dMartin v. Löwis """Does the node represent a tuple literal?""" 160ef04c44e29a8276a484f58d03a75a2dec516302dMartin v. Löwis if isinstance(node, Node) and node.children == [LParen(), RParen()]: 161ef04c44e29a8276a484f58d03a75a2dec516302dMartin v. Löwis return True 162ef04c44e29a8276a484f58d03a75a2dec516302dMartin v. Löwis return (isinstance(node, Node) 163ef04c44e29a8276a484f58d03a75a2dec516302dMartin v. Löwis and len(node.children) == 3 164ef04c44e29a8276a484f58d03a75a2dec516302dMartin v. Löwis and isinstance(node.children[0], Leaf) 165ef04c44e29a8276a484f58d03a75a2dec516302dMartin v. Löwis and isinstance(node.children[1], Node) 166ef04c44e29a8276a484f58d03a75a2dec516302dMartin v. Löwis and isinstance(node.children[2], Leaf) 167ef04c44e29a8276a484f58d03a75a2dec516302dMartin v. Löwis and node.children[0].value == "(" 168ef04c44e29a8276a484f58d03a75a2dec516302dMartin v. Löwis and node.children[2].value == ")") 169ef04c44e29a8276a484f58d03a75a2dec516302dMartin v. Löwis 170ef04c44e29a8276a484f58d03a75a2dec516302dMartin v. Löwisdef is_list(node): 171ef04c44e29a8276a484f58d03a75a2dec516302dMartin v. Löwis """Does the node represent a list literal?""" 172ef04c44e29a8276a484f58d03a75a2dec516302dMartin v. Löwis return (isinstance(node, Node) 173ef04c44e29a8276a484f58d03a75a2dec516302dMartin v. Löwis and len(node.children) > 1 174ef04c44e29a8276a484f58d03a75a2dec516302dMartin v. Löwis and isinstance(node.children[0], Leaf) 175ef04c44e29a8276a484f58d03a75a2dec516302dMartin v. Löwis and isinstance(node.children[-1], Leaf) 176ef04c44e29a8276a484f58d03a75a2dec516302dMartin v. Löwis and node.children[0].value == "[" 177ef04c44e29a8276a484f58d03a75a2dec516302dMartin v. Löwis and node.children[-1].value == "]") 178ef04c44e29a8276a484f58d03a75a2dec516302dMartin v. Löwis 179ef04c44e29a8276a484f58d03a75a2dec516302dMartin v. Löwis 180ef04c44e29a8276a484f58d03a75a2dec516302dMartin v. Löwis########################################################### 181ef04c44e29a8276a484f58d03a75a2dec516302dMartin v. Löwis### Misc 182ef04c44e29a8276a484f58d03a75a2dec516302dMartin v. Löwis########################################################### 183ef04c44e29a8276a484f58d03a75a2dec516302dMartin v. Löwis 1840b24b3d9511f8f838da4dfb9a59255aff5fcf72eBenjamin Petersondef parenthesize(node): 1850b24b3d9511f8f838da4dfb9a59255aff5fcf72eBenjamin Peterson return Node(syms.atom, [LParen(), node, RParen()]) 1860b24b3d9511f8f838da4dfb9a59255aff5fcf72eBenjamin Peterson 1873de92bf155a1fff6e48b32c5d5f0071f6669ebf0Martin v. Löwis 188db9b65d9e5d9d39199714cfd15a26e46ab0eaae1Serhiy Storchakaconsuming_calls = {"sorted", "list", "set", "any", "all", "tuple", "sum", 189db9b65d9e5d9d39199714cfd15a26e46ab0eaae1Serhiy Storchaka "min", "max", "enumerate"} 1903de92bf155a1fff6e48b32c5d5f0071f6669ebf0Martin v. Löwis 191ef04c44e29a8276a484f58d03a75a2dec516302dMartin v. Löwisdef attr_chain(obj, attr): 192ef04c44e29a8276a484f58d03a75a2dec516302dMartin v. Löwis """Follow an attribute chain. 193f733c60d9aea123a46cd41dbe4dedee7aa2f20f3Martin v. Löwis 194ef04c44e29a8276a484f58d03a75a2dec516302dMartin v. Löwis If you have a chain of objects where a.foo -> b, b.foo-> c, etc, 195ef04c44e29a8276a484f58d03a75a2dec516302dMartin v. Löwis use this to iterate over all objects in the chain. Iteration is 196ef04c44e29a8276a484f58d03a75a2dec516302dMartin v. Löwis terminated by getattr(x, attr) is None. 197f733c60d9aea123a46cd41dbe4dedee7aa2f20f3Martin v. Löwis 198ef04c44e29a8276a484f58d03a75a2dec516302dMartin v. Löwis Args: 199ef04c44e29a8276a484f58d03a75a2dec516302dMartin v. Löwis obj: the starting object 200ef04c44e29a8276a484f58d03a75a2dec516302dMartin v. Löwis attr: the name of the chaining attribute 201f733c60d9aea123a46cd41dbe4dedee7aa2f20f3Martin v. Löwis 202ef04c44e29a8276a484f58d03a75a2dec516302dMartin v. Löwis Yields: 203ef04c44e29a8276a484f58d03a75a2dec516302dMartin v. Löwis Each successive object in the chain. 204ef04c44e29a8276a484f58d03a75a2dec516302dMartin v. Löwis """ 205ef04c44e29a8276a484f58d03a75a2dec516302dMartin v. Löwis next = getattr(obj, attr) 206ef04c44e29a8276a484f58d03a75a2dec516302dMartin v. Löwis while next: 207ef04c44e29a8276a484f58d03a75a2dec516302dMartin v. Löwis yield next 208ef04c44e29a8276a484f58d03a75a2dec516302dMartin v. Löwis next = getattr(next, attr) 209ef04c44e29a8276a484f58d03a75a2dec516302dMartin v. Löwis 210f733c60d9aea123a46cd41dbe4dedee7aa2f20f3Martin v. Löwisp0 = """for_stmt< 'for' any 'in' node=any ':' any* > 211f733c60d9aea123a46cd41dbe4dedee7aa2f20f3Martin v. Löwis | comp_for< 'for' any 'in' node=any any* > 212f733c60d9aea123a46cd41dbe4dedee7aa2f20f3Martin v. Löwis """ 213f733c60d9aea123a46cd41dbe4dedee7aa2f20f3Martin v. Löwisp1 = """ 214f733c60d9aea123a46cd41dbe4dedee7aa2f20f3Martin v. Löwispower< 215f733c60d9aea123a46cd41dbe4dedee7aa2f20f3Martin v. Löwis ( 'iter' | 'list' | 'tuple' | 'sorted' | 'set' | 'sum' | 21657af38700a42f7c037eab85f107800d50d641198Benjamin Peterson 'any' | 'all' | 'enumerate' | (any* trailer< '.' 'join' >) ) 217f733c60d9aea123a46cd41dbe4dedee7aa2f20f3Martin v. Löwis trailer< '(' node=any ')' > 218f733c60d9aea123a46cd41dbe4dedee7aa2f20f3Martin v. Löwis any* 219f733c60d9aea123a46cd41dbe4dedee7aa2f20f3Martin v. Löwis> 220f733c60d9aea123a46cd41dbe4dedee7aa2f20f3Martin v. Löwis""" 221f733c60d9aea123a46cd41dbe4dedee7aa2f20f3Martin v. Löwisp2 = """ 222f733c60d9aea123a46cd41dbe4dedee7aa2f20f3Martin v. Löwispower< 22357af38700a42f7c037eab85f107800d50d641198Benjamin Peterson ( 'sorted' | 'enumerate' ) 224f733c60d9aea123a46cd41dbe4dedee7aa2f20f3Martin v. Löwis trailer< '(' arglist<node=any any*> ')' > 225f733c60d9aea123a46cd41dbe4dedee7aa2f20f3Martin v. Löwis any* 226f733c60d9aea123a46cd41dbe4dedee7aa2f20f3Martin v. Löwis> 227f733c60d9aea123a46cd41dbe4dedee7aa2f20f3Martin v. Löwis""" 228f733c60d9aea123a46cd41dbe4dedee7aa2f20f3Martin v. Löwispats_built = False 229f733c60d9aea123a46cd41dbe4dedee7aa2f20f3Martin v. Löwisdef in_special_context(node): 230f733c60d9aea123a46cd41dbe4dedee7aa2f20f3Martin v. Löwis """ Returns true if node is in an environment where all that is required 23157af38700a42f7c037eab85f107800d50d641198Benjamin Peterson of it is being iterable (ie, it doesn't matter if it returns a list 23257af38700a42f7c037eab85f107800d50d641198Benjamin Peterson or an iterator). 233f733c60d9aea123a46cd41dbe4dedee7aa2f20f3Martin v. Löwis See test_map_nochange in test_fixers.py for some examples and tests. 234f733c60d9aea123a46cd41dbe4dedee7aa2f20f3Martin v. Löwis """ 235f733c60d9aea123a46cd41dbe4dedee7aa2f20f3Martin v. Löwis global p0, p1, p2, pats_built 236f733c60d9aea123a46cd41dbe4dedee7aa2f20f3Martin v. Löwis if not pats_built: 237f733c60d9aea123a46cd41dbe4dedee7aa2f20f3Martin v. Löwis p0 = patcomp.compile_pattern(p0) 23857af38700a42f7c037eab85f107800d50d641198Benjamin Peterson p1 = patcomp.compile_pattern(p1) 239f733c60d9aea123a46cd41dbe4dedee7aa2f20f3Martin v. Löwis p2 = patcomp.compile_pattern(p2) 240f733c60d9aea123a46cd41dbe4dedee7aa2f20f3Martin v. Löwis pats_built = True 241f733c60d9aea123a46cd41dbe4dedee7aa2f20f3Martin v. Löwis patterns = [p0, p1, p2] 242f733c60d9aea123a46cd41dbe4dedee7aa2f20f3Martin v. Löwis for pattern, parent in zip(patterns, attr_chain(node, "parent")): 243f733c60d9aea123a46cd41dbe4dedee7aa2f20f3Martin v. Löwis results = {} 244f733c60d9aea123a46cd41dbe4dedee7aa2f20f3Martin v. Löwis if pattern.match(parent, results) and results["node"] is node: 245f733c60d9aea123a46cd41dbe4dedee7aa2f20f3Martin v. Löwis return True 246f733c60d9aea123a46cd41dbe4dedee7aa2f20f3Martin v. Löwis return False 247f733c60d9aea123a46cd41dbe4dedee7aa2f20f3Martin v. Löwis 2488bcddcabd770dd424b97d7c667ef8e5337436215Benjamin Petersondef is_probably_builtin(node): 2498bcddcabd770dd424b97d7c667ef8e5337436215Benjamin Peterson """ 2508bcddcabd770dd424b97d7c667ef8e5337436215Benjamin Peterson Check that something isn't an attribute or function name etc. 2518bcddcabd770dd424b97d7c667ef8e5337436215Benjamin Peterson """ 252608d8bcdfc1762cde31efdf02a6379a06d7964b1Benjamin Peterson prev = node.prev_sibling 2538bcddcabd770dd424b97d7c667ef8e5337436215Benjamin Peterson if prev is not None and prev.type == token.DOT: 2548bcddcabd770dd424b97d7c667ef8e5337436215Benjamin Peterson # Attribute lookup. 2558bcddcabd770dd424b97d7c667ef8e5337436215Benjamin Peterson return False 2568bcddcabd770dd424b97d7c667ef8e5337436215Benjamin Peterson parent = node.parent 2578bcddcabd770dd424b97d7c667ef8e5337436215Benjamin Peterson if parent.type in (syms.funcdef, syms.classdef): 2588bcddcabd770dd424b97d7c667ef8e5337436215Benjamin Peterson return False 2598bcddcabd770dd424b97d7c667ef8e5337436215Benjamin Peterson if parent.type == syms.expr_stmt and parent.children[0] is node: 2608bcddcabd770dd424b97d7c667ef8e5337436215Benjamin Peterson # Assignment. 2618bcddcabd770dd424b97d7c667ef8e5337436215Benjamin Peterson return False 2628bcddcabd770dd424b97d7c667ef8e5337436215Benjamin Peterson if parent.type == syms.parameters or \ 2638bcddcabd770dd424b97d7c667ef8e5337436215Benjamin Peterson (parent.type == syms.typedargslist and ( 2648bcddcabd770dd424b97d7c667ef8e5337436215Benjamin Peterson (prev is not None and prev.type == token.COMMA) or 2658bcddcabd770dd424b97d7c667ef8e5337436215Benjamin Peterson parent.children[0] is node 2668bcddcabd770dd424b97d7c667ef8e5337436215Benjamin Peterson )): 2678bcddcabd770dd424b97d7c667ef8e5337436215Benjamin Peterson # The name of an argument. 2688bcddcabd770dd424b97d7c667ef8e5337436215Benjamin Peterson return False 2698bcddcabd770dd424b97d7c667ef8e5337436215Benjamin Peterson return True 2708bcddcabd770dd424b97d7c667ef8e5337436215Benjamin Peterson 2714eb5fa56e4b84a3a76ddc9487f19764dfab52a7fBenjamin Petersondef find_indentation(node): 2724eb5fa56e4b84a3a76ddc9487f19764dfab52a7fBenjamin Peterson """Find the indentation of *node*.""" 2734eb5fa56e4b84a3a76ddc9487f19764dfab52a7fBenjamin Peterson while node is not None: 2744eb5fa56e4b84a3a76ddc9487f19764dfab52a7fBenjamin Peterson if node.type == syms.suite and len(node.children) > 2: 2754eb5fa56e4b84a3a76ddc9487f19764dfab52a7fBenjamin Peterson indent = node.children[1] 2764eb5fa56e4b84a3a76ddc9487f19764dfab52a7fBenjamin Peterson if indent.type == token.INDENT: 2774eb5fa56e4b84a3a76ddc9487f19764dfab52a7fBenjamin Peterson return indent.value 2784eb5fa56e4b84a3a76ddc9487f19764dfab52a7fBenjamin Peterson node = node.parent 2794eb5fa56e4b84a3a76ddc9487f19764dfab52a7fBenjamin Peterson return "" 2804eb5fa56e4b84a3a76ddc9487f19764dfab52a7fBenjamin Peterson 281ef04c44e29a8276a484f58d03a75a2dec516302dMartin v. Löwis########################################################### 282ef04c44e29a8276a484f58d03a75a2dec516302dMartin v. Löwis### The following functions are to find bindings in a suite 283ef04c44e29a8276a484f58d03a75a2dec516302dMartin v. Löwis########################################################### 284ef04c44e29a8276a484f58d03a75a2dec516302dMartin v. Löwis 285ef04c44e29a8276a484f58d03a75a2dec516302dMartin v. Löwisdef make_suite(node): 286ef04c44e29a8276a484f58d03a75a2dec516302dMartin v. Löwis if node.type == syms.suite: 287ef04c44e29a8276a484f58d03a75a2dec516302dMartin v. Löwis return node 288ef04c44e29a8276a484f58d03a75a2dec516302dMartin v. Löwis node = node.clone() 289ef04c44e29a8276a484f58d03a75a2dec516302dMartin v. Löwis parent, node.parent = node.parent, None 290ef04c44e29a8276a484f58d03a75a2dec516302dMartin v. Löwis suite = Node(syms.suite, [node]) 291ef04c44e29a8276a484f58d03a75a2dec516302dMartin v. Löwis suite.parent = parent 292ef04c44e29a8276a484f58d03a75a2dec516302dMartin v. Löwis return suite 293ef04c44e29a8276a484f58d03a75a2dec516302dMartin v. Löwis 2940b24b3d9511f8f838da4dfb9a59255aff5fcf72eBenjamin Petersondef find_root(node): 2950b24b3d9511f8f838da4dfb9a59255aff5fcf72eBenjamin Peterson """Find the top level namespace.""" 296ef04c44e29a8276a484f58d03a75a2dec516302dMartin v. Löwis # Scamper up to the top level namespace 297ef04c44e29a8276a484f58d03a75a2dec516302dMartin v. Löwis while node.type != syms.file_input: 298ef04c44e29a8276a484f58d03a75a2dec516302dMartin v. Löwis node = node.parent 2991654d74e9a6250491bfb923424c4cf83af2af7a6Benjamin Peterson if not node: 3001654d74e9a6250491bfb923424c4cf83af2af7a6Benjamin Peterson raise ValueError("root found before file_input node was found.") 3010b24b3d9511f8f838da4dfb9a59255aff5fcf72eBenjamin Peterson return node 302ef04c44e29a8276a484f58d03a75a2dec516302dMartin v. Löwis 3030b24b3d9511f8f838da4dfb9a59255aff5fcf72eBenjamin Petersondef does_tree_import(package, name, node): 3040b24b3d9511f8f838da4dfb9a59255aff5fcf72eBenjamin Peterson """ Returns true if name is imported from package at the 3050b24b3d9511f8f838da4dfb9a59255aff5fcf72eBenjamin Peterson top level of the tree which node belongs to. 3060b24b3d9511f8f838da4dfb9a59255aff5fcf72eBenjamin Peterson To cover the case of an import like 'import foo', use 3070b24b3d9511f8f838da4dfb9a59255aff5fcf72eBenjamin Peterson None for the package and 'foo' for the name. """ 3080b24b3d9511f8f838da4dfb9a59255aff5fcf72eBenjamin Peterson binding = find_binding(name, find_root(node), package) 309ef04c44e29a8276a484f58d03a75a2dec516302dMartin v. Löwis return bool(binding) 310ef04c44e29a8276a484f58d03a75a2dec516302dMartin v. Löwis 3110b24b3d9511f8f838da4dfb9a59255aff5fcf72eBenjamin Petersondef is_import(node): 3120b24b3d9511f8f838da4dfb9a59255aff5fcf72eBenjamin Peterson """Returns true if the node is an import statement.""" 3130b24b3d9511f8f838da4dfb9a59255aff5fcf72eBenjamin Peterson return node.type in (syms.import_name, syms.import_from) 3140b24b3d9511f8f838da4dfb9a59255aff5fcf72eBenjamin Peterson 3150b24b3d9511f8f838da4dfb9a59255aff5fcf72eBenjamin Petersondef touch_import(package, name, node): 3160b24b3d9511f8f838da4dfb9a59255aff5fcf72eBenjamin Peterson """ Works like `does_tree_import` but adds an import statement 3170b24b3d9511f8f838da4dfb9a59255aff5fcf72eBenjamin Peterson if it was not imported. """ 3180b24b3d9511f8f838da4dfb9a59255aff5fcf72eBenjamin Peterson def is_import_stmt(node): 319f37eb3a184c344a5ee7355a1c1a9acc6d4835c6eBenjamin Peterson return (node.type == syms.simple_stmt and node.children and 320f37eb3a184c344a5ee7355a1c1a9acc6d4835c6eBenjamin Peterson is_import(node.children[0])) 3210b24b3d9511f8f838da4dfb9a59255aff5fcf72eBenjamin Peterson 3220b24b3d9511f8f838da4dfb9a59255aff5fcf72eBenjamin Peterson root = find_root(node) 3230b24b3d9511f8f838da4dfb9a59255aff5fcf72eBenjamin Peterson 3240b24b3d9511f8f838da4dfb9a59255aff5fcf72eBenjamin Peterson if does_tree_import(package, name, root): 3250b24b3d9511f8f838da4dfb9a59255aff5fcf72eBenjamin Peterson return 3260b24b3d9511f8f838da4dfb9a59255aff5fcf72eBenjamin Peterson 3270b24b3d9511f8f838da4dfb9a59255aff5fcf72eBenjamin Peterson # figure out where to insert the new import. First try to find 3280b24b3d9511f8f838da4dfb9a59255aff5fcf72eBenjamin Peterson # the first import and then skip to the last one. 3290b24b3d9511f8f838da4dfb9a59255aff5fcf72eBenjamin Peterson insert_pos = offset = 0 3300b24b3d9511f8f838da4dfb9a59255aff5fcf72eBenjamin Peterson for idx, node in enumerate(root.children): 3310b24b3d9511f8f838da4dfb9a59255aff5fcf72eBenjamin Peterson if not is_import_stmt(node): 3320b24b3d9511f8f838da4dfb9a59255aff5fcf72eBenjamin Peterson continue 3330b24b3d9511f8f838da4dfb9a59255aff5fcf72eBenjamin Peterson for offset, node2 in enumerate(root.children[idx:]): 3340b24b3d9511f8f838da4dfb9a59255aff5fcf72eBenjamin Peterson if not is_import_stmt(node2): 3350b24b3d9511f8f838da4dfb9a59255aff5fcf72eBenjamin Peterson break 3360b24b3d9511f8f838da4dfb9a59255aff5fcf72eBenjamin Peterson insert_pos = idx + offset 3370b24b3d9511f8f838da4dfb9a59255aff5fcf72eBenjamin Peterson break 3380b24b3d9511f8f838da4dfb9a59255aff5fcf72eBenjamin Peterson 3390b24b3d9511f8f838da4dfb9a59255aff5fcf72eBenjamin Peterson # if there are no imports where we can insert, find the docstring. 3400b24b3d9511f8f838da4dfb9a59255aff5fcf72eBenjamin Peterson # if that also fails, we stick to the beginning of the file 3410b24b3d9511f8f838da4dfb9a59255aff5fcf72eBenjamin Peterson if insert_pos == 0: 3420b24b3d9511f8f838da4dfb9a59255aff5fcf72eBenjamin Peterson for idx, node in enumerate(root.children): 343f37eb3a184c344a5ee7355a1c1a9acc6d4835c6eBenjamin Peterson if (node.type == syms.simple_stmt and node.children and 344f37eb3a184c344a5ee7355a1c1a9acc6d4835c6eBenjamin Peterson node.children[0].type == token.STRING): 3450b24b3d9511f8f838da4dfb9a59255aff5fcf72eBenjamin Peterson insert_pos = idx + 1 3460b24b3d9511f8f838da4dfb9a59255aff5fcf72eBenjamin Peterson break 3470b24b3d9511f8f838da4dfb9a59255aff5fcf72eBenjamin Peterson 3480b24b3d9511f8f838da4dfb9a59255aff5fcf72eBenjamin Peterson if package is None: 3490b24b3d9511f8f838da4dfb9a59255aff5fcf72eBenjamin Peterson import_ = Node(syms.import_name, [ 3504eb5fa56e4b84a3a76ddc9487f19764dfab52a7fBenjamin Peterson Leaf(token.NAME, "import"), 3514eb5fa56e4b84a3a76ddc9487f19764dfab52a7fBenjamin Peterson Leaf(token.NAME, name, prefix=" ") 3520b24b3d9511f8f838da4dfb9a59255aff5fcf72eBenjamin Peterson ]) 3530b24b3d9511f8f838da4dfb9a59255aff5fcf72eBenjamin Peterson else: 3544eb5fa56e4b84a3a76ddc9487f19764dfab52a7fBenjamin Peterson import_ = FromImport(package, [Leaf(token.NAME, name, prefix=" ")]) 3550b24b3d9511f8f838da4dfb9a59255aff5fcf72eBenjamin Peterson 3560b24b3d9511f8f838da4dfb9a59255aff5fcf72eBenjamin Peterson children = [import_, Newline()] 3570b24b3d9511f8f838da4dfb9a59255aff5fcf72eBenjamin Peterson root.insert_child(insert_pos, Node(syms.simple_stmt, children)) 3580b24b3d9511f8f838da4dfb9a59255aff5fcf72eBenjamin Peterson 3590b24b3d9511f8f838da4dfb9a59255aff5fcf72eBenjamin Peterson 360db9b65d9e5d9d39199714cfd15a26e46ab0eaae1Serhiy Storchaka_def_syms = {syms.classdef, syms.funcdef} 361ef04c44e29a8276a484f58d03a75a2dec516302dMartin v. Löwisdef find_binding(name, node, package=None): 362ef04c44e29a8276a484f58d03a75a2dec516302dMartin v. Löwis """ Returns the node which binds variable name, otherwise None. 363ef04c44e29a8276a484f58d03a75a2dec516302dMartin v. Löwis If optional argument package is supplied, only imports will 364ef04c44e29a8276a484f58d03a75a2dec516302dMartin v. Löwis be returned. 365ef04c44e29a8276a484f58d03a75a2dec516302dMartin v. Löwis See test cases for examples.""" 366ef04c44e29a8276a484f58d03a75a2dec516302dMartin v. Löwis for child in node.children: 367ef04c44e29a8276a484f58d03a75a2dec516302dMartin v. Löwis ret = None 368ef04c44e29a8276a484f58d03a75a2dec516302dMartin v. Löwis if child.type == syms.for_stmt: 369ef04c44e29a8276a484f58d03a75a2dec516302dMartin v. Löwis if _find(name, child.children[1]): 370ef04c44e29a8276a484f58d03a75a2dec516302dMartin v. Löwis return child 371ef04c44e29a8276a484f58d03a75a2dec516302dMartin v. Löwis n = find_binding(name, make_suite(child.children[-1]), package) 372ef04c44e29a8276a484f58d03a75a2dec516302dMartin v. Löwis if n: ret = n 373ef04c44e29a8276a484f58d03a75a2dec516302dMartin v. Löwis elif child.type in (syms.if_stmt, syms.while_stmt): 374ef04c44e29a8276a484f58d03a75a2dec516302dMartin v. Löwis n = find_binding(name, make_suite(child.children[-1]), package) 375ef04c44e29a8276a484f58d03a75a2dec516302dMartin v. Löwis if n: ret = n 376ef04c44e29a8276a484f58d03a75a2dec516302dMartin v. Löwis elif child.type == syms.try_stmt: 377ef04c44e29a8276a484f58d03a75a2dec516302dMartin v. Löwis n = find_binding(name, make_suite(child.children[2]), package) 378ef04c44e29a8276a484f58d03a75a2dec516302dMartin v. Löwis if n: 379ef04c44e29a8276a484f58d03a75a2dec516302dMartin v. Löwis ret = n 380ef04c44e29a8276a484f58d03a75a2dec516302dMartin v. Löwis else: 381ef04c44e29a8276a484f58d03a75a2dec516302dMartin v. Löwis for i, kid in enumerate(child.children[3:]): 382ef04c44e29a8276a484f58d03a75a2dec516302dMartin v. Löwis if kid.type == token.COLON and kid.value == ":": 383ef04c44e29a8276a484f58d03a75a2dec516302dMartin v. Löwis # i+3 is the colon, i+4 is the suite 384ef04c44e29a8276a484f58d03a75a2dec516302dMartin v. Löwis n = find_binding(name, make_suite(child.children[i+4]), package) 385ef04c44e29a8276a484f58d03a75a2dec516302dMartin v. Löwis if n: ret = n 386ef04c44e29a8276a484f58d03a75a2dec516302dMartin v. Löwis elif child.type in _def_syms and child.children[1].value == name: 387ef04c44e29a8276a484f58d03a75a2dec516302dMartin v. Löwis ret = child 388ef04c44e29a8276a484f58d03a75a2dec516302dMartin v. Löwis elif _is_import_binding(child, name, package): 389ef04c44e29a8276a484f58d03a75a2dec516302dMartin v. Löwis ret = child 390ef04c44e29a8276a484f58d03a75a2dec516302dMartin v. Löwis elif child.type == syms.simple_stmt: 391ef04c44e29a8276a484f58d03a75a2dec516302dMartin v. Löwis ret = find_binding(name, child, package) 392ef04c44e29a8276a484f58d03a75a2dec516302dMartin v. Löwis elif child.type == syms.expr_stmt: 393f733c60d9aea123a46cd41dbe4dedee7aa2f20f3Martin v. Löwis if _find(name, child.children[0]): 394f733c60d9aea123a46cd41dbe4dedee7aa2f20f3Martin v. Löwis ret = child 395ef04c44e29a8276a484f58d03a75a2dec516302dMartin v. Löwis 396ef04c44e29a8276a484f58d03a75a2dec516302dMartin v. Löwis if ret: 397ef04c44e29a8276a484f58d03a75a2dec516302dMartin v. Löwis if not package: 398ef04c44e29a8276a484f58d03a75a2dec516302dMartin v. Löwis return ret 3990b24b3d9511f8f838da4dfb9a59255aff5fcf72eBenjamin Peterson if is_import(ret): 400ef04c44e29a8276a484f58d03a75a2dec516302dMartin v. Löwis return ret 401ef04c44e29a8276a484f58d03a75a2dec516302dMartin v. Löwis return None 402ef04c44e29a8276a484f58d03a75a2dec516302dMartin v. Löwis 403db9b65d9e5d9d39199714cfd15a26e46ab0eaae1Serhiy Storchaka_block_syms = {syms.funcdef, syms.classdef, syms.trailer} 404ef04c44e29a8276a484f58d03a75a2dec516302dMartin v. Löwisdef _find(name, node): 405ef04c44e29a8276a484f58d03a75a2dec516302dMartin v. Löwis nodes = [node] 406ef04c44e29a8276a484f58d03a75a2dec516302dMartin v. Löwis while nodes: 407ef04c44e29a8276a484f58d03a75a2dec516302dMartin v. Löwis node = nodes.pop() 408ef04c44e29a8276a484f58d03a75a2dec516302dMartin v. Löwis if node.type > 256 and node.type not in _block_syms: 409ef04c44e29a8276a484f58d03a75a2dec516302dMartin v. Löwis nodes.extend(node.children) 410ef04c44e29a8276a484f58d03a75a2dec516302dMartin v. Löwis elif node.type == token.NAME and node.value == name: 411ef04c44e29a8276a484f58d03a75a2dec516302dMartin v. Löwis return node 412ef04c44e29a8276a484f58d03a75a2dec516302dMartin v. Löwis return None 413ef04c44e29a8276a484f58d03a75a2dec516302dMartin v. Löwis 414ef04c44e29a8276a484f58d03a75a2dec516302dMartin v. Löwisdef _is_import_binding(node, name, package=None): 415ef04c44e29a8276a484f58d03a75a2dec516302dMartin v. Löwis """ Will reuturn node if node will import name, or node 416ef04c44e29a8276a484f58d03a75a2dec516302dMartin v. Löwis will import * from package. None is returned otherwise. 417ef04c44e29a8276a484f58d03a75a2dec516302dMartin v. Löwis See test cases for examples. """ 418ef04c44e29a8276a484f58d03a75a2dec516302dMartin v. Löwis 419ef04c44e29a8276a484f58d03a75a2dec516302dMartin v. Löwis if node.type == syms.import_name and not package: 420ef04c44e29a8276a484f58d03a75a2dec516302dMartin v. Löwis imp = node.children[1] 421ef04c44e29a8276a484f58d03a75a2dec516302dMartin v. Löwis if imp.type == syms.dotted_as_names: 422ef04c44e29a8276a484f58d03a75a2dec516302dMartin v. Löwis for child in imp.children: 423ef04c44e29a8276a484f58d03a75a2dec516302dMartin v. Löwis if child.type == syms.dotted_as_name: 424ef04c44e29a8276a484f58d03a75a2dec516302dMartin v. Löwis if child.children[2].value == name: 425ef04c44e29a8276a484f58d03a75a2dec516302dMartin v. Löwis return node 426ef04c44e29a8276a484f58d03a75a2dec516302dMartin v. Löwis elif child.type == token.NAME and child.value == name: 427ef04c44e29a8276a484f58d03a75a2dec516302dMartin v. Löwis return node 428ef04c44e29a8276a484f58d03a75a2dec516302dMartin v. Löwis elif imp.type == syms.dotted_as_name: 429ef04c44e29a8276a484f58d03a75a2dec516302dMartin v. Löwis last = imp.children[-1] 430ef04c44e29a8276a484f58d03a75a2dec516302dMartin v. Löwis if last.type == token.NAME and last.value == name: 431ef04c44e29a8276a484f58d03a75a2dec516302dMartin v. Löwis return node 432ef04c44e29a8276a484f58d03a75a2dec516302dMartin v. Löwis elif imp.type == token.NAME and imp.value == name: 433ef04c44e29a8276a484f58d03a75a2dec516302dMartin v. Löwis return node 434ef04c44e29a8276a484f58d03a75a2dec516302dMartin v. Löwis elif node.type == syms.import_from: 435df6dc8f10770f92db68c69d87abe7c89774d128cBenjamin Peterson # str(...) is used to make life easier here, because 436ef04c44e29a8276a484f58d03a75a2dec516302dMartin v. Löwis # from a.b import parses to ['import', ['a', '.', 'b'], ...] 4378a5f8ca33b56db9af973d1f34a9b3df5271b56c0Martin v. Löwis if package and str(node.children[1]).strip() != package: 438ef04c44e29a8276a484f58d03a75a2dec516302dMartin v. Löwis return None 439ef04c44e29a8276a484f58d03a75a2dec516302dMartin v. Löwis n = node.children[3] 4404eb5fa56e4b84a3a76ddc9487f19764dfab52a7fBenjamin Peterson if package and _find("as", n): 441ef04c44e29a8276a484f58d03a75a2dec516302dMartin v. Löwis # See test_from_import_as for explanation 442ef04c44e29a8276a484f58d03a75a2dec516302dMartin v. Löwis return None 443ef04c44e29a8276a484f58d03a75a2dec516302dMartin v. Löwis elif n.type == syms.import_as_names and _find(name, n): 444ef04c44e29a8276a484f58d03a75a2dec516302dMartin v. Löwis return node 445ef04c44e29a8276a484f58d03a75a2dec516302dMartin v. Löwis elif n.type == syms.import_as_name: 446ef04c44e29a8276a484f58d03a75a2dec516302dMartin v. Löwis child = n.children[2] 447ef04c44e29a8276a484f58d03a75a2dec516302dMartin v. Löwis if child.type == token.NAME and child.value == name: 448ef04c44e29a8276a484f58d03a75a2dec516302dMartin v. Löwis return node 449ef04c44e29a8276a484f58d03a75a2dec516302dMartin v. Löwis elif n.type == token.NAME and n.value == name: 450ef04c44e29a8276a484f58d03a75a2dec516302dMartin v. Löwis return node 451ef04c44e29a8276a484f58d03a75a2dec516302dMartin v. Löwis elif package and n.type == token.STAR: 452ef04c44e29a8276a484f58d03a75a2dec516302dMartin v. Löwis return node 453ef04c44e29a8276a484f58d03a75a2dec516302dMartin v. Löwis return None 454