14710c53dcad1ebf3755f3efb9e80ac24bd72a9b2darylm"""Utility functions, node construction macros, etc.""" 24710c53dcad1ebf3755f3efb9e80ac24bd72a9b2darylm# Author: Collin Winter 34710c53dcad1ebf3755f3efb9e80ac24bd72a9b2darylm 44710c53dcad1ebf3755f3efb9e80ac24bd72a9b2darylmfrom itertools import islice 54710c53dcad1ebf3755f3efb9e80ac24bd72a9b2darylm 64710c53dcad1ebf3755f3efb9e80ac24bd72a9b2darylm# Local imports 74710c53dcad1ebf3755f3efb9e80ac24bd72a9b2darylmfrom .pgen2 import token 84710c53dcad1ebf3755f3efb9e80ac24bd72a9b2darylmfrom .pytree import Leaf, Node 94710c53dcad1ebf3755f3efb9e80ac24bd72a9b2darylmfrom .pygram import python_symbols as syms 104710c53dcad1ebf3755f3efb9e80ac24bd72a9b2darylmfrom . import patcomp 114710c53dcad1ebf3755f3efb9e80ac24bd72a9b2darylm 124710c53dcad1ebf3755f3efb9e80ac24bd72a9b2darylm 134710c53dcad1ebf3755f3efb9e80ac24bd72a9b2darylm########################################################### 144710c53dcad1ebf3755f3efb9e80ac24bd72a9b2darylm### Common node-construction "macros" 154710c53dcad1ebf3755f3efb9e80ac24bd72a9b2darylm########################################################### 164710c53dcad1ebf3755f3efb9e80ac24bd72a9b2darylm 174710c53dcad1ebf3755f3efb9e80ac24bd72a9b2darylmdef KeywordArg(keyword, value): 184710c53dcad1ebf3755f3efb9e80ac24bd72a9b2darylm return Node(syms.argument, 194710c53dcad1ebf3755f3efb9e80ac24bd72a9b2darylm [keyword, Leaf(token.EQUAL, u"="), value]) 204710c53dcad1ebf3755f3efb9e80ac24bd72a9b2darylm 214710c53dcad1ebf3755f3efb9e80ac24bd72a9b2darylmdef LParen(): 224710c53dcad1ebf3755f3efb9e80ac24bd72a9b2darylm return Leaf(token.LPAR, u"(") 234710c53dcad1ebf3755f3efb9e80ac24bd72a9b2darylm 244710c53dcad1ebf3755f3efb9e80ac24bd72a9b2darylmdef RParen(): 254710c53dcad1ebf3755f3efb9e80ac24bd72a9b2darylm return Leaf(token.RPAR, u")") 264710c53dcad1ebf3755f3efb9e80ac24bd72a9b2darylm 274710c53dcad1ebf3755f3efb9e80ac24bd72a9b2darylmdef Assign(target, source): 284710c53dcad1ebf3755f3efb9e80ac24bd72a9b2darylm """Build an assignment statement""" 294710c53dcad1ebf3755f3efb9e80ac24bd72a9b2darylm if not isinstance(target, list): 304710c53dcad1ebf3755f3efb9e80ac24bd72a9b2darylm target = [target] 314710c53dcad1ebf3755f3efb9e80ac24bd72a9b2darylm if not isinstance(source, list): 324710c53dcad1ebf3755f3efb9e80ac24bd72a9b2darylm source.prefix = u" " 334710c53dcad1ebf3755f3efb9e80ac24bd72a9b2darylm source = [source] 344710c53dcad1ebf3755f3efb9e80ac24bd72a9b2darylm 354710c53dcad1ebf3755f3efb9e80ac24bd72a9b2darylm return Node(syms.atom, 364710c53dcad1ebf3755f3efb9e80ac24bd72a9b2darylm target + [Leaf(token.EQUAL, u"=", prefix=u" ")] + source) 374710c53dcad1ebf3755f3efb9e80ac24bd72a9b2darylm 384710c53dcad1ebf3755f3efb9e80ac24bd72a9b2darylmdef Name(name, prefix=None): 394710c53dcad1ebf3755f3efb9e80ac24bd72a9b2darylm """Return a NAME leaf""" 404710c53dcad1ebf3755f3efb9e80ac24bd72a9b2darylm return Leaf(token.NAME, name, prefix=prefix) 414710c53dcad1ebf3755f3efb9e80ac24bd72a9b2darylm 424710c53dcad1ebf3755f3efb9e80ac24bd72a9b2darylmdef Attr(obj, attr): 434710c53dcad1ebf3755f3efb9e80ac24bd72a9b2darylm """A node tuple for obj.attr""" 444710c53dcad1ebf3755f3efb9e80ac24bd72a9b2darylm return [obj, Node(syms.trailer, [Dot(), attr])] 454710c53dcad1ebf3755f3efb9e80ac24bd72a9b2darylm 464710c53dcad1ebf3755f3efb9e80ac24bd72a9b2darylmdef Comma(): 474710c53dcad1ebf3755f3efb9e80ac24bd72a9b2darylm """A comma leaf""" 484710c53dcad1ebf3755f3efb9e80ac24bd72a9b2darylm return Leaf(token.COMMA, u",") 494710c53dcad1ebf3755f3efb9e80ac24bd72a9b2darylm 504710c53dcad1ebf3755f3efb9e80ac24bd72a9b2darylmdef Dot(): 514710c53dcad1ebf3755f3efb9e80ac24bd72a9b2darylm """A period (.) leaf""" 524710c53dcad1ebf3755f3efb9e80ac24bd72a9b2darylm return Leaf(token.DOT, u".") 534710c53dcad1ebf3755f3efb9e80ac24bd72a9b2darylm 544710c53dcad1ebf3755f3efb9e80ac24bd72a9b2darylmdef ArgList(args, lparen=LParen(), rparen=RParen()): 554710c53dcad1ebf3755f3efb9e80ac24bd72a9b2darylm """A parenthesised argument list, used by Call()""" 564710c53dcad1ebf3755f3efb9e80ac24bd72a9b2darylm node = Node(syms.trailer, [lparen.clone(), rparen.clone()]) 574710c53dcad1ebf3755f3efb9e80ac24bd72a9b2darylm if args: 584710c53dcad1ebf3755f3efb9e80ac24bd72a9b2darylm node.insert_child(1, Node(syms.arglist, args)) 594710c53dcad1ebf3755f3efb9e80ac24bd72a9b2darylm return node 604710c53dcad1ebf3755f3efb9e80ac24bd72a9b2darylm 614710c53dcad1ebf3755f3efb9e80ac24bd72a9b2darylmdef Call(func_name, args=None, prefix=None): 624710c53dcad1ebf3755f3efb9e80ac24bd72a9b2darylm """A function call""" 634710c53dcad1ebf3755f3efb9e80ac24bd72a9b2darylm node = Node(syms.power, [func_name, ArgList(args)]) 644710c53dcad1ebf3755f3efb9e80ac24bd72a9b2darylm if prefix is not None: 654710c53dcad1ebf3755f3efb9e80ac24bd72a9b2darylm node.prefix = prefix 664710c53dcad1ebf3755f3efb9e80ac24bd72a9b2darylm return node 674710c53dcad1ebf3755f3efb9e80ac24bd72a9b2darylm 684710c53dcad1ebf3755f3efb9e80ac24bd72a9b2darylmdef Newline(): 694710c53dcad1ebf3755f3efb9e80ac24bd72a9b2darylm """A newline literal""" 704710c53dcad1ebf3755f3efb9e80ac24bd72a9b2darylm return Leaf(token.NEWLINE, u"\n") 714710c53dcad1ebf3755f3efb9e80ac24bd72a9b2darylm 724710c53dcad1ebf3755f3efb9e80ac24bd72a9b2darylmdef BlankLine(): 734710c53dcad1ebf3755f3efb9e80ac24bd72a9b2darylm """A blank line""" 744710c53dcad1ebf3755f3efb9e80ac24bd72a9b2darylm return Leaf(token.NEWLINE, u"") 754710c53dcad1ebf3755f3efb9e80ac24bd72a9b2darylm 764710c53dcad1ebf3755f3efb9e80ac24bd72a9b2darylmdef Number(n, prefix=None): 774710c53dcad1ebf3755f3efb9e80ac24bd72a9b2darylm return Leaf(token.NUMBER, n, prefix=prefix) 784710c53dcad1ebf3755f3efb9e80ac24bd72a9b2darylm 794710c53dcad1ebf3755f3efb9e80ac24bd72a9b2darylmdef Subscript(index_node): 804710c53dcad1ebf3755f3efb9e80ac24bd72a9b2darylm """A numeric or string subscript""" 814710c53dcad1ebf3755f3efb9e80ac24bd72a9b2darylm return Node(syms.trailer, [Leaf(token.LBRACE, u"["), 824710c53dcad1ebf3755f3efb9e80ac24bd72a9b2darylm index_node, 834710c53dcad1ebf3755f3efb9e80ac24bd72a9b2darylm Leaf(token.RBRACE, u"]")]) 844710c53dcad1ebf3755f3efb9e80ac24bd72a9b2darylm 854710c53dcad1ebf3755f3efb9e80ac24bd72a9b2darylmdef String(string, prefix=None): 864710c53dcad1ebf3755f3efb9e80ac24bd72a9b2darylm """A string leaf""" 874710c53dcad1ebf3755f3efb9e80ac24bd72a9b2darylm return Leaf(token.STRING, string, prefix=prefix) 884710c53dcad1ebf3755f3efb9e80ac24bd72a9b2darylm 894710c53dcad1ebf3755f3efb9e80ac24bd72a9b2darylmdef ListComp(xp, fp, it, test=None): 904710c53dcad1ebf3755f3efb9e80ac24bd72a9b2darylm """A list comprehension of the form [xp for fp in it if test]. 914710c53dcad1ebf3755f3efb9e80ac24bd72a9b2darylm 924710c53dcad1ebf3755f3efb9e80ac24bd72a9b2darylm If test is None, the "if test" part is omitted. 934710c53dcad1ebf3755f3efb9e80ac24bd72a9b2darylm """ 944710c53dcad1ebf3755f3efb9e80ac24bd72a9b2darylm xp.prefix = u"" 954710c53dcad1ebf3755f3efb9e80ac24bd72a9b2darylm fp.prefix = u" " 964710c53dcad1ebf3755f3efb9e80ac24bd72a9b2darylm it.prefix = u" " 974710c53dcad1ebf3755f3efb9e80ac24bd72a9b2darylm for_leaf = Leaf(token.NAME, u"for") 984710c53dcad1ebf3755f3efb9e80ac24bd72a9b2darylm for_leaf.prefix = u" " 994710c53dcad1ebf3755f3efb9e80ac24bd72a9b2darylm in_leaf = Leaf(token.NAME, u"in") 1004710c53dcad1ebf3755f3efb9e80ac24bd72a9b2darylm in_leaf.prefix = u" " 1014710c53dcad1ebf3755f3efb9e80ac24bd72a9b2darylm inner_args = [for_leaf, fp, in_leaf, it] 1024710c53dcad1ebf3755f3efb9e80ac24bd72a9b2darylm if test: 1034710c53dcad1ebf3755f3efb9e80ac24bd72a9b2darylm test.prefix = u" " 1044710c53dcad1ebf3755f3efb9e80ac24bd72a9b2darylm if_leaf = Leaf(token.NAME, u"if") 1054710c53dcad1ebf3755f3efb9e80ac24bd72a9b2darylm if_leaf.prefix = u" " 1064710c53dcad1ebf3755f3efb9e80ac24bd72a9b2darylm inner_args.append(Node(syms.comp_if, [if_leaf, test])) 1074710c53dcad1ebf3755f3efb9e80ac24bd72a9b2darylm inner = Node(syms.listmaker, [xp, Node(syms.comp_for, inner_args)]) 1084710c53dcad1ebf3755f3efb9e80ac24bd72a9b2darylm return Node(syms.atom, 1094710c53dcad1ebf3755f3efb9e80ac24bd72a9b2darylm [Leaf(token.LBRACE, u"["), 1104710c53dcad1ebf3755f3efb9e80ac24bd72a9b2darylm inner, 1114710c53dcad1ebf3755f3efb9e80ac24bd72a9b2darylm Leaf(token.RBRACE, u"]")]) 1124710c53dcad1ebf3755f3efb9e80ac24bd72a9b2darylm 1134710c53dcad1ebf3755f3efb9e80ac24bd72a9b2darylmdef FromImport(package_name, name_leafs): 1144710c53dcad1ebf3755f3efb9e80ac24bd72a9b2darylm """ Return an import statement in the form: 1154710c53dcad1ebf3755f3efb9e80ac24bd72a9b2darylm from package import name_leafs""" 1164710c53dcad1ebf3755f3efb9e80ac24bd72a9b2darylm # XXX: May not handle dotted imports properly (eg, package_name='foo.bar') 1174710c53dcad1ebf3755f3efb9e80ac24bd72a9b2darylm #assert package_name == '.' or '.' not in package_name, "FromImport has "\ 1184710c53dcad1ebf3755f3efb9e80ac24bd72a9b2darylm # "not been tested with dotted package names -- use at your own "\ 1194710c53dcad1ebf3755f3efb9e80ac24bd72a9b2darylm # "peril!" 1204710c53dcad1ebf3755f3efb9e80ac24bd72a9b2darylm 1214710c53dcad1ebf3755f3efb9e80ac24bd72a9b2darylm for leaf in name_leafs: 1224710c53dcad1ebf3755f3efb9e80ac24bd72a9b2darylm # Pull the leaves out of their old tree 1234710c53dcad1ebf3755f3efb9e80ac24bd72a9b2darylm leaf.remove() 1244710c53dcad1ebf3755f3efb9e80ac24bd72a9b2darylm 1254710c53dcad1ebf3755f3efb9e80ac24bd72a9b2darylm children = [Leaf(token.NAME, u"from"), 1264710c53dcad1ebf3755f3efb9e80ac24bd72a9b2darylm Leaf(token.NAME, package_name, prefix=u" "), 1274710c53dcad1ebf3755f3efb9e80ac24bd72a9b2darylm Leaf(token.NAME, u"import", prefix=u" "), 1284710c53dcad1ebf3755f3efb9e80ac24bd72a9b2darylm Node(syms.import_as_names, name_leafs)] 1294710c53dcad1ebf3755f3efb9e80ac24bd72a9b2darylm imp = Node(syms.import_from, children) 1304710c53dcad1ebf3755f3efb9e80ac24bd72a9b2darylm return imp 1314710c53dcad1ebf3755f3efb9e80ac24bd72a9b2darylm 1324710c53dcad1ebf3755f3efb9e80ac24bd72a9b2darylm 1334710c53dcad1ebf3755f3efb9e80ac24bd72a9b2darylm########################################################### 1344710c53dcad1ebf3755f3efb9e80ac24bd72a9b2darylm### Determine whether a node represents a given literal 1354710c53dcad1ebf3755f3efb9e80ac24bd72a9b2darylm########################################################### 1364710c53dcad1ebf3755f3efb9e80ac24bd72a9b2darylm 1374710c53dcad1ebf3755f3efb9e80ac24bd72a9b2darylmdef is_tuple(node): 1384710c53dcad1ebf3755f3efb9e80ac24bd72a9b2darylm """Does the node represent a tuple literal?""" 1394710c53dcad1ebf3755f3efb9e80ac24bd72a9b2darylm if isinstance(node, Node) and node.children == [LParen(), RParen()]: 1404710c53dcad1ebf3755f3efb9e80ac24bd72a9b2darylm return True 1414710c53dcad1ebf3755f3efb9e80ac24bd72a9b2darylm return (isinstance(node, Node) 1424710c53dcad1ebf3755f3efb9e80ac24bd72a9b2darylm and len(node.children) == 3 1434710c53dcad1ebf3755f3efb9e80ac24bd72a9b2darylm and isinstance(node.children[0], Leaf) 1444710c53dcad1ebf3755f3efb9e80ac24bd72a9b2darylm and isinstance(node.children[1], Node) 1454710c53dcad1ebf3755f3efb9e80ac24bd72a9b2darylm and isinstance(node.children[2], Leaf) 1464710c53dcad1ebf3755f3efb9e80ac24bd72a9b2darylm and node.children[0].value == u"(" 1474710c53dcad1ebf3755f3efb9e80ac24bd72a9b2darylm and node.children[2].value == u")") 1484710c53dcad1ebf3755f3efb9e80ac24bd72a9b2darylm 1494710c53dcad1ebf3755f3efb9e80ac24bd72a9b2darylmdef is_list(node): 1504710c53dcad1ebf3755f3efb9e80ac24bd72a9b2darylm """Does the node represent a list literal?""" 1514710c53dcad1ebf3755f3efb9e80ac24bd72a9b2darylm return (isinstance(node, Node) 1524710c53dcad1ebf3755f3efb9e80ac24bd72a9b2darylm and len(node.children) > 1 1534710c53dcad1ebf3755f3efb9e80ac24bd72a9b2darylm and isinstance(node.children[0], Leaf) 1544710c53dcad1ebf3755f3efb9e80ac24bd72a9b2darylm and isinstance(node.children[-1], Leaf) 1554710c53dcad1ebf3755f3efb9e80ac24bd72a9b2darylm and node.children[0].value == u"[" 1564710c53dcad1ebf3755f3efb9e80ac24bd72a9b2darylm and node.children[-1].value == u"]") 1574710c53dcad1ebf3755f3efb9e80ac24bd72a9b2darylm 1584710c53dcad1ebf3755f3efb9e80ac24bd72a9b2darylm 1594710c53dcad1ebf3755f3efb9e80ac24bd72a9b2darylm########################################################### 1604710c53dcad1ebf3755f3efb9e80ac24bd72a9b2darylm### Misc 1614710c53dcad1ebf3755f3efb9e80ac24bd72a9b2darylm########################################################### 1624710c53dcad1ebf3755f3efb9e80ac24bd72a9b2darylm 1634710c53dcad1ebf3755f3efb9e80ac24bd72a9b2darylmdef parenthesize(node): 1644710c53dcad1ebf3755f3efb9e80ac24bd72a9b2darylm return Node(syms.atom, [LParen(), node, RParen()]) 1654710c53dcad1ebf3755f3efb9e80ac24bd72a9b2darylm 1664710c53dcad1ebf3755f3efb9e80ac24bd72a9b2darylm 1674710c53dcad1ebf3755f3efb9e80ac24bd72a9b2darylmconsuming_calls = set(["sorted", "list", "set", "any", "all", "tuple", "sum", 1684710c53dcad1ebf3755f3efb9e80ac24bd72a9b2darylm "min", "max"]) 1694710c53dcad1ebf3755f3efb9e80ac24bd72a9b2darylm 1704710c53dcad1ebf3755f3efb9e80ac24bd72a9b2darylmdef attr_chain(obj, attr): 1714710c53dcad1ebf3755f3efb9e80ac24bd72a9b2darylm """Follow an attribute chain. 1724710c53dcad1ebf3755f3efb9e80ac24bd72a9b2darylm 1734710c53dcad1ebf3755f3efb9e80ac24bd72a9b2darylm If you have a chain of objects where a.foo -> b, b.foo-> c, etc, 1744710c53dcad1ebf3755f3efb9e80ac24bd72a9b2darylm use this to iterate over all objects in the chain. Iteration is 1754710c53dcad1ebf3755f3efb9e80ac24bd72a9b2darylm terminated by getattr(x, attr) is None. 1764710c53dcad1ebf3755f3efb9e80ac24bd72a9b2darylm 1774710c53dcad1ebf3755f3efb9e80ac24bd72a9b2darylm Args: 1784710c53dcad1ebf3755f3efb9e80ac24bd72a9b2darylm obj: the starting object 1794710c53dcad1ebf3755f3efb9e80ac24bd72a9b2darylm attr: the name of the chaining attribute 1804710c53dcad1ebf3755f3efb9e80ac24bd72a9b2darylm 1814710c53dcad1ebf3755f3efb9e80ac24bd72a9b2darylm Yields: 1824710c53dcad1ebf3755f3efb9e80ac24bd72a9b2darylm Each successive object in the chain. 1834710c53dcad1ebf3755f3efb9e80ac24bd72a9b2darylm """ 1844710c53dcad1ebf3755f3efb9e80ac24bd72a9b2darylm next = getattr(obj, attr) 1854710c53dcad1ebf3755f3efb9e80ac24bd72a9b2darylm while next: 1864710c53dcad1ebf3755f3efb9e80ac24bd72a9b2darylm yield next 1874710c53dcad1ebf3755f3efb9e80ac24bd72a9b2darylm next = getattr(next, attr) 1884710c53dcad1ebf3755f3efb9e80ac24bd72a9b2darylm 1894710c53dcad1ebf3755f3efb9e80ac24bd72a9b2darylmp0 = """for_stmt< 'for' any 'in' node=any ':' any* > 1904710c53dcad1ebf3755f3efb9e80ac24bd72a9b2darylm | comp_for< 'for' any 'in' node=any any* > 1914710c53dcad1ebf3755f3efb9e80ac24bd72a9b2darylm """ 1924710c53dcad1ebf3755f3efb9e80ac24bd72a9b2darylmp1 = """ 1934710c53dcad1ebf3755f3efb9e80ac24bd72a9b2darylmpower< 1944710c53dcad1ebf3755f3efb9e80ac24bd72a9b2darylm ( 'iter' | 'list' | 'tuple' | 'sorted' | 'set' | 'sum' | 1954710c53dcad1ebf3755f3efb9e80ac24bd72a9b2darylm 'any' | 'all' | (any* trailer< '.' 'join' >) ) 1964710c53dcad1ebf3755f3efb9e80ac24bd72a9b2darylm trailer< '(' node=any ')' > 1974710c53dcad1ebf3755f3efb9e80ac24bd72a9b2darylm any* 1984710c53dcad1ebf3755f3efb9e80ac24bd72a9b2darylm> 1994710c53dcad1ebf3755f3efb9e80ac24bd72a9b2darylm""" 2004710c53dcad1ebf3755f3efb9e80ac24bd72a9b2darylmp2 = """ 2014710c53dcad1ebf3755f3efb9e80ac24bd72a9b2darylmpower< 2024710c53dcad1ebf3755f3efb9e80ac24bd72a9b2darylm 'sorted' 2034710c53dcad1ebf3755f3efb9e80ac24bd72a9b2darylm trailer< '(' arglist<node=any any*> ')' > 2044710c53dcad1ebf3755f3efb9e80ac24bd72a9b2darylm any* 2054710c53dcad1ebf3755f3efb9e80ac24bd72a9b2darylm> 2064710c53dcad1ebf3755f3efb9e80ac24bd72a9b2darylm""" 2074710c53dcad1ebf3755f3efb9e80ac24bd72a9b2darylmpats_built = False 2084710c53dcad1ebf3755f3efb9e80ac24bd72a9b2darylmdef in_special_context(node): 2094710c53dcad1ebf3755f3efb9e80ac24bd72a9b2darylm """ Returns true if node is in an environment where all that is required 2104710c53dcad1ebf3755f3efb9e80ac24bd72a9b2darylm of it is being itterable (ie, it doesn't matter if it returns a list 2114710c53dcad1ebf3755f3efb9e80ac24bd72a9b2darylm or an itterator). 2124710c53dcad1ebf3755f3efb9e80ac24bd72a9b2darylm See test_map_nochange in test_fixers.py for some examples and tests. 2134710c53dcad1ebf3755f3efb9e80ac24bd72a9b2darylm """ 2144710c53dcad1ebf3755f3efb9e80ac24bd72a9b2darylm global p0, p1, p2, pats_built 2154710c53dcad1ebf3755f3efb9e80ac24bd72a9b2darylm if not pats_built: 2164710c53dcad1ebf3755f3efb9e80ac24bd72a9b2darylm p1 = patcomp.compile_pattern(p1) 2174710c53dcad1ebf3755f3efb9e80ac24bd72a9b2darylm p0 = patcomp.compile_pattern(p0) 2184710c53dcad1ebf3755f3efb9e80ac24bd72a9b2darylm p2 = patcomp.compile_pattern(p2) 2194710c53dcad1ebf3755f3efb9e80ac24bd72a9b2darylm pats_built = True 2204710c53dcad1ebf3755f3efb9e80ac24bd72a9b2darylm patterns = [p0, p1, p2] 2214710c53dcad1ebf3755f3efb9e80ac24bd72a9b2darylm for pattern, parent in zip(patterns, attr_chain(node, "parent")): 2224710c53dcad1ebf3755f3efb9e80ac24bd72a9b2darylm results = {} 2234710c53dcad1ebf3755f3efb9e80ac24bd72a9b2darylm if pattern.match(parent, results) and results["node"] is node: 2244710c53dcad1ebf3755f3efb9e80ac24bd72a9b2darylm return True 2254710c53dcad1ebf3755f3efb9e80ac24bd72a9b2darylm return False 2264710c53dcad1ebf3755f3efb9e80ac24bd72a9b2darylm 2274710c53dcad1ebf3755f3efb9e80ac24bd72a9b2darylmdef is_probably_builtin(node): 2284710c53dcad1ebf3755f3efb9e80ac24bd72a9b2darylm """ 2294710c53dcad1ebf3755f3efb9e80ac24bd72a9b2darylm Check that something isn't an attribute or function name etc. 2304710c53dcad1ebf3755f3efb9e80ac24bd72a9b2darylm """ 2314710c53dcad1ebf3755f3efb9e80ac24bd72a9b2darylm prev = node.prev_sibling 2324710c53dcad1ebf3755f3efb9e80ac24bd72a9b2darylm if prev is not None and prev.type == token.DOT: 2334710c53dcad1ebf3755f3efb9e80ac24bd72a9b2darylm # Attribute lookup. 2344710c53dcad1ebf3755f3efb9e80ac24bd72a9b2darylm return False 2354710c53dcad1ebf3755f3efb9e80ac24bd72a9b2darylm parent = node.parent 2364710c53dcad1ebf3755f3efb9e80ac24bd72a9b2darylm if parent.type in (syms.funcdef, syms.classdef): 2374710c53dcad1ebf3755f3efb9e80ac24bd72a9b2darylm return False 2384710c53dcad1ebf3755f3efb9e80ac24bd72a9b2darylm if parent.type == syms.expr_stmt and parent.children[0] is node: 2394710c53dcad1ebf3755f3efb9e80ac24bd72a9b2darylm # Assignment. 2404710c53dcad1ebf3755f3efb9e80ac24bd72a9b2darylm return False 2414710c53dcad1ebf3755f3efb9e80ac24bd72a9b2darylm if parent.type == syms.parameters or \ 2424710c53dcad1ebf3755f3efb9e80ac24bd72a9b2darylm (parent.type == syms.typedargslist and ( 2434710c53dcad1ebf3755f3efb9e80ac24bd72a9b2darylm (prev is not None and prev.type == token.COMMA) or 2444710c53dcad1ebf3755f3efb9e80ac24bd72a9b2darylm parent.children[0] is node 2454710c53dcad1ebf3755f3efb9e80ac24bd72a9b2darylm )): 2464710c53dcad1ebf3755f3efb9e80ac24bd72a9b2darylm # The name of an argument. 2474710c53dcad1ebf3755f3efb9e80ac24bd72a9b2darylm return False 2484710c53dcad1ebf3755f3efb9e80ac24bd72a9b2darylm return True 2494710c53dcad1ebf3755f3efb9e80ac24bd72a9b2darylm 2504710c53dcad1ebf3755f3efb9e80ac24bd72a9b2darylmdef find_indentation(node): 2514710c53dcad1ebf3755f3efb9e80ac24bd72a9b2darylm """Find the indentation of *node*.""" 2524710c53dcad1ebf3755f3efb9e80ac24bd72a9b2darylm while node is not None: 2534710c53dcad1ebf3755f3efb9e80ac24bd72a9b2darylm if node.type == syms.suite and len(node.children) > 2: 2544710c53dcad1ebf3755f3efb9e80ac24bd72a9b2darylm indent = node.children[1] 2554710c53dcad1ebf3755f3efb9e80ac24bd72a9b2darylm if indent.type == token.INDENT: 2564710c53dcad1ebf3755f3efb9e80ac24bd72a9b2darylm return indent.value 2574710c53dcad1ebf3755f3efb9e80ac24bd72a9b2darylm node = node.parent 2584710c53dcad1ebf3755f3efb9e80ac24bd72a9b2darylm return u"" 2594710c53dcad1ebf3755f3efb9e80ac24bd72a9b2darylm 2604710c53dcad1ebf3755f3efb9e80ac24bd72a9b2darylm########################################################### 2614710c53dcad1ebf3755f3efb9e80ac24bd72a9b2darylm### The following functions are to find bindings in a suite 2624710c53dcad1ebf3755f3efb9e80ac24bd72a9b2darylm########################################################### 2634710c53dcad1ebf3755f3efb9e80ac24bd72a9b2darylm 2644710c53dcad1ebf3755f3efb9e80ac24bd72a9b2darylmdef make_suite(node): 2654710c53dcad1ebf3755f3efb9e80ac24bd72a9b2darylm if node.type == syms.suite: 2664710c53dcad1ebf3755f3efb9e80ac24bd72a9b2darylm return node 2674710c53dcad1ebf3755f3efb9e80ac24bd72a9b2darylm node = node.clone() 2684710c53dcad1ebf3755f3efb9e80ac24bd72a9b2darylm parent, node.parent = node.parent, None 2694710c53dcad1ebf3755f3efb9e80ac24bd72a9b2darylm suite = Node(syms.suite, [node]) 2704710c53dcad1ebf3755f3efb9e80ac24bd72a9b2darylm suite.parent = parent 2714710c53dcad1ebf3755f3efb9e80ac24bd72a9b2darylm return suite 2724710c53dcad1ebf3755f3efb9e80ac24bd72a9b2darylm 2734710c53dcad1ebf3755f3efb9e80ac24bd72a9b2darylmdef find_root(node): 2744710c53dcad1ebf3755f3efb9e80ac24bd72a9b2darylm """Find the top level namespace.""" 2754710c53dcad1ebf3755f3efb9e80ac24bd72a9b2darylm # Scamper up to the top level namespace 2764710c53dcad1ebf3755f3efb9e80ac24bd72a9b2darylm while node.type != syms.file_input: 2774710c53dcad1ebf3755f3efb9e80ac24bd72a9b2darylm assert node.parent, "Tree is insane! root found before "\ 2784710c53dcad1ebf3755f3efb9e80ac24bd72a9b2darylm "file_input node was found." 2794710c53dcad1ebf3755f3efb9e80ac24bd72a9b2darylm node = node.parent 2804710c53dcad1ebf3755f3efb9e80ac24bd72a9b2darylm return node 2814710c53dcad1ebf3755f3efb9e80ac24bd72a9b2darylm 2824710c53dcad1ebf3755f3efb9e80ac24bd72a9b2darylmdef does_tree_import(package, name, node): 2834710c53dcad1ebf3755f3efb9e80ac24bd72a9b2darylm """ Returns true if name is imported from package at the 2844710c53dcad1ebf3755f3efb9e80ac24bd72a9b2darylm top level of the tree which node belongs to. 2854710c53dcad1ebf3755f3efb9e80ac24bd72a9b2darylm To cover the case of an import like 'import foo', use 2864710c53dcad1ebf3755f3efb9e80ac24bd72a9b2darylm None for the package and 'foo' for the name. """ 2874710c53dcad1ebf3755f3efb9e80ac24bd72a9b2darylm binding = find_binding(name, find_root(node), package) 2884710c53dcad1ebf3755f3efb9e80ac24bd72a9b2darylm return bool(binding) 2894710c53dcad1ebf3755f3efb9e80ac24bd72a9b2darylm 2904710c53dcad1ebf3755f3efb9e80ac24bd72a9b2darylmdef is_import(node): 2914710c53dcad1ebf3755f3efb9e80ac24bd72a9b2darylm """Returns true if the node is an import statement.""" 2924710c53dcad1ebf3755f3efb9e80ac24bd72a9b2darylm return node.type in (syms.import_name, syms.import_from) 2934710c53dcad1ebf3755f3efb9e80ac24bd72a9b2darylm 2944710c53dcad1ebf3755f3efb9e80ac24bd72a9b2darylmdef touch_import(package, name, node): 2954710c53dcad1ebf3755f3efb9e80ac24bd72a9b2darylm """ Works like `does_tree_import` but adds an import statement 2964710c53dcad1ebf3755f3efb9e80ac24bd72a9b2darylm if it was not imported. """ 2974710c53dcad1ebf3755f3efb9e80ac24bd72a9b2darylm def is_import_stmt(node): 2984710c53dcad1ebf3755f3efb9e80ac24bd72a9b2darylm return (node.type == syms.simple_stmt and node.children and 2994710c53dcad1ebf3755f3efb9e80ac24bd72a9b2darylm is_import(node.children[0])) 3004710c53dcad1ebf3755f3efb9e80ac24bd72a9b2darylm 3014710c53dcad1ebf3755f3efb9e80ac24bd72a9b2darylm root = find_root(node) 3024710c53dcad1ebf3755f3efb9e80ac24bd72a9b2darylm 3034710c53dcad1ebf3755f3efb9e80ac24bd72a9b2darylm if does_tree_import(package, name, root): 3044710c53dcad1ebf3755f3efb9e80ac24bd72a9b2darylm return 3054710c53dcad1ebf3755f3efb9e80ac24bd72a9b2darylm 3064710c53dcad1ebf3755f3efb9e80ac24bd72a9b2darylm # figure out where to insert the new import. First try to find 3074710c53dcad1ebf3755f3efb9e80ac24bd72a9b2darylm # the first import and then skip to the last one. 3084710c53dcad1ebf3755f3efb9e80ac24bd72a9b2darylm insert_pos = offset = 0 3094710c53dcad1ebf3755f3efb9e80ac24bd72a9b2darylm for idx, node in enumerate(root.children): 3104710c53dcad1ebf3755f3efb9e80ac24bd72a9b2darylm if not is_import_stmt(node): 3114710c53dcad1ebf3755f3efb9e80ac24bd72a9b2darylm continue 3124710c53dcad1ebf3755f3efb9e80ac24bd72a9b2darylm for offset, node2 in enumerate(root.children[idx:]): 3134710c53dcad1ebf3755f3efb9e80ac24bd72a9b2darylm if not is_import_stmt(node2): 3144710c53dcad1ebf3755f3efb9e80ac24bd72a9b2darylm break 3154710c53dcad1ebf3755f3efb9e80ac24bd72a9b2darylm insert_pos = idx + offset 3164710c53dcad1ebf3755f3efb9e80ac24bd72a9b2darylm break 3174710c53dcad1ebf3755f3efb9e80ac24bd72a9b2darylm 3184710c53dcad1ebf3755f3efb9e80ac24bd72a9b2darylm # if there are no imports where we can insert, find the docstring. 3194710c53dcad1ebf3755f3efb9e80ac24bd72a9b2darylm # if that also fails, we stick to the beginning of the file 3204710c53dcad1ebf3755f3efb9e80ac24bd72a9b2darylm if insert_pos == 0: 3214710c53dcad1ebf3755f3efb9e80ac24bd72a9b2darylm for idx, node in enumerate(root.children): 3224710c53dcad1ebf3755f3efb9e80ac24bd72a9b2darylm if (node.type == syms.simple_stmt and node.children and 3234710c53dcad1ebf3755f3efb9e80ac24bd72a9b2darylm node.children[0].type == token.STRING): 3244710c53dcad1ebf3755f3efb9e80ac24bd72a9b2darylm insert_pos = idx + 1 3254710c53dcad1ebf3755f3efb9e80ac24bd72a9b2darylm break 3264710c53dcad1ebf3755f3efb9e80ac24bd72a9b2darylm 3274710c53dcad1ebf3755f3efb9e80ac24bd72a9b2darylm if package is None: 3284710c53dcad1ebf3755f3efb9e80ac24bd72a9b2darylm import_ = Node(syms.import_name, [ 3294710c53dcad1ebf3755f3efb9e80ac24bd72a9b2darylm Leaf(token.NAME, u"import"), 3304710c53dcad1ebf3755f3efb9e80ac24bd72a9b2darylm Leaf(token.NAME, name, prefix=u" ") 3314710c53dcad1ebf3755f3efb9e80ac24bd72a9b2darylm ]) 3324710c53dcad1ebf3755f3efb9e80ac24bd72a9b2darylm else: 3334710c53dcad1ebf3755f3efb9e80ac24bd72a9b2darylm import_ = FromImport(package, [Leaf(token.NAME, name, prefix=u" ")]) 3344710c53dcad1ebf3755f3efb9e80ac24bd72a9b2darylm 3354710c53dcad1ebf3755f3efb9e80ac24bd72a9b2darylm children = [import_, Newline()] 3364710c53dcad1ebf3755f3efb9e80ac24bd72a9b2darylm root.insert_child(insert_pos, Node(syms.simple_stmt, children)) 3374710c53dcad1ebf3755f3efb9e80ac24bd72a9b2darylm 3384710c53dcad1ebf3755f3efb9e80ac24bd72a9b2darylm 3394710c53dcad1ebf3755f3efb9e80ac24bd72a9b2darylm_def_syms = set([syms.classdef, syms.funcdef]) 3404710c53dcad1ebf3755f3efb9e80ac24bd72a9b2darylmdef find_binding(name, node, package=None): 3414710c53dcad1ebf3755f3efb9e80ac24bd72a9b2darylm """ Returns the node which binds variable name, otherwise None. 3424710c53dcad1ebf3755f3efb9e80ac24bd72a9b2darylm If optional argument package is supplied, only imports will 3434710c53dcad1ebf3755f3efb9e80ac24bd72a9b2darylm be returned. 3444710c53dcad1ebf3755f3efb9e80ac24bd72a9b2darylm See test cases for examples.""" 3454710c53dcad1ebf3755f3efb9e80ac24bd72a9b2darylm for child in node.children: 3464710c53dcad1ebf3755f3efb9e80ac24bd72a9b2darylm ret = None 3474710c53dcad1ebf3755f3efb9e80ac24bd72a9b2darylm if child.type == syms.for_stmt: 3484710c53dcad1ebf3755f3efb9e80ac24bd72a9b2darylm if _find(name, child.children[1]): 3494710c53dcad1ebf3755f3efb9e80ac24bd72a9b2darylm return child 3504710c53dcad1ebf3755f3efb9e80ac24bd72a9b2darylm n = find_binding(name, make_suite(child.children[-1]), package) 3514710c53dcad1ebf3755f3efb9e80ac24bd72a9b2darylm if n: ret = n 3524710c53dcad1ebf3755f3efb9e80ac24bd72a9b2darylm elif child.type in (syms.if_stmt, syms.while_stmt): 3534710c53dcad1ebf3755f3efb9e80ac24bd72a9b2darylm n = find_binding(name, make_suite(child.children[-1]), package) 3544710c53dcad1ebf3755f3efb9e80ac24bd72a9b2darylm if n: ret = n 3554710c53dcad1ebf3755f3efb9e80ac24bd72a9b2darylm elif child.type == syms.try_stmt: 3564710c53dcad1ebf3755f3efb9e80ac24bd72a9b2darylm n = find_binding(name, make_suite(child.children[2]), package) 3574710c53dcad1ebf3755f3efb9e80ac24bd72a9b2darylm if n: 3584710c53dcad1ebf3755f3efb9e80ac24bd72a9b2darylm ret = n 3594710c53dcad1ebf3755f3efb9e80ac24bd72a9b2darylm else: 3604710c53dcad1ebf3755f3efb9e80ac24bd72a9b2darylm for i, kid in enumerate(child.children[3:]): 3614710c53dcad1ebf3755f3efb9e80ac24bd72a9b2darylm if kid.type == token.COLON and kid.value == ":": 3624710c53dcad1ebf3755f3efb9e80ac24bd72a9b2darylm # i+3 is the colon, i+4 is the suite 3634710c53dcad1ebf3755f3efb9e80ac24bd72a9b2darylm n = find_binding(name, make_suite(child.children[i+4]), package) 3644710c53dcad1ebf3755f3efb9e80ac24bd72a9b2darylm if n: ret = n 3654710c53dcad1ebf3755f3efb9e80ac24bd72a9b2darylm elif child.type in _def_syms and child.children[1].value == name: 3664710c53dcad1ebf3755f3efb9e80ac24bd72a9b2darylm ret = child 3674710c53dcad1ebf3755f3efb9e80ac24bd72a9b2darylm elif _is_import_binding(child, name, package): 3684710c53dcad1ebf3755f3efb9e80ac24bd72a9b2darylm ret = child 3694710c53dcad1ebf3755f3efb9e80ac24bd72a9b2darylm elif child.type == syms.simple_stmt: 3704710c53dcad1ebf3755f3efb9e80ac24bd72a9b2darylm ret = find_binding(name, child, package) 3714710c53dcad1ebf3755f3efb9e80ac24bd72a9b2darylm elif child.type == syms.expr_stmt: 3724710c53dcad1ebf3755f3efb9e80ac24bd72a9b2darylm if _find(name, child.children[0]): 3734710c53dcad1ebf3755f3efb9e80ac24bd72a9b2darylm ret = child 3744710c53dcad1ebf3755f3efb9e80ac24bd72a9b2darylm 3754710c53dcad1ebf3755f3efb9e80ac24bd72a9b2darylm if ret: 3764710c53dcad1ebf3755f3efb9e80ac24bd72a9b2darylm if not package: 3774710c53dcad1ebf3755f3efb9e80ac24bd72a9b2darylm return ret 3784710c53dcad1ebf3755f3efb9e80ac24bd72a9b2darylm if is_import(ret): 3794710c53dcad1ebf3755f3efb9e80ac24bd72a9b2darylm return ret 3804710c53dcad1ebf3755f3efb9e80ac24bd72a9b2darylm return None 3814710c53dcad1ebf3755f3efb9e80ac24bd72a9b2darylm 3824710c53dcad1ebf3755f3efb9e80ac24bd72a9b2darylm_block_syms = set([syms.funcdef, syms.classdef, syms.trailer]) 3834710c53dcad1ebf3755f3efb9e80ac24bd72a9b2darylmdef _find(name, node): 3844710c53dcad1ebf3755f3efb9e80ac24bd72a9b2darylm nodes = [node] 3854710c53dcad1ebf3755f3efb9e80ac24bd72a9b2darylm while nodes: 3864710c53dcad1ebf3755f3efb9e80ac24bd72a9b2darylm node = nodes.pop() 3874710c53dcad1ebf3755f3efb9e80ac24bd72a9b2darylm if node.type > 256 and node.type not in _block_syms: 3884710c53dcad1ebf3755f3efb9e80ac24bd72a9b2darylm nodes.extend(node.children) 3894710c53dcad1ebf3755f3efb9e80ac24bd72a9b2darylm elif node.type == token.NAME and node.value == name: 3904710c53dcad1ebf3755f3efb9e80ac24bd72a9b2darylm return node 3914710c53dcad1ebf3755f3efb9e80ac24bd72a9b2darylm return None 3924710c53dcad1ebf3755f3efb9e80ac24bd72a9b2darylm 3934710c53dcad1ebf3755f3efb9e80ac24bd72a9b2darylmdef _is_import_binding(node, name, package=None): 3944710c53dcad1ebf3755f3efb9e80ac24bd72a9b2darylm """ Will reuturn node if node will import name, or node 3954710c53dcad1ebf3755f3efb9e80ac24bd72a9b2darylm will import * from package. None is returned otherwise. 3964710c53dcad1ebf3755f3efb9e80ac24bd72a9b2darylm See test cases for examples. """ 3974710c53dcad1ebf3755f3efb9e80ac24bd72a9b2darylm 3984710c53dcad1ebf3755f3efb9e80ac24bd72a9b2darylm if node.type == syms.import_name and not package: 3994710c53dcad1ebf3755f3efb9e80ac24bd72a9b2darylm imp = node.children[1] 4004710c53dcad1ebf3755f3efb9e80ac24bd72a9b2darylm if imp.type == syms.dotted_as_names: 4014710c53dcad1ebf3755f3efb9e80ac24bd72a9b2darylm for child in imp.children: 4024710c53dcad1ebf3755f3efb9e80ac24bd72a9b2darylm if child.type == syms.dotted_as_name: 4034710c53dcad1ebf3755f3efb9e80ac24bd72a9b2darylm if child.children[2].value == name: 4044710c53dcad1ebf3755f3efb9e80ac24bd72a9b2darylm return node 4054710c53dcad1ebf3755f3efb9e80ac24bd72a9b2darylm elif child.type == token.NAME and child.value == name: 4064710c53dcad1ebf3755f3efb9e80ac24bd72a9b2darylm return node 4074710c53dcad1ebf3755f3efb9e80ac24bd72a9b2darylm elif imp.type == syms.dotted_as_name: 4084710c53dcad1ebf3755f3efb9e80ac24bd72a9b2darylm last = imp.children[-1] 4094710c53dcad1ebf3755f3efb9e80ac24bd72a9b2darylm if last.type == token.NAME and last.value == name: 4104710c53dcad1ebf3755f3efb9e80ac24bd72a9b2darylm return node 4114710c53dcad1ebf3755f3efb9e80ac24bd72a9b2darylm elif imp.type == token.NAME and imp.value == name: 4124710c53dcad1ebf3755f3efb9e80ac24bd72a9b2darylm return node 4134710c53dcad1ebf3755f3efb9e80ac24bd72a9b2darylm elif node.type == syms.import_from: 4144710c53dcad1ebf3755f3efb9e80ac24bd72a9b2darylm # unicode(...) is used to make life easier here, because 4154710c53dcad1ebf3755f3efb9e80ac24bd72a9b2darylm # from a.b import parses to ['import', ['a', '.', 'b'], ...] 4164710c53dcad1ebf3755f3efb9e80ac24bd72a9b2darylm if package and unicode(node.children[1]).strip() != package: 4174710c53dcad1ebf3755f3efb9e80ac24bd72a9b2darylm return None 4184710c53dcad1ebf3755f3efb9e80ac24bd72a9b2darylm n = node.children[3] 4194710c53dcad1ebf3755f3efb9e80ac24bd72a9b2darylm if package and _find(u"as", n): 4204710c53dcad1ebf3755f3efb9e80ac24bd72a9b2darylm # See test_from_import_as for explanation 4214710c53dcad1ebf3755f3efb9e80ac24bd72a9b2darylm return None 4224710c53dcad1ebf3755f3efb9e80ac24bd72a9b2darylm elif n.type == syms.import_as_names and _find(name, n): 4234710c53dcad1ebf3755f3efb9e80ac24bd72a9b2darylm return node 4244710c53dcad1ebf3755f3efb9e80ac24bd72a9b2darylm elif n.type == syms.import_as_name: 4254710c53dcad1ebf3755f3efb9e80ac24bd72a9b2darylm child = n.children[2] 4264710c53dcad1ebf3755f3efb9e80ac24bd72a9b2darylm if child.type == token.NAME and child.value == name: 4274710c53dcad1ebf3755f3efb9e80ac24bd72a9b2darylm return node 4284710c53dcad1ebf3755f3efb9e80ac24bd72a9b2darylm elif n.type == token.NAME and n.value == name: 4294710c53dcad1ebf3755f3efb9e80ac24bd72a9b2darylm return node 4304710c53dcad1ebf3755f3efb9e80ac24bd72a9b2darylm elif package and n.type == token.STAR: 4314710c53dcad1ebf3755f3efb9e80ac24bd72a9b2darylm return node 4324710c53dcad1ebf3755f3efb9e80ac24bd72a9b2darylm return None 433