def test_primitive_float(assertion_to_ast):
assertion = MagicMock(value=1.5)
assertion_to_ast.visit_primitive_assertion(assertion)
assert (
astor.to_source(Module(body=assertion_to_ast.nodes))
== "assert var0 == pytest.approx(1.5, abs=0.01, rel=0.01)\n"
)
def test_not_none(assertion_to_ast):
assertion = MagicMock(value=False)
assertion_to_ast.visit_none_assertion(assertion)
assert (
astor.to_source(Module(body=assertion_to_ast.nodes))
== "assert var0 is not None\n"
)
def test_primitive_non_bool(assertion_to_ast):
assertion = MagicMock(value=42)
assertion_to_ast.visit_primitive_assertion(assertion)
assert astor.to_source(
Module(body=assertion_to_ast.nodes)) == "assert var0 == 42\n"
def test_primitive_float(assertion_to_ast):
assertion = MagicMock(value=1.5)
assertion_to_ast.visit_primitive_assertion(assertion)
assert (astor.to_source(Module(body=assertion_to_ast.nodes)) ==
"assert math.isclose(var0, 1.5, abs_tol=0.01)\n")
def test_primitive_bool(assertion_to_ast):
assertion = MagicMock(value=True)
assertion_to_ast.visit_primitive_assertion(assertion)
assert (astor.to_source(
Module(body=assertion_to_ast.nodes)) == "assert var0 is True\n")
def module_factory(doc):
node = Module()
node.body = []
node.doc = doc
return node
def visit_Module(self, node):
"""This is the main entry point of the source AST.
We want to wrap the expression the user sent us
with a function definition, a loop, and then
an initialization of all references variables.
This might be easier to read if we just modified the source,
but it is far more dependable to work with raw ast nodes,
and never even think of manually parsing any strings."""
# we'll assemble an array of statements to form the body of the generated function
fun_body = []
for_body = []
arg_format = '{}' if self.mode is VecMode.vectorize else '{}_'
# Lets prepare the arguments. First, all column references are translated from colname into colname_
args = list(map(lambda x: arg_or_name(arg_format.format(x.id)), values(self.analyzed.all_names)))
# Next, add any external references that are simply passed through to the function
args.extend(list(map(lambda x: arg_or_name(x[0]), self.external)))
if self.mode is not VecMode.vectorize:
# Add the returned array. We will give it a default value of None, and initialize it when it is None
args.append(arg_or_name(self.key_result))
default_args = [read_none()] if self.mode is VecMode.jit else []
# Add the desired decorator.
decorators = [self.choose_decorator()]
if self.mode is VecMode.jit:
# only jit-mode has to worry about an empty output array being passed.
# vectorize and guvectorize both handle output arrays "for free".
# assemble: dtype=np.return_type
ret_dtype = keyword_arg('dtype', field_read('np', self.return_type_name()))
# assemble: np.empty(arg0.shape, dtype=np.return_type)
create_empty = method_call(
scope=field_read('np', 'empty'),
args=[field_read(read(args[0]), 'shape')],
keywords=[ret_dtype]
)
# assemble: out = np.empty(arg0.shape, dtype=np.return_type)
create_result = assign(self.key_result, create_empty)
# assemble: if out is None:
# out = np.empty(arg0.shape, dtype=np.return_type)
check_result = check_if_none(self.key_result, [create_result])
fun_body.append(check_result)
# Fill in the body of the function. All types except vectorize use the same code path.
if self.mode is VecMode.vectorize:
if self.analyzed.last_assign is None:
fun_body.extend(node.body[0:-1])
arr_assign = assign(self.key_result, node.body[-1].value)
fun_body.append(arr_assign)
else:
fun_body.extend(node.body)
else:
# jit, guvectorize
# assemble: for i__ in range(arg0.shape[0]):
# a = a__[i__]
for arg in values(self.analyzed.all_names):
# Create assignments to the expected variable names w/in the loop
self.extra_lines += 1
# assemble: a = a[i__]
a = assign(arg.id, array_access(arg_format.format(arg.id), self.index))
for_body.append(a)
# visit the original body of source code, to do transformation on assignments
# This O(n) visit is only necessary if the original expression contained an assignment.
fixed_body = node.body
if self.analyzed.last_assign is None:
# The user did not supply an assignment statement; build one for them.
for_body.extend(fixed_body[0:-1])
arr_assign = assign(write_array(self.key_result, self.index), fixed_body[-1].value)
for_body.append(arr_assign)
else:
fixed_body = list(map(lambda x: increment_lineno(x, self.extra_lines), fixed_body))
fixed_body = list(map(lambda x: self.generic_visit(x), fixed_body))
# append the fixed statements to the new body
for_body.extend(fixed_body)
# assemble: r__.shape[0]
shape = array_access(field_read(self.key_result, 'shape'), 0)
# assemble: range(r__.shape[0])
shape_range = method_call('range', [shape])
# assemble: for i__ in range(r__.shape[0]):
for_loop = For(target=write(Transformer._key_index), iter=shape_range, body=for_body, orelse=[])
# print our loop
fun_body.append(for_loop)
# print the return statement
fun_body.append(Return(read(self.key_result)))
# create the whole function
func = new_func(self.func_name, args, fun_body, decorators, default_args)
# build a new module w/ just the import and the function (so we can call it later)
# if we can forgo the imports, or do them manually to the environment,
# then we could use `single` mode instead of `exec` mode when compiling,
# and just return a single function (effectively what the java version did)
fixed = Module(body=[
# # import numpy as np
# Import(names=[alias(name='numpy', asname='np')]),
# # from numba import jit
# ImportFrom(module='numba', names=[alias(name='jit', asname=None)], level=0),
# the function we just built
func])
# built-in module to generate source information in the generated results.
# TODO: verify these are actually correct in unit test
node = fix_missing_locations(fixed)
# print dump(node)
return node