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 = parsing.ast_to_object(node)
expected_source = """
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 test_keywords_to_dict(self):
keywords = parsing.parse_expression('f(a=b, c=1, d=\'e\')').keywords
d = ast_util.keywords_to_dict(keywords)
# Make sure we generate a usable dict node by attaching it to a variable and
# compiling everything.
node = parsing.parse_str('def f(b): pass').body[0]
node.body.append(ast.Return(d))
result, _ = parsing.ast_to_object(node)
self.assertDictEqual(result.f(3), {'a': 3, 'c': 1, 'd': 'e'})
def test_replace_tuple(self):
template = """
def test_fn(a, c):
return b,
"""
node = templates.replace(template, b=('a', 'c'))[0]
result, _ = parsing.ast_to_object(node)
self.assertEqual((2, 3), result.test_fn(2, 3))
def test_replace_name_with_dict(self):
template = """
def test_fn():
return foo['bar']
"""
source = parsing.parse_expression('{\'bar\': 3}')
node = templates.replace(template, foo=source)[0]
result, _ = parsing.ast_to_object(node)
self.assertEqual(3, result.test_fn())
def test_replace_function_name(self):
template = """
def fname(a):
a += 1
a = 2 * a + 1
return a
"""
node = templates.replace(template, fname='test_fn')[0]
result, _ = parsing.ast_to_object(node)
self.assertEqual(7, result.test_fn(2))
def test_replace_variable(self):
template = """
def test_fn(a):
a += 1
a = 2 * a + 1
return b
"""
node = templates.replace(template, a='b')[0]
result, _ = parsing.ast_to_object(node)
self.assertEqual(7, result.test_fn(2))
def test_parsing_compile_idempotent(self):
def test_fn(x):
a = True
b = ''
if a:
b = x + 1
return b
self.assertEqual(
textwrap.dedent(inspect.getsource(test_fn)),
inspect.getsource(
parsing.ast_to_object(
parsing.parse_entity(test_fn)[0].body[0])[0].test_fn))
def test_replace_attribute(self):
template = """
def test_fn(a):
return a.foo
"""
node = templates.replace(template, foo='b')[0]
result, _ = parsing.ast_to_object(node)
mod = imp.new_module('test')
mod.b = 3
self.assertEqual(3, result.test_fn(mod))
with self.assertRaises(ValueError):
templates.replace(template, foo=1)
def test_replace_name_with_call(self):
template = """
def test_fn():
b = 5
def g(a):
return 3 * a
def f():
return g
return foo
"""
source = parsing.parse_expression('f()(b)')
node = templates.replace(template, foo=source)[0]
result, _ = parsing.ast_to_object(node)
self.assertEqual(15, result.test_fn())
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, _ = parsing.ast_to_object(node)
self.assertEqual(3, result.test_fn(1))
def test_replace_call_keyword(self):
template = """
def test_fn():
def f(a, d, f):
return a + d + f
return f(1, kws=None)
"""
source = parsing.parse_expression('f(d=3, f=5)')
node = templates.replace(template, kws=source.keywords)[0]
result, _ = parsing.ast_to_object(node)
self.assertEqual(9, result.test_fn())
with self.assertRaises(ValueError):
templates.replace(template, kws=[])
templates.replace(template, kws=1)
def _wrap_in_generator(func, source, namer, overload):
"""Wraps the source code in a generated function.
Args:
func: the original function
source: the generated source code
namer: naming.Namer, used for naming vars
overload: config.VirtualizationConfig
Returns:
The generated function with a new closure variable.
"""
nonlocals = []
for var in six.get_function_code(func).co_freevars:
# We must generate dummy vars so the generated function has the same closure
# as the original function.
free_template = 'var = None'
nonlocal_node = templates.replace(free_template, var=var)
nonlocals.extend(nonlocal_node)
gen_fun_name = namer.new_symbol('gen_fun', set())
template = """
def gen_fun(overload):
nonlocals
program
return f_name
"""
ret = templates.replace(
template,
gen_fun=gen_fun_name,
nonlocals=nonlocals,
overload=overload.symbol_name,
program=source,
f_name=func.__name__)
converted_module, _ = parsing.ast_to_object(ret)
outer_func = getattr(converted_module, gen_fun_name)
return outer_func(overload.module)