def get_for_args_stmts(self, iter_name, args_list):
'''
Returns 3 gast stmt nodes for argument.
1. Initailize of iterate variable
2. Condition for the loop
3. Statement for changing of iterate variable during the loop
NOTE(TODO): Python allows to access iteration variable after loop, such
as "for i in range(10)" will create i = 9 after the loop. But using
current conversion will make i = 10. We should find a way to change it
'''
len_range_args = len(args_list)
assert len_range_args >= 1 and len_range_args <= 3, "range() function takes 1 to 3 arguments"
if len_range_args == 1:
init_stmt = get_constant_variable_node(iter_name, 0)
else:
init_stmt = gast.Assign(
targets=[
gast.Name(
id=iter_name,
ctx=gast.Store(),
annotation=None,
type_comment=None)
],
value=args_list[0])
range_max_node = args_list[0] if len_range_args == 1 else args_list[1]
step_node = args_list[2] if len_range_args == 3 else gast.Constant(
value=1, kind=None)
cond_stmt = gast.Compare(
left=gast.BinOp(
left=gast.Name(
id=iter_name,
ctx=gast.Load(),
annotation=None,
type_comment=None),
op=gast.Add(),
right=step_node),
ops=[gast.LtE()],
comparators=[range_max_node])
change_stmt = gast.AugAssign(
target=gast.Name(
id=iter_name,
ctx=gast.Store(),
annotation=None,
type_comment=None),
op=gast.Add(),
value=step_node)
return init_stmt, cond_stmt, change_stmt
def test_ast_to_object(self):
node = gast.FunctionDef(
name='f',
args=gast.arguments(
args=[gast.Name('a', gast.Param(), None)],
vararg=None,
kwonlyargs=[],
kwarg=None,
defaults=[],
kw_defaults=[]),
body=[
gast.Return(
gast.BinOp(
op=gast.Add(),
left=gast.Name('a', gast.Load(), None),
right=gast.Num(1)))
],
decorator_list=[],
returns=None)
module, source, _ = compiler.ast_to_object(node)
expected_source = """
# coding=utf-8
def f(a):
return a + 1
"""
self.assertEqual(
textwrap.dedent(expected_source).strip(),
source.strip())
self.assertEqual(2, module.f(1))
with open(module.__file__, 'r') as temp_output:
self.assertEqual(
textwrap.dedent(expected_source).strip(),
temp_output.read().strip())
def _build_enum_increase_node(self):
return gast.AugAssign(target=gast.Name(id=self.enum_idx_name,
ctx=gast.Store(),
annotation=None,
type_comment=None),
op=gast.Add(),
value=gast.Constant(value=1, kind=None))
def test_replace_code_block(self):
template = """
def test_fn(a):
block
return a
"""
class ShouldBeReplaced(object):
pass
node = templates.replace(
template,
block=[
gast.Assign(
[
gast.Name('a',
ctx=ShouldBeReplaced,
annotation=None,
type_comment=None)
],
gast.BinOp(
gast.Name('a',
ctx=ShouldBeReplaced,
annotation=None,
type_comment=None), gast.Add(),
gast.Constant(1, kind=None)),
),
] * 2)[0]
result, _, _ = loader.load_ast(node)
self.assertEqual(3, result.test_fn(1))
def to_ast(self):
assert self._finalized
if self._argspec:
result = self._argspec[0]
for i in range(1, len(self._argspec)):
result = gast.BinOp(result, gast.Add(), self._argspec[i])
return result
return gast.Tuple([], None)
def _build_index_increase_node(self, step_node):
return gast.AugAssign(target=gast.Name(
id=self.iter_var_name
if self.is_for_range_iter() else self.iter_idx_name,
ctx=gast.Store(),
annotation=None,
type_comment=None),
op=gast.Add(),
value=step_node)
class StrJoinPattern(Pattern):
# a + "..." + b => "...".join((a, b))
pattern = ast.BinOp(left=ast.BinOp(left=Placeholder(0),
op=ast.Add(),
right=ast.Constant(
Placeholder(1, str), None)),
op=ast.Add(),
right=Placeholder(2))
@staticmethod
def sub():
return ast.Call(
func=ast.Attribute(
ast.Attribute(ast.Name('__builtin__', ast.Load(), None, None),
'str', ast.Load()), 'join', ast.Load()),
args=[
ast.Constant(Placeholder(1), None),
ast.Tuple([Placeholder(0), Placeholder(2)], ast.Load())
],
keywords=[])
def get_for_args_stmts(self, iter_name, args_list):
'''
Returns 3 gast stmt nodes for argument.
1. Initailize of iterate variable
2. Condition for the loop
3. Statement for changing of iterate variable during the loop
'''
len_range_args = len(args_list)
assert len_range_args >= 1 and len_range_args <= 3, "range() function takes 1 to 3 arguments"
if len_range_args == 1:
init_stmt = get_constant_variable_node(iter_name, 0)
else:
init_stmt = gast.Assign(targets=[
gast.Name(id=iter_name,
ctx=gast.Store(),
annotation=None,
type_comment=None)
],
value=args_list[0])
range_max_node = args_list[0] if len_range_args == 1 else args_list[1]
step_node = args_list[2] if len_range_args == 3 else gast.Constant(
value=1, kind=None)
old_cond_stmt = gast.Compare(left=gast.BinOp(left=gast.Name(
id=iter_name, ctx=gast.Load(), annotation=None, type_comment=None),
op=gast.Add(),
right=step_node),
ops=[gast.LtE()],
comparators=[range_max_node])
cond_stmt = gast.BoolOp(op=gast.And(),
values=[old_cond_stmt, self.condition_node])
change_stmt = gast.AugAssign(target=gast.Name(id=iter_name,
ctx=gast.Store(),
annotation=None,
type_comment=None),
op=gast.Add(),
value=step_node)
return init_stmt, cond_stmt, change_stmt
def visit_Subscript(self, node):
"""
>>> import gast as ast
>>> from pythran import passmanager, backend
>>> pm = passmanager.PassManager("test")
>>> node = ast.parse("def foo(a): a[1:][3]")
>>> _, node = pm.apply(PartialConstantFolding, node)
>>> _, node = pm.apply(ConstantFolding, node)
>>> print(pm.dump(backend.Python, node))
def foo(a):
a[4]
>>> node = ast.parse("def foo(a): a[::2][3]")
>>> _, node = pm.apply(PartialConstantFolding, node)
>>> _, node = pm.apply(ConstantFolding, node)
>>> print(pm.dump(backend.Python, node))
def foo(a):
a[6]
>>> node = ast.parse("def foo(a): a[-4:][5]")
>>> _, node = pm.apply(PartialConstantFolding, node)
>>> _, node = pm.apply(ConstantFolding, node)
>>> print(pm.dump(backend.Python, node))
def foo(a):
a[1]
"""
self.generic_visit(node)
if not isinstance(node.value, ast.Subscript):
return node
if not isinstance(node.value.slice, ast.Slice):
return node
if not isinstance(node.slice, ast.Index):
return node
if not isnum(node.slice.value):
return node
slice_ = node.value.slice
index = node.slice
node = node.value
node.slice = index
lower = slice_.lower or ast.Constant(0, None)
step = slice_.step or ast.Constant(1, None)
node.slice.value = ast.BinOp(lower,
ast.Add(),
ast.BinOp(index.value,
ast.Mult(),
step))
self.update = True
return node
def _build_cond_stmt(self, step_node, compare_node):
return gast.Compare(
left=gast.BinOp(
left=gast.Name(
id=self.iter_var_name
if self.is_for_range_iter() else self.iter_idx_name,
ctx=gast.Load(),
annotation=None,
type_comment=None),
op=gast.Add(),
right=step_node),
ops=[gast.LtE()],
comparators=[compare_node])
def test_code_block(self):
def template(block): # pylint:disable=unused-argument
def test_fn(a): # pylint:disable=unused-variable
block # pylint:disable=pointless-statement
return a
node = templates.replace(
template,
block=[
gast.Assign([gast.Name('a', gast.Store(), None)],
gast.BinOp(gast.Name('a', gast.Load(), None),
gast.Add(), gast.Num(1))),
] * 2)[0]
result = compiler.ast_to_object(node)
self.assertEquals(3, result.test_fn(1))
def test_replace_code_block(self):
template = """
def test_fn(a):
block
return a
"""
node = templates.replace(
template,
block=[
gast.Assign([gast.Name('a', None, None)],
gast.BinOp(gast.Name('a', None, None), gast.Add(),
gast.Num(1))),
] * 2)[0]
result, _ = compiler.ast_to_object(node)
self.assertEquals(3, result.test_fn(1))
def test_load_ast(self):
node = gast.FunctionDef(
name='f',
args=gast.arguments(
args=[
gast.Name(
'a', ctx=gast.Param(), annotation=None, type_comment=None)
],
posonlyargs=[],
vararg=None,
kwonlyargs=[],
kw_defaults=[],
kwarg=None,
defaults=[]),
body=[
gast.Return(
gast.BinOp(
op=gast.Add(),
left=gast.Name(
'a',
ctx=gast.Load(),
annotation=None,
type_comment=None),
right=gast.Constant(1, kind=None)))
],
decorator_list=[],
returns=None,
type_comment=None)
module, source, _ = loader.load_ast(node)
expected_source = """
# coding=utf-8
def f(a):
return (a + 1)
"""
self.assertEqual(
textwrap.dedent(expected_source).strip(),
source.strip())
self.assertEqual(2, module.f(1))
with open(module.__file__, 'r') as temp_output:
self.assertEqual(
textwrap.dedent(expected_source).strip(),
temp_output.read().strip())
args=[
ast.BinOp(left=Placeholder(0),
op=ast.Sub(),
right=ast.Num(n=1)),
ast.Num(n=-1),
ast.Num(n=-1)
],
keywords=[])),
# X * X => X ** 2
(ast.BinOp(left=Placeholder(0), op=ast.Mult(), right=Placeholder(0)),
lambda: ast.BinOp(left=Placeholder(0), op=ast.Pow(), right=ast.Num(n=2))),
# a + "..." + b => "...".join((a, b))
(ast.BinOp(left=ast.BinOp(left=Placeholder(0),
op=ast.Add(),
right=ast.Str(Placeholder(1))),
op=ast.Add(),
right=Placeholder(2)),
lambda: ast.Call(func=ast.Attribute(
ast.Attribute(ast.Name('__builtin__', ast.Load(), None), 'str',
ast.Load()), 'join', ast.Load()),
args=[
ast.Str(Placeholder(1)),
ast.Tuple([Placeholder(0),
Placeholder(2)], ast.Load())
],
keywords=[])),
]
def visit_BinOp(self, node):
if not isinstance(node.op, ast.Mod):
return self.generic_visit(node)
# check that right is a name defined once outside of loop
# TODO: handle expression instead of names
if not isinstance(node.right, ast.Name):
return self.generic_visit(node)
right_def = self.single_def(node.right)
if not right_def:
return self.generic_visit(node)
if self.range_values[node.right.id].low < 0:
return self.generic_visit(node)
# same for lhs
if not isinstance(node.left, ast.Name):
return self.generic_visit(node)
head = self.single_def(node.left)
if not head:
return self.generic_visit(node)
# check lhs is the actual index of a loop
head = head['name']
loop = self.ancestors[head][-1]
if not isinstance(loop, ast.For):
return self.generic_visit(node)
if not isinstance(loop.iter, ast.Call):
return self.generic_visit(node)
# make sure rhs is defined out of the loop
if loop in self.ancestors[right_def['name']]:
return self.generic_visit(node)
# gather range informations
range_ = None
for alias in self.aliases[loop.iter.func]:
if alias is MODULES['__builtin__']['range']:
range_ = alias
elif alias is MODULES['__builtin__']['xrange']:
range_ = alias
else:
break
if range_ is None:
return self.generic_visit(node)
# everything is setup for the transformation!
new_id = node.left.id + '_m'
i = 0
while new_id in self.identifiers:
new_id = '{}_m{}'.format(node.left.id, i)
i += 1
rargs = range_.args.args
lower = rargs[0] if len(rargs) > 1 else ast.Num(0)
header = ast.Assign([ast.Name(new_id, ast.Store(), None)],
ast.BinOp(
ast.BinOp(deepcopy(lower), ast.Sub(),
ast.Num(1)), ast.Mod(),
deepcopy(node.right)))
incr = ast.BinOp(ast.Name(new_id, ast.Load(), None), ast.Add(),
ast.Num(1))
step = ast.Assign([ast.Name(new_id, ast.Store(), None)],
ast.IfExp(
ast.Compare(incr, [ast.Eq()],
[deepcopy(node.right)]), ast.Num(0),
deepcopy(incr)))
self.loops_mod.setdefault(loop, []).append((header, step))
self.update = True
return ast.Name(new_id, ast.Load(), None)
