def test_test_case_to_ast_twice(simple_test_case):
visitor = tc_to_ast.TestCaseToAstVisitor()
simple_test_case.accept(visitor)
simple_test_case.accept(visitor)
assert (astor.to_source(Module(body=visitor.test_case_asts[0])) ==
"var0 = 5\nvar1 = module0.SomeType(var0)\nassert var1 == 3\n")
assert (astor.to_source(Module(body=visitor.test_case_asts[1])) ==
"var0 = 5\nvar1 = module0.SomeType(var0)\nassert var1 == 3\n")
def _build_thunk(node: ast.Module, fn_cls: Type, globals: Dict, locals: Dict):
# Temporarily polluting the `locals` scope with the thunk function is not
# ideal, but choosing an illegal name means that only very sneaky code will
# be aware of it.
thunkname = '-aeval-' + ''.join(random.sample(string.hexdigits, 8))
global_names = exposed_names_with_store(node)
if global_names:
node.body.insert(0, ast.Global(list(global_names)))
if len(node.body) > 0 and isinstance(node.body[-1], ast.Expr):
node.body[-1] = ast.Return(value=node.body[-1].value)
node = RewriteExposedAnnotatedNames().visit(node)
mod = ast.parse('')
mod.body = [
fn_cls(
name=thunkname,
args=_EMPTY_ARGS,
body=node.body,
decorator_list=[],
returns=None,
)
]
ast.fix_missing_locations(mod)
exec(compile(mod, '<string>', 'exec'), globals, locals)
thunk = locals[thunkname]
thunk.__name__ = 'aeval_thunk'
del locals[thunkname]
return thunk
def file(node, filename, mode="a", skip_black=False):
"""
Convert AST to a file
:param node: AST node
:type node: ```Union[Module, ClassDef, FunctionDef]```
:param filename: emit to this file
:type filename: ```str```
:param mode: Mode to open the file in, defaults to append
:type mode: ```str```
:param skip_black: Skip formatting with black
:type skip_black: ```bool```
:return: None
:rtype: ```NoneType```
"""
if isinstance(node, (ClassDef, FunctionDef)):
node = Module(body=[node], type_ignores=[], stmt=None)
src = to_code(node)
if not skip_black:
src = format_str(
src,
mode=Mode(
target_versions=set(),
line_length=119,
is_pyi=False,
string_normalization=False,
),
)
with open(filename, mode) as f:
f.write(src)
def visit_With(self, node):
# With(items=[withitem(context_expr=Call(func=Name(id='foo', ctx=Load()),
# args=[], keywords=[], starargs=None,
# kwargs=None),
# optional_vars=Name(id='bar', ctx=Store()))], body=[Pass()])
for item in node.items:
if item.optional_vars is not None:
self.bind(item.optional_vars.id)
self.generic_visit(node)
# handle 'remote'
sub_node = node.items[0].context_expr
if (type(sub_node) == Call and hasattr(sub_node.func, 'id')
and sub_node.func.id == 'remote'):
# capture the body and put it as the first argument to ssh()
# but within a module, and already pickled;
# otherwise we need to create an AST for the call of all the
# constructors in the body we capture... it's complicated
m = Module(body=node.body)
data = pickle.dumps(m)
s = Bytes(s=data)
s.lineno = node.lineno
s.col_offset = node.col_offset
sub_node.args.insert(0, s)
p = Pass()
p.lineno = node.lineno + 1
p.col_offset = node.col_offset + 4
node.body = [p]
return node
def _structure_changed(self, name, new):
if not self._updating_structure:
try:
self._updating_structure = True
if name == 'ast_tree':
# Compute our new sub-blocks and give them our filename
sub_blocks = Block._decompose(self.ast_tree)
self.sub_blocks = sub_blocks
if self.sub_blocks is not None:
for b in self.sub_blocks:
b.filename = self.filename
# Policy: Only atomic blocks have filenames
self.filename = None
elif name in ('sub_blocks', 'sub_blocks_items'):
self.ast_tree = Module([ b.ast_tree for b in self.sub_blocks ])
else:
assert False
# Invalidate caches
self.__dep_graph_is_valid = False
self.__restrictions.clear()
self._stored_string = ''
# update inputs and outputs
self._clear_cache_inputs_and_outputs()
finally:
self._updating_structure = False
def __init__(self, filename):
if not filename.endswith(".py"):
raise ValueError("FunctionStore can only work with pure Python source code with .py extension")
super().__init__()
self._filename = fullname = os.path.abspath(os.path.realpath(filename))
path, filename = os.path.split(fullname)
modname = filename[:-3]
if sys.path[0] != path:
if path in sys.path:
sys.path.remove(path)
sys.path.insert(0, path)
try:
if modname in sys.modules:
self._module = sys.modules[modname]
else:
self._module = importlib.import_module(modname)
self._ast = parse(self._module.__loader__.get_data(fullname))
self._ast_idx = {}
for i, node in enumerate(self._ast.body):
self._ast_idx[node.name] = i
except (FileNotFoundError, ModuleNotFoundError):
self._module = None
self._ast = Module(body=[])
self._ast_idx = {}
self._need_save = False
self._locl = {}
def visit_Module(self, node):
imports = []
for subnode in node.body:
if isinstance(subnode, (Import, ImportFrom)):
imports.append(self.visit(subnode))
sorted_imports = self.sort_imports(imports)
return Module(body=sorted_imports)
def test_annotate_ancestry(self) -> None:
"""Tests that `annotate_ancestry` properly decorates"""
node = Module(
body=[
AnnAssign(
annotation=Name(
"str",
Load(),
),
simple=1,
target=Name("dataset_name", Store()),
value=set_value("~/tensorflow_datasets"),
expr=None,
expr_target=None,
expr_annotation=None,
),
Assign(
annotation=None,
simple=1,
targets=[Name("epochs", Store())],
value=set_value("333"),
expr=None,
expr_target=None,
expr_annotation=None,
**maybe_type_comment
),
],
stmt=None,
)
self.assertFalse(hasattr(node.body[0], "_location"))
self.assertFalse(hasattr(node.body[1], "_location"))
annotate_ancestry(node)
self.assertEqual(node.body[0]._location, ["dataset_name"])
self.assertEqual(node.body[1]._location, ["epochs"])
def _exec(excode):
"""this code is boiled down from
https://github.com/ipython/ipython/blob/master/IPython/core/interactiveshell.py
it makes sure that expressions on the
last line of a nbplots directive are evaluated and printed
"""
import ast
if sys.version_info > (3, 8):
from ast import Module
else:
# mock the new API, ignore second argument
# see https://github.com/ipython/ipython/issues/11590
from ast import Module as OriginalModule
Module = lambda nodelist, type_ignores: OriginalModule(nodelist)
nodelist = ast.parse(excode).body
if not nodelist:
return
if isinstance(nodelist[-1], ast.Expr):
to_run_exec, to_run_interactive = nodelist[:-1], nodelist[-1:]
else:
to_run_exec, to_run_interactive = nodelist, []
to_run = [(node, 'exec') for node in to_run_exec]
for node in to_run_interactive:
to_run.append((node, 'single'))
for node, mode in to_run:
if mode == 'exec':
mod = Module([node], [])
elif mode == 'single':
mod = ast.Interactive([node])
six.exec_(compile(mod, '<dummyfile>', mode), ns)
def getModuleScope(n: ast.Module, surrounding: tydict):
from . import classes
try:
aliases = gather_aliases(n, {})
aliases.update(n.retic_import_aliases.copy())
theclasses, classenv, aliasenv = classes.get_class_scope(
n.body, surrounding, n.retic_import_env, aliases)
aliases.update(aliasenv)
local = InitialScopeFinder().preorder(n.body, aliases)
local.update(classenv)
local.update(n.retic_import_env) # We probably want to make
# sure there's no conflict
# between imports and
# annotated locals
except InconsistentAssignment as e:
raise exc.StaticTypeError(
None,
'Multiple bindings of {} occur at the top level with differing types: {} and {}'
.format(e.args[0], e.args[1], e.args[2]))
modscope = surrounding.copy() if surrounding else {}
modscope.update(env.module_env())
modscope.update(local)
inferred = infer_types(modscope, local, n.body, theclasses)
n.retic_aliases = aliases
return inferred, aliases
def test_statement_to_ast_int(statement_to_ast_visitor):
int_stmt = MagicMock(stmt.Statement)
int_stmt.value = 5
statement_to_ast_visitor.visit_int_primitive_statement(int_stmt)
assert (
astor.to_source(Module(body=statement_to_ast_visitor.ast_nodes)) == "var0 = 5\n"
)
def getter(item):
"""Construct the getter function.
partof: #SPC-asts.getter
"""
func_name = f"{item.var}_getter"
self_arg = arg(arg="self", annotation=None)
func_args = arguments(
args=[self_arg],
kwonlyargs=[],
vararg=None,
kwarg=None,
defaults=[],
kw_defaults=[],
)
inst_var = Attribute(value=Name(id="self", ctx=ast.Load()),
attr=f"_{item.var}",
ctx=ast.Load())
ret_stmt = Return(value=inst_var)
func_node = FunctionDef(name=func_name,
args=func_args,
body=[ret_stmt],
decorator_list=[],
returns=None)
mod_node = Module(body=[func_node])
return ast_to_func(mod_node, func_name)
def inner(tree: ast.Module, use_pypprint: bool = False) -> bool:
if not tree.body:
return False
if not isinstance(tree.body[-1], ast.Expr):
if not isinstance(tree.body[-1], ast.Pass):
return False
del tree.body[-1]
return True
if isinstance(tree.body[-1].value, ast.Name):
output = tree.body[-1].value.id
tree.body.pop()
else:
output = self.get_valid_name("output")
self.define(output)
tree.body[-1] = ast.Assign(
targets=[ast.Name(id=output, ctx=ast.Store())], value=tree.body[-1].value
)
print_fn = "print"
if use_pypprint:
print_fn = "pypprint"
self.undefined.add("pypprint")
if_print = ast.parse(f"if {output} is not None: {print_fn}({output})").body[0]
tree.body.append(if_print)
self.implicit_print = if_print.body[0].value # type: ignore
return True
def setter(item):
"""Construct the setter function.
partof: #SPC-asts.setter
"""
func_name = f"{item.var}_setter"
self_arg = arg(arg="self", annotation=None)
new_arg = arg(arg="new", annotation=None)
func_args = arguments(
args=[self_arg, new_arg],
kwonlyargs=[],
vararg=None,
kwarg=None,
defaults=[],
kw_defaults=[],
)
func_node = FunctionDef(
name=func_name,
args=func_args,
body=[setter_body(item)],
decorator_list=[],
returns=None,
)
mod_node = Module(body=[func_node])
return ast_to_func(mod_node, func_name)
def test_to_named_class_def(self) -> None:
"""Test that find_ast_type gives the wrapped named class back"""
class_def = ClassDef(
name="foo",
bases=tuple(),
keywords=tuple(),
decorator_list=[],
body=[],
expr=None,
identifier_name=None,
)
run_ast_test(
self,
find_ast_type(
Module(
body=[
ClassDef(
name="bar",
bases=tuple(),
keywords=tuple(),
decorator_list=[],
body=[],
expr=None,
identifier_name=None,
),
class_def,
],
stmt=None,
),
node_name="foo",
),
class_def,
skip_black=True,
)
def generic_visit(self, node):
# first handle some early exits
if self._next_comment is None:
return node
elif node is None:
# this can happen if the module contains nothing but comments
comments = [self._next_comment]
self._next_comment = None
comments.extend(reversed(self._comments))
self._comments.clear()
new_node = Module(body=comments)
return new_node
# grab prior body comments
self._grab_prior_body_comments(node)
if self._next_comment is None:
# can early exit again
return node
new_node = self._attach_comment(node)
if hasattr(node, "body"):
self._comments_in_body.append([])
for i, n in enumerate(node.body):
node.body[i] = self.visit(n)
# grab trainling body comments
while (self._next_comment is not None
and self._next_comment.col_offset >= n.col_offset):
self._comments_in_body[-1].append(self._next_comment)
self._next_comment = self._comments.pop(
) if self._comments else None
merge_body_comments(node.body, self._comments_in_body.pop())
elif node is not new_node:
self.visit(node)
return new_node
def test_statement_to_ast_function_args(statement_to_ast_visitor,
test_case_mock, function_mock):
function_stmt = param_stmt.FunctionStatement(
test_case_mock, function_mock, [MagicMock(vr.VariableReference)])
statement_to_ast_visitor.visit_function_statement(function_stmt)
assert (astor.to_source(Module(body=statement_to_ast_visitor.ast_nodes)) ==
"var1 = module0.simple_function(var0)\n")
def test_statement_to_ast_assignment(variable_reference_mock,
statement_to_ast_visitor):
assign_stmt = MagicMock(stmt.Statement)
assign_stmt.ret_val = variable_reference_mock
assign_stmt.rhs = variable_reference_mock
statement_to_ast_visitor.visit_assignment_statement(assign_stmt)
assert (astor.to_source(
Module(body=statement_to_ast_visitor.ast_nodes)) == "var0 = var0\n")
def test_statement_to_ast_none(statement_to_ast_visitor):
none_stmt = MagicMock(stmt.Statement)
none_stmt.value = None
statement_to_ast_visitor.visit_none_statement(none_stmt)
assert (
astor.to_source(Module(body=statement_to_ast_visitor.ast_nodes))
== "var0 = None\n"
)
def test_statement_to_ast_str(statement_to_ast_visitor):
str_stmt = MagicMock(stmt.Statement)
str_stmt.value = "TestMe"
statement_to_ast_visitor.visit_string_primitive_statement(str_stmt)
assert (
astor.to_source(Module(body=statement_to_ast_visitor.ast_nodes))
== "var0 = 'TestMe'\n"
)
def test_statement_to_ast_bool(statement_to_ast_visitor):
bool_stmt = MagicMock(stmt.Statement)
bool_stmt.value = True
statement_to_ast_visitor.visit_boolean_primitive_statement(bool_stmt)
assert (
astor.to_source(Module(body=statement_to_ast_visitor.ast_nodes))
== "var0 = True\n"
)
def test_statement_to_ast_method_no_args(statement_to_ast_visitor,
test_case_mock,
variable_reference_mock, method_mock):
method_stmt = param_stmt.MethodStatement(test_case_mock, method_mock,
variable_reference_mock)
statement_to_ast_visitor.visit_method_statement(method_stmt)
assert (astor.to_source(Module(body=statement_to_ast_visitor.ast_nodes)) ==
"var1 = var0.simple_method()\n")
def test_statement_to_ast_constructor_no_args(statement_to_ast_visitor,
test_case_mock,
constructor_mock):
constr_stmt = param_stmt.ConstructorStatement(test_case_mock,
constructor_mock)
statement_to_ast_visitor.visit_constructor_statement(constr_stmt)
assert (astor.to_source(Module(body=statement_to_ast_visitor.ast_nodes)) ==
"var0 = module0.SomeType()\n")
def clean_top_level(tree: ast.Module) -> ast.AST:
tree.body = [
node for node in tree.body
if isinstance(node, _ast.FunctionDef) \
or isinstance(node, _ast.ClassDef) \
or isinstance(node, _ast.Import) \
or isinstance(node, _ast.ImportFrom)
]
def __init__(self,
grid,
ptype,
pyfunc=None,
funcname=None,
funccode=None,
py_ast=None,
funcvars=None):
self.grid = grid
self.ptype = ptype
# Derive meta information from pyfunc, if not given
self.funcname = funcname or pyfunc.__name__
self.funcvars = funcvars or list(pyfunc.__code__.co_varnames)
self.funccode = funccode or inspect.getsource(pyfunc.__code__)
# Parse AST if it is not provided explicitly
self.py_ast = py_ast or parse(fix_indentation(self.funccode)).body[0]
if pyfunc is None:
# Extract user context by inspecting the call stack
stack = inspect.stack()
try:
user_ctx = stack[-1][0].f_globals
user_ctx['math'] = globals()['math']
user_ctx['random'] = globals()['random']
except:
print(
"Warning: Could not access user context when merging kernels"
)
user_ctx = globals()
finally:
del stack # Remove cyclic references
# Compile and generate Python function from AST
py_mod = Module(body=[self.py_ast])
exec(compile(py_mod, "<ast>", "exec"), user_ctx)
self.pyfunc = user_ctx[self.funcname]
else:
self.pyfunc = pyfunc
self.name = "%s%s" % (ptype.name, self.funcname)
# Generate the kernel function and add the outer loop
if self.ptype.uses_jit:
kernelgen = KernelGenerator(grid, ptype)
self.field_args = kernelgen.field_args
kernel_ccode = kernelgen.generate(deepcopy(self.py_ast),
self.funcvars)
self.field_args = kernelgen.field_args
loopgen = LoopGenerator(grid, ptype)
adaptive = 'AdvectionRK45' in self.funcname
self.ccode = loopgen.generate(self.funcname,
self.field_args,
kernel_ccode,
adaptive=adaptive)
basename = path.join(get_cache_dir(), self._cache_key)
self.src_file = "%s.c" % basename
self.lib_file = "%s.so" % basename
self.log_file = "%s.log" % basename
self._lib = None
def visit(self, node: ast.Module) -> ast.Module:
"""Transform a Python module AST node.
:param node: ast.Module object
:return: ast.Module object
"""
node.body = [*self._transform_module_body(node.body)]
ast.fix_missing_locations(node)
return node
def start(self, tree):
from ast import Module, expr, Expr, Call, fix_missing_locations
stmts = [
Expr(s) if isinstance(s, expr) else s
for s in reversed(self.statements)
]
if not stmts:
stmts.append(tree[0])
return fix_missing_locations(Module(body=stmts, type_ignores=[]))
def test_statement_to_ast_function_no_args(
statement_to_ast_visitor, test_case_mock, function_mock
):
function_stmt = param_stmt.FunctionStatement(test_case_mock, function_mock)
statement_to_ast_visitor.visit_function_statement(function_stmt)
assert (
astor.to_source(Module(body=statement_to_ast_visitor.ast_nodes))
== "var0 = module0.simple_function()\n"
)
def annotate_funcs_in_module(module: ast.Module, decorator: ast.expr, copy: bool = False) -> ast.Module:
if copy:
module = deepcopy(module)
module.body = [
annotate_funcs_in_stmt(stmt, decorator, copy=False)
# no need to copy, since we are already copied the whole module
for stmt in module.body
]
return module
def run(self, mod: ast.Module) -> None:
"""Find all assert statements in *mod* and rewrite them."""
if not mod.body:
# Nothing to do.
return
# Insert some special imports at the top of the module but after any
# docstrings and __future__ imports.
aliases = [
ast.alias("builtins", "@py_builtins"),
ast.alias("_pytest.assertion.rewrite", "@pytest_ar"),
]
doc = getattr(mod, "docstring", None)
expect_docstring = doc is None
if doc is not None and self.is_rewrite_disabled(doc):
return
pos = 0
lineno = 1
for item in mod.body:
if (expect_docstring and isinstance(item, ast.Expr)
and isinstance(item.value, ast.Str)):
doc = item.value.s
if self.is_rewrite_disabled(doc):
return
expect_docstring = False
elif (not isinstance(item, ast.ImportFrom) or item.level > 0
or item.module != "__future__"):
lineno = item.lineno
break
pos += 1
else:
lineno = item.lineno
imports = [
ast.Import([alias], lineno=lineno, col_offset=0)
for alias in aliases
]
mod.body[pos:pos] = imports
# Collect asserts.
nodes = [mod] # type: List[ast.AST]
while nodes:
node = nodes.pop()
for name, field in ast.iter_fields(node):
if isinstance(field, list):
new = [] # type: List[ast.AST]
for i, child in enumerate(field):
if isinstance(child, ast.Assert):
# Transform assert.
new.extend(self.visit(child))
else:
new.append(child)
if isinstance(child, ast.AST):
nodes.append(child)
setattr(node, name, new)
elif (isinstance(field, ast.AST)
# Don't recurse into expressions as they can't contain
# asserts.
and not isinstance(field, ast.expr)):
nodes.append(field)
def getModuleScope(n: ast.Module, surrounding:tydict):
theclasses, classenv, ctbl, st = get_class_scope(n.body, surrounding, n.retic_import_cenv)
n.retic_cctbl = ctbl
local = InitialScopeFinder().preorder(n.body, False)
local.update(classenv)
local.update(n.retic_import_cenv)
modscope = surrounding.copy() if surrounding else {}
modscope.update(env.module_cenv())
modscope.update(local)
local = local_types(modscope, local, n.body)
return local, st
def getModuleScope(n: ast.Module, surrounding:tydict):
from . import classes
try:
aliases = gather_aliases(n, {})
aliases.update(n.retic_import_aliases.copy())
theclasses, classenv, aliasenv = classes.get_class_scope(n.body, surrounding, n.retic_import_env, aliases)
aliases.update(aliasenv)
local = InitialScopeFinder().preorder(n.body, aliases)
local.update(classenv)
local.update(n.retic_import_env) # We probably want to make
# sure there's no conflict
# between imports and
# annotated locals
except InconsistentAssignment as e:
raise exc.StaticTypeError(None, 'Multiple bindings of {} occur at the top level with differing types: {} and {}'.format(e.args[0], e.args[1], e.args[2]))
modscope = surrounding.copy() if surrounding else {}
modscope.update(env.module_env())
modscope.update(local)
inferred = infer_types(modscope, local, n.body, theclasses)
n.retic_aliases = aliases
return inferred, aliases
def parseModule(parser):
"""parse module statements"""
parser.check('Module ModuleID', lexeme=keywords['MODULE'] )
#create a new ast module.
#parse the module identifier.
parser.next()
module = Module(parseModuleId(parser))
# while there are more tokens
# parse for Site , Function and Import statements
while(parser.hasnext()):
if parser.matchLexeme(keywords['IMPORT']):
module.addImport(parseImport(parser))
elif parser.matchLexeme(keywords['DEF']):
module.addFunction(parseFunction(parser))
elif parser.matchLexeme(keywords['SITE']):
module.addSite(parseSite(parser))
elif parser.matchTokensort(ENDMARKER):
return
else:
raise SyntaxError(parser.currentToken,
expected="import, def, site, newline" )
return module
