def test_alias():
assert eq(
alias.foo,
Alias(name='foo', asname=None)
)
assert sourcify(alias.foo) == 'foo'
assert eq(
alias['foo'].as_('bar'),
Alias(name='foo', asname='bar')
)
assert sourcify(alias['foo'].as_('bar')) == 'foo as bar'
def test_try_except_else_finally():
stmt = try_(
var.x.store(1)
).except_(var.TypeError)[
var.y.store(2)
].except_(var.ValueError)[
var.y.store(3)
].except_(var.OSError, var.RuntimeError, as_='foo')[
pass_,
].else_(
var.x.iadd(1)
).finally_(
var.print(1, 2, 3)
)
expected = """\
try:
x = 1
except TypeError:
y = 2
except ValueError:
y = 3
except (OSError, RuntimeError) as foo:
pass
else:
x += 1
finally:
print(1, 2, 3)"""
result = sourcify(stmt)
assert result == expected
def test_import():
assert eq(
import_(alias.bar),
ast.Import(names=[Alias(name='bar', asname=None)])
)
assert eq(import_('bar'), import_(alias.bar))
assert sourcify(import_('bar')) == 'import bar'
def test_for():
loop = for_(var.x).in_(var.y)[
var.print(1)
]
assert sourcify(loop) == """\
for x in y:
print(1)"""
assert loop is not None
def test_while():
loop = while_(var.x < var.y)[
pass_
]
assert sourcify(loop) == """\
while x < y:
pass"""
assert loop is not None
def test_classdef():
myklass = class_.Yuge(var.object, metaclass=var.object)[
def_.method1(arg.self, arg.a)[
if_(var.a == 1)[return_(var.a + 1)].else_[return_(1)]
]
]
s = sourcify(myklass)
assert s == """\
def test_try_except():
stmt = try_(
var.print(2)
).except_(var.TypeError)[
var.print(1)
]
result = sourcify(stmt)
expected = """\
try:
print(2)
except TypeError:
print(1)"""
assert result == expected
def test_try_finally():
stmt = try_(
var.print(2)
).finally_(
var.print(1)
)
result = sourcify(stmt)
expected = """\
try:
print(2)
finally:
print(1)"""
assert result == expected
def test_try_except_finally():
stmt = try_(
var.x.store(1)
).except_(var.TypeError)[
var.y.store(2)
].finally_(
return_(3)
)
result = sourcify(stmt)
expected = """\
try:
x = 1
except TypeError:
y = 2
finally:
return 3"""
assert result == expected
def test_complex_class():
klass = class_.Average[
def_.__init__(arg.self)[
var.self.value.store(0.0),
var.self.count.store(0),
],
def_.step(arg.self, arg.value)[
if_(var.value.is_not(NONE))[
var.self.value.iadd(var.value),
var.self.count.iadd(1)
]
],
def_.finalize(arg.self)[
if_(var.self.count)[
return_(var.self.value / var.self.count)
]
]
]
assert sourcify(klass) == """\
def test_multiple_assignment(expr, expected):
result = sourcify(expr)
assert result == expected
def test_yield_from():
func = def_.my_func(arg.x)[
yield_from(var.x)
]
result = sourcify(func)
assert result == """
def test_args_kwargs():
func = def_.my_func(*arg.args, **arg.kwargs)[
var.print(1)
]
result = sourcify(func)
assert result == """
def test_lambda():
func = lambda_(arg.x, arg.y)[var.x + 1 * var.y]
assert sourcify(func) == 'lambda x, y: x + 1 * y'
def test_store_compile():
expr = mod(var.self.x.store(0))
assert sourcify(expr) == 'self.x = 0'
assert ast.dump(expr) == ast.dump(ast.parse('self.x = 0'))