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