1
"""Utility functions, node construction macros, etc."""
2
# Author: Collin Winter
5
from .pgen2 import token
6
from .pytree import Leaf, Node
7
from .pygram import python_symbols as syms
11
###########################################################
12
### Common node-construction "macros"
13
###########################################################
15
def KeywordArg(keyword, value):
16
return Node(syms.argument,
17
[keyword, Leaf(token.EQUAL, '='), value])
20
return Leaf(token.LPAR, "(")
23
return Leaf(token.RPAR, ")")
25
def Assign(target, source):
26
"""Build an assignment statement"""
27
if not isinstance(target, list):
29
if not isinstance(source, list):
30
source.set_prefix(" ")
33
return Node(syms.atom,
34
target + [Leaf(token.EQUAL, "=", prefix=" ")] + source)
36
def Name(name, prefix=None):
37
"""Return a NAME leaf"""
38
return Leaf(token.NAME, name, prefix=prefix)
41
"""A node tuple for obj.attr"""
42
return [obj, Node(syms.trailer, [Dot(), attr])]
46
return Leaf(token.COMMA, ",")
49
"""A period (.) leaf"""
50
return Leaf(token.DOT, ".")
52
def ArgList(args, lparen=LParen(), rparen=RParen()):
53
"""A parenthesised argument list, used by Call()"""
54
node = Node(syms.trailer, [lparen.clone(), rparen.clone()])
56
node.insert_child(1, Node(syms.arglist, args))
59
def Call(func_name, args=None, prefix=None):
61
node = Node(syms.power, [func_name, ArgList(args)])
62
if prefix is not None:
63
node.set_prefix(prefix)
67
"""A newline literal"""
68
return Leaf(token.NEWLINE, "\n")
72
return Leaf(token.NEWLINE, "")
74
def Number(n, prefix=None):
75
return Leaf(token.NUMBER, n, prefix=prefix)
77
def Subscript(index_node):
78
"""A numeric or string subscript"""
79
return Node(syms.trailer, [Leaf(token.LBRACE, '['),
81
Leaf(token.RBRACE, ']')])
83
def String(string, prefix=None):
85
return Leaf(token.STRING, string, prefix=prefix)
87
def ListComp(xp, fp, it, test=None):
88
"""A list comprehension of the form [xp for fp in it if test].
90
If test is None, the "if test" part is omitted.
95
for_leaf = Leaf(token.NAME, "for")
96
for_leaf.set_prefix(" ")
97
in_leaf = Leaf(token.NAME, "in")
98
in_leaf.set_prefix(" ")
99
inner_args = [for_leaf, fp, in_leaf, it]
102
if_leaf = Leaf(token.NAME, "if")
103
if_leaf.set_prefix(" ")
104
inner_args.append(Node(syms.comp_if, [if_leaf, test]))
105
inner = Node(syms.listmaker, [xp, Node(syms.comp_for, inner_args)])
106
return Node(syms.atom,
107
[Leaf(token.LBRACE, "["),
109
Leaf(token.RBRACE, "]")])
111
def FromImport(package_name, name_leafs):
112
""" Return an import statement in the form:
113
from package import name_leafs"""
114
# XXX: May not handle dotted imports properly (eg, package_name='foo.bar')
115
#assert package_name == '.' or '.' not in package_name, "FromImport has "\
116
# "not been tested with dotted package names -- use at your own "\
119
for leaf in name_leafs:
120
# Pull the leaves out of their old tree
123
children = [Leaf(token.NAME, 'from'),
124
Leaf(token.NAME, package_name, prefix=" "),
125
Leaf(token.NAME, 'import', prefix=" "),
126
Node(syms.import_as_names, name_leafs)]
127
imp = Node(syms.import_from, children)
131
###########################################################
132
### Determine whether a node represents a given literal
133
###########################################################
136
"""Does the node represent a tuple literal?"""
137
if isinstance(node, Node) and node.children == [LParen(), RParen()]:
139
return (isinstance(node, Node)
140
and len(node.children) == 3
141
and isinstance(node.children[0], Leaf)
142
and isinstance(node.children[1], Node)
143
and isinstance(node.children[2], Leaf)
144
and node.children[0].value == "("
145
and node.children[2].value == ")")
148
"""Does the node represent a list literal?"""
149
return (isinstance(node, Node)
150
and len(node.children) > 1
151
and isinstance(node.children[0], Leaf)
152
and isinstance(node.children[-1], Leaf)
153
and node.children[0].value == "["
154
and node.children[-1].value == "]")
156
###########################################################
157
### Common portability code. This allows fixers to do, eg,
158
### "from .util import set" and forget about it.
159
###########################################################
173
from sets import Set as set
181
###########################################################
183
###########################################################
186
consuming_calls = set(["sorted", "list", "set", "any", "all", "tuple", "sum",
189
def attr_chain(obj, attr):
190
"""Follow an attribute chain.
192
If you have a chain of objects where a.foo -> b, b.foo-> c, etc,
193
use this to iterate over all objects in the chain. Iteration is
194
terminated by getattr(x, attr) is None.
197
obj: the starting object
198
attr: the name of the chaining attribute
201
Each successive object in the chain.
203
next = getattr(obj, attr)
206
next = getattr(next, attr)
208
p0 = """for_stmt< 'for' any 'in' node=any ':' any* >
209
| comp_for< 'for' any 'in' node=any any* >
213
( 'iter' | 'list' | 'tuple' | 'sorted' | 'set' | 'sum' |
214
'any' | 'all' | (any* trailer< '.' 'join' >) )
215
trailer< '(' node=any ')' >
222
trailer< '(' arglist<node=any any*> ')' >
227
def in_special_context(node):
228
""" Returns true if node is in an environment where all that is required
229
of it is being itterable (ie, it doesn't matter if it returns a list
231
See test_map_nochange in test_fixers.py for some examples and tests.
233
global p0, p1, p2, pats_built
235
p1 = patcomp.compile_pattern(p1)
236
p0 = patcomp.compile_pattern(p0)
237
p2 = patcomp.compile_pattern(p2)
239
patterns = [p0, p1, p2]
240
for pattern, parent in zip(patterns, attr_chain(node, "parent")):
242
if pattern.match(parent, results) and results["node"] is node:
246
###########################################################
247
### The following functions are to find bindings in a suite
248
###########################################################
250
def make_suite(node):
251
if node.type == syms.suite:
254
parent, node.parent = node.parent, None
255
suite = Node(syms.suite, [node])
256
suite.parent = parent
259
def does_tree_import(package, name, node):
260
""" Returns true if name is imported from package at the
261
top level of the tree which node belongs to.
262
To cover the case of an import like 'import foo', use
263
Null for the package and 'foo' for the name. """
264
# Scamper up to the top level namespace
265
while node.type != syms.file_input:
266
assert node.parent, "Tree is insane! root found before "\
267
"file_input node was found."
270
binding = find_binding(name, node, package)
273
_def_syms = set([syms.classdef, syms.funcdef])
274
def find_binding(name, node, package=None):
275
""" Returns the node which binds variable name, otherwise None.
276
If optional argument package is supplied, only imports will
278
See test cases for examples."""
279
for child in node.children:
281
if child.type == syms.for_stmt:
282
if _find(name, child.children[1]):
284
n = find_binding(name, make_suite(child.children[-1]), package)
286
elif child.type in (syms.if_stmt, syms.while_stmt):
287
n = find_binding(name, make_suite(child.children[-1]), package)
289
elif child.type == syms.try_stmt:
290
n = find_binding(name, make_suite(child.children[2]), package)
294
for i, kid in enumerate(child.children[3:]):
295
if kid.type == token.COLON and kid.value == ":":
296
# i+3 is the colon, i+4 is the suite
297
n = find_binding(name, make_suite(child.children[i+4]), package)
299
elif child.type in _def_syms and child.children[1].value == name:
301
elif _is_import_binding(child, name, package):
303
elif child.type == syms.simple_stmt:
304
ret = find_binding(name, child, package)
305
elif child.type == syms.expr_stmt:
306
if _find(name, child.children[0]):
312
if ret.type in (syms.import_name, syms.import_from):
316
_block_syms = set([syms.funcdef, syms.classdef, syms.trailer])
317
def _find(name, node):
321
if node.type > 256 and node.type not in _block_syms:
322
nodes.extend(node.children)
323
elif node.type == token.NAME and node.value == name:
327
def _is_import_binding(node, name, package=None):
328
""" Will reuturn node if node will import name, or node
329
will import * from package. None is returned otherwise.
330
See test cases for examples. """
332
if node.type == syms.import_name and not package:
333
imp = node.children[1]
334
if imp.type == syms.dotted_as_names:
335
for child in imp.children:
336
if child.type == syms.dotted_as_name:
337
if child.children[2].value == name:
339
elif child.type == token.NAME and child.value == name:
341
elif imp.type == syms.dotted_as_name:
342
last = imp.children[-1]
343
if last.type == token.NAME and last.value == name:
345
elif imp.type == token.NAME and imp.value == name:
347
elif node.type == syms.import_from:
348
# unicode(...) is used to make life easier here, because
349
# from a.b import parses to ['import', ['a', '.', 'b'], ...]
350
if package and unicode(node.children[1]).strip() != package:
353
if package and _find('as', n):
354
# See test_from_import_as for explanation
356
elif n.type == syms.import_as_names and _find(name, n):
358
elif n.type == syms.import_as_name:
359
child = n.children[2]
360
if child.type == token.NAME and child.value == name:
362
elif n.type == token.NAME and n.value == name:
364
elif package and n.type == token.STAR: