"""Utility functions, node construction macros, etc."""
from itertools import islice
from .pytree import Leaf, Node
from .pygram import python_symbols as syms
###########################################################
### Common node-construction "macros"
###########################################################
def KeywordArg(keyword, value):
return Node(syms.argument,
[keyword, Leaf(token.EQUAL, u"="), value])
return Leaf(token.LPAR, u"(")
return Leaf(token.RPAR, u")")
def Assign(target, source):
"""Build an assignment statement"""
if not isinstance(target, list):
if not isinstance(source, list):
target + [Leaf(token.EQUAL, u"=", prefix=u" ")] + source)
def Name(name, prefix=None):
return Leaf(token.NAME, name, prefix=prefix)
"""A node tuple for obj.attr"""
return [obj, Node(syms.trailer, [Dot(), attr])]
return Leaf(token.COMMA, u",")
return Leaf(token.DOT, u".")
def ArgList(args, lparen=LParen(), rparen=RParen()):
"""A parenthesised argument list, used by Call()"""
node = Node(syms.trailer, [lparen.clone(), rparen.clone()])
node.insert_child(1, Node(syms.arglist, args))
def Call(func_name, args=None, prefix=None):
node = Node(syms.power, [func_name, ArgList(args)])
return Leaf(token.NEWLINE, u"\n")
return Leaf(token.NEWLINE, u"")
def Number(n, prefix=None):
return Leaf(token.NUMBER, n, prefix=prefix)
def Subscript(index_node):
"""A numeric or string subscript"""
return Node(syms.trailer, [Leaf(token.LBRACE, u"["),
Leaf(token.RBRACE, u"]")])
def String(string, prefix=None):
return Leaf(token.STRING, string, prefix=prefix)
def ListComp(xp, fp, it, test=None):
"""A list comprehension of the form [xp for fp in it if test].
If test is None, the "if test" part is omitted.
for_leaf = Leaf(token.NAME, u"for")
in_leaf = Leaf(token.NAME, u"in")
inner_args = [for_leaf, fp, in_leaf, it]
if_leaf = Leaf(token.NAME, u"if")
inner_args.append(Node(syms.comp_if, [if_leaf, test]))
inner = Node(syms.listmaker, [xp, Node(syms.comp_for, inner_args)])
[Leaf(token.LBRACE, u"["),
Leaf(token.RBRACE, u"]")])
def FromImport(package_name, name_leafs):
""" Return an import statement in the form:
from package import name_leafs"""
# XXX: May not handle dotted imports properly (eg, package_name='foo.bar')
#assert package_name == '.' or '.' not in package_name, "FromImport has "\
# "not been tested with dotted package names -- use at your own "\
# Pull the leaves out of their old tree
children = [Leaf(token.NAME, u"from"),
Leaf(token.NAME, package_name, prefix=u" "),
Leaf(token.NAME, u"import", prefix=u" "),
Node(syms.import_as_names, name_leafs)]
imp = Node(syms.import_from, children)
###########################################################
### Determine whether a node represents a given literal
###########################################################
"""Does the node represent a tuple literal?"""
if isinstance(node, Node) and node.children == [LParen(), RParen()]:
return (isinstance(node, Node)
and len(node.children) == 3
and isinstance(node.children[0], Leaf)
and isinstance(node.children[1], Node)
and isinstance(node.children[2], Leaf)
and node.children[0].value == u"("
and node.children[2].value == u")")
"""Does the node represent a list literal?"""
return (isinstance(node, Node)
and len(node.children) > 1
and isinstance(node.children[0], Leaf)
and isinstance(node.children[-1], Leaf)
and node.children[0].value == u"["
and node.children[-1].value == u"]")
###########################################################
###########################################################
return Node(syms.atom, [LParen(), node, RParen()])
consuming_calls = set(["sorted", "list", "set", "any", "all", "tuple", "sum",
"min", "max", "enumerate"])
def attr_chain(obj, attr):
"""Follow an attribute chain.
If you have a chain of objects where a.foo -> b, b.foo-> c, etc,
use this to iterate over all objects in the chain. Iteration is
terminated by getattr(x, attr) is None.
attr: the name of the chaining attribute
Each successive object in the chain.
next = getattr(obj, attr)
next = getattr(next, attr)
p0 = """for_stmt< 'for' any 'in' node=any ':' any* >
| comp_for< 'for' any 'in' node=any any* >
( 'iter' | 'list' | 'tuple' | 'sorted' | 'set' | 'sum' |
'any' | 'all' | 'enumerate' | (any* trailer< '.' 'join' >) )
trailer< '(' node=any ')' >
( 'sorted' | 'enumerate' )
trailer< '(' arglist<node=any any*> ')' >
def in_special_context(node):
""" Returns true if node is in an environment where all that is required
of it is being iterable (ie, it doesn't matter if it returns a list
See test_map_nochange in test_fixers.py for some examples and tests.
global p0, p1, p2, pats_built
p0 = patcomp.compile_pattern(p0)
p1 = patcomp.compile_pattern(p1)
p2 = patcomp.compile_pattern(p2)
for pattern, parent in zip(patterns, attr_chain(node, "parent")):
if pattern.match(parent, results) and results["node"] is node:
def is_probably_builtin(node):
Check that something isn't an attribute or function name etc.
if prev is not None and prev.type == token.DOT:
if parent.type in (syms.funcdef, syms.classdef):
if parent.type == syms.expr_stmt and parent.children[0] is node:
if parent.type == syms.parameters or \
(parent.type == syms.typedargslist and (
(prev is not None and prev.type == token.COMMA) or
parent.children[0] is node
# The name of an argument.
def find_indentation(node):
"""Find the indentation of *node*."""
if node.type == syms.suite and len(node.children) > 2:
indent = node.children[1]
if indent.type == token.INDENT:
###########################################################
### The following functions are to find bindings in a suite
###########################################################
if node.type == syms.suite:
parent, node.parent = node.parent, None
suite = Node(syms.suite, [node])
"""Find the top level namespace."""
# Scamper up to the top level namespace
while node.type != syms.file_input:
raise ValueError("root found before file_input node was found.")
def does_tree_import(package, name, node):
""" Returns true if name is imported from package at the
top level of the tree which node belongs to.
To cover the case of an import like 'import foo', use
None for the package and 'foo' for the name. """
binding = find_binding(name, find_root(node), package)
"""Returns true if the node is an import statement."""
return node.type in (syms.import_name, syms.import_from)
def touch_import(package, name, node):
""" Works like `does_tree_import` but adds an import statement
if it was not imported. """
def is_import_stmt(node):
return (node.type == syms.simple_stmt and node.children and
is_import(node.children[0]))
if does_tree_import(package, name, root):
# figure out where to insert the new import. First try to find
# the first import and then skip to the last one.
for idx, node in enumerate(root.children):
if not is_import_stmt(node):
for offset, node2 in enumerate(root.children[idx:]):
if not is_import_stmt(node2):
insert_pos = idx + offset
# if there are no imports where we can insert, find the docstring.
# if that also fails, we stick to the beginning of the file
for idx, node in enumerate(root.children):
if (node.type == syms.simple_stmt and node.children and
node.children[0].type == token.STRING):
import_ = Node(syms.import_name, [
Leaf(token.NAME, u"import"),
Leaf(token.NAME, name, prefix=u" ")
import_ = FromImport(package, [Leaf(token.NAME, name, prefix=u" ")])
children = [import_, Newline()]
root.insert_child(insert_pos, Node(syms.simple_stmt, children))
_def_syms = set([syms.classdef, syms.funcdef])
def find_binding(name, node, package=None):
""" Returns the node which binds variable name, otherwise None.
If optional argument package is supplied, only imports will
See test cases for examples."""
for child in node.children:
if child.type == syms.for_stmt:
if _find(name, child.children[1]):
n = find_binding(name, make_suite(child.children[-1]), package)
elif child.type in (syms.if_stmt, syms.while_stmt):
n = find_binding(name, make_suite(child.children[-1]), package)
elif child.type == syms.try_stmt:
n = find_binding(name, make_suite(child.children[2]), package)
for i, kid in enumerate(child.children[3:]):
if kid.type == token.COLON and kid.value == ":":
# i+3 is the colon, i+4 is the suite
n = find_binding(name, make_suite(child.children[i+4]), package)
elif child.type in _def_syms and child.children[1].value == name:
elif _is_import_binding(child, name, package):
elif child.type == syms.simple_stmt:
ret = find_binding(name, child, package)
elif child.type == syms.expr_stmt:
if _find(name, child.children[0]):
_block_syms = set([syms.funcdef, syms.classdef, syms.trailer])
if node.type > 256 and node.type not in _block_syms:
nodes.extend(node.children)
elif node.type == token.NAME and node.value == name:
def _is_import_binding(node, name, package=None):
""" Will reuturn node if node will import name, or node
will import * from package. None is returned otherwise.
See test cases for examples. """
if node.type == syms.import_name and not package:
if imp.type == syms.dotted_as_names:
for child in imp.children:
if child.type == syms.dotted_as_name:
if child.children[2].value == name:
elif child.type == token.NAME and child.value == name:
elif imp.type == syms.dotted_as_name:
if last.type == token.NAME and last.value == name:
elif imp.type == token.NAME and imp.value == name:
elif node.type == syms.import_from:
# unicode(...) is used to make life easier here, because
# from a.b import parses to ['import', ['a', '.', 'b'], ...]
if package and unicode(node.children[1]).strip() != package:
if package and _find(u"as", n):
# See test_from_import_as for explanation
elif n.type == syms.import_as_names and _find(name, n):
elif n.type == syms.import_as_name:
if child.type == token.NAME and child.value == name:
elif n.type == token.NAME and n.value == name:
elif package and n.type == token.STAR: