def _create_loop_options(self, node):
if not anno.hasanno(node, anno.Basic.DIRECTIVES):
return gast.Dict([], [])
loop_directives = anno.getanno(node, anno.Basic.DIRECTIVES)
if directives.set_loop_options not in loop_directives:
return gast.Dict([], [])
opts_dict = loop_directives[directives.set_loop_options]
str_keys, values = zip(*opts_dict.items())
keys = [gast.Constant(s, kind=None) for s in str_keys]
values = list(values) # ast and gast don't play well with tuples.
return gast.Dict(keys, values)
def keywords_to_dict(keywords):
keys = []
values = []
for kw in keywords:
keys.append(gast.Str(kw.arg))
values.append(kw.value)
return gast.Dict(keys=keys, values=values)
def size_container_folding(value):
"""
Convert value to ast expression if size is not too big.
Converter for sized container.
"""
if len(value) < MAX_LEN:
if isinstance(value, list):
return ast.List([to_ast(elt) for elt in value], ast.Load())
elif isinstance(value, tuple):
return ast.Tuple([to_ast(elt) for elt in value], ast.Load())
elif isinstance(value, set):
return ast.Set([to_ast(elt) for elt in value])
elif isinstance(value, dict):
keys = [to_ast(elt) for elt in value.keys()]
values = [to_ast(elt) for elt in value.values()]
return ast.Dict(keys, values)
elif isinstance(value, numpy.ndarray):
return ast.Call(func=ast.Attribute(
ast.Name(mangle('numpy'), ast.Load(), None), 'array',
ast.Load()),
args=[to_ast(value.tolist())],
keywords=[])
else:
raise ConversionError()
else:
raise ToNotEval()
def get_annotations(object_def, namespace):
"""Create the annotations from a definition node"""
# print_dump(object_def)
ast_annotations = ast.Assign(
targets=[extast.Name("annotations", ast.Store())],
value=ast.Dict(keys=[], values=[]),
type_comment=None,
)
if isinstance(object_def, ast.FunctionDef):
_fill_ast_annotations_function(object_def, ast_annotations)
elif isinstance(object_def, ast.ClassDef):
_fill_ast_annotations_class(object_def, ast_annotations)
else:
raise NotImplementedError
# print_dump(ast_annotations)
source = extast.unparse(ast_annotations)
try:
del namespace["__builtins__"]
except KeyError:
pass
exec(source, namespace)
return namespace["annotations"]
def keywords_to_dict(keywords):
"""Converts a list of ast.keyword objects to a dict."""
keys = []
values = []
for kw in keywords:
keys.append(gast.Str(kw.arg))
values.append(kw.value)
return gast.Dict(keys=keys, values=values)
def size_container_folding(value):
"""
Convert value to ast expression if size is not too big.
Converter for sized container.
"""
def size(x):
return len(getattr(x, 'flatten', lambda: x)())
if size(value) < MAX_LEN:
if isinstance(value, list):
return ast.List([to_ast(elt) for elt in value], ast.Load())
elif isinstance(value, tuple):
return ast.Tuple([to_ast(elt) for elt in value], ast.Load())
elif isinstance(value, set):
if value:
return ast.Set([to_ast(elt) for elt in value])
else:
return ast.Call(func=ast.Attribute(
ast.Name(mangle('builtins'), ast.Load(), None, None),
'set', ast.Load()),
args=[],
keywords=[])
elif isinstance(value, dict):
keys = [to_ast(elt) for elt in value.keys()]
values = [to_ast(elt) for elt in value.values()]
return ast.Dict(keys, values)
elif isinstance(value, np.ndarray):
if len(value) == 0:
return ast.Call(
func=ast.Attribute(
ast.Name(mangle('numpy'), ast.Load(), None, None),
'empty', ast.Load()),
args=[to_ast(value.shape),
dtype_to_ast(value.dtype.name)],
keywords=[])
else:
return ast.Call(func=ast.Attribute(
ast.Name(mangle('numpy'), ast.Load(), None, None), 'array',
ast.Load()),
args=[
to_ast(totuple(value.tolist())),
dtype_to_ast(value.dtype.name)
],
keywords=[])
else:
raise ConversionError()
else:
raise ToNotEval()
def visit_For(self, node):
node = self.generic_visit(node)
(basic_loop_vars, composite_loop_vars,
reserved_symbols, possibly_undefs) = self._get_loop_vars(
node,
(anno.getanno(node, annos.NodeAnno.BODY_SCOPE).modified
| anno.getanno(node, annos.NodeAnno.ITERATE_SCOPE).modified))
loop_vars, loop_vars_ast_tuple = self._loop_var_constructs(
basic_loop_vars)
body_name = self.ctx.namer.new_symbol('loop_body', reserved_symbols)
state_getter_name = self.ctx.namer.new_symbol('get_state',
reserved_symbols)
state_setter_name = self.ctx.namer.new_symbol('set_state',
reserved_symbols)
state_functions = self._create_state_functions(composite_loop_vars,
state_getter_name,
state_setter_name)
if anno.hasanno(node, 'extra_test'):
extra_test = anno.getanno(node, 'extra_test')
extra_test_name = self.ctx.namer.new_symbol(
'extra_test', reserved_symbols)
template = """
def extra_test_name(loop_vars):
return extra_test_expr
"""
extra_test_function = templates.replace(
template,
extra_test_name=extra_test_name,
loop_vars=loop_vars,
extra_test_expr=extra_test)
else:
extra_test_name = parser.parse_expression('None')
extra_test_function = []
# Workaround for PEP-3113
# iterates_var holds a single variable with the iterates, which may be a
# tuple.
iterates_var_name = self.ctx.namer.new_symbol('iterates',
reserved_symbols)
template = """
iterates = iterates_var_name
"""
iterate_expansion = templates.replace(
template,
iterates=node.target,
iterates_var_name=iterates_var_name)
undefined_assigns = self._create_undefined_assigns(possibly_undefs)
basic_symbol_names = tuple(
gast.Str(str(symbol)) for symbol in basic_loop_vars)
composite_symbol_names = tuple(
gast.Str(str(symbol)) for symbol in composite_loop_vars)
# TODO(b/140125096): Populate.
opts = gast.Dict([], [])
# TODO(mdan): Use a single template.
# If the body and test functions took a single tuple for loop_vars, instead
# of *loop_vars, then a single template could be used.
if loop_vars:
template = """
undefined_assigns
state_functions
def body_name(iterates_var_name, loop_vars):
iterate_expansion
body
return loop_vars,
extra_test_function
loop_vars_ast_tuple = ag__.for_stmt(
iter_,
extra_test_name,
body_name,
state_getter_name,
state_setter_name,
(loop_vars,),
(basic_symbol_names,),
(composite_symbol_names,),
opts)
"""
return templates.replace(
template,
undefined_assigns=undefined_assigns,
loop_vars=loop_vars,
loop_vars_ast_tuple=loop_vars_ast_tuple,
iter_=node.iter,
iterate_expansion=iterate_expansion,
iterates_var_name=iterates_var_name,
extra_test_name=extra_test_name,
extra_test_function=extra_test_function,
body_name=body_name,
body=node.body,
state_functions=state_functions,
state_getter_name=state_getter_name,
state_setter_name=state_setter_name,
basic_symbol_names=basic_symbol_names,
composite_symbol_names=composite_symbol_names,
opts=opts)
else:
template = """
undefined_assigns
state_functions
def body_name(iterates_var_name):
iterate_expansion
body
return ()
extra_test_function
ag__.for_stmt(
iter_,
extra_test_name,
body_name,
state_getter_name,
state_setter_name,
(),
(),
(composite_symbol_names,),
opts)
"""
return templates.replace(
template,
undefined_assigns=undefined_assigns,
iter_=node.iter,
iterate_expansion=iterate_expansion,
iterates_var_name=iterates_var_name,
extra_test_name=extra_test_name,
extra_test_function=extra_test_function,
body_name=body_name,
body=node.body,
state_functions=state_functions,
state_getter_name=state_getter_name,
state_setter_name=state_setter_name,
composite_symbol_names=composite_symbol_names,
opts=opts)
def visit_While(self, node):
node = self.generic_visit(node)
(basic_loop_vars, composite_loop_vars, reserved_symbols,
possibly_undefs) = self._get_loop_vars(
node,
anno.getanno(node, annos.NodeAnno.BODY_SCOPE).modified)
loop_vars, loop_vars_ast_tuple = self._loop_var_constructs(
basic_loop_vars)
state_getter_name = self.ctx.namer.new_symbol('get_state',
reserved_symbols)
state_setter_name = self.ctx.namer.new_symbol('set_state',
reserved_symbols)
state_functions = self._create_state_functions(composite_loop_vars,
state_getter_name,
state_setter_name)
basic_symbol_names = tuple(
gast.Str(str(symbol)) for symbol in basic_loop_vars)
composite_symbol_names = tuple(
gast.Str(str(symbol)) for symbol in composite_loop_vars)
# TODO(b/140125096): Populate.
opts = gast.Dict([], [])
# TODO(mdan): Use a single template.
# If the body and test functions took a single tuple for loop_vars, instead
# of *loop_vars, then a single template could be used.
if loop_vars:
template = """
state_functions
def body_name(loop_vars):
body
return loop_vars,
def test_name(loop_vars):
return test
loop_vars_ast_tuple = ag__.while_stmt(
test_name,
body_name,
state_getter_name,
state_setter_name,
(loop_vars,),
(basic_symbol_names,),
(composite_symbol_names,),
opts)
"""
node = templates.replace(
template,
loop_vars=loop_vars,
loop_vars_ast_tuple=loop_vars_ast_tuple,
test_name=self.ctx.namer.new_symbol('loop_test',
reserved_symbols),
test=node.test,
body_name=self.ctx.namer.new_symbol('loop_body',
reserved_symbols),
body=node.body,
state_functions=state_functions,
state_getter_name=state_getter_name,
state_setter_name=state_setter_name,
basic_symbol_names=basic_symbol_names,
composite_symbol_names=composite_symbol_names,
opts=opts)
else:
template = """
state_functions
def body_name():
body
return ()
def test_name():
return test
ag__.while_stmt(
test_name,
body_name,
state_getter_name,
state_setter_name,
(),
(),
(composite_symbol_names,),
opts)
"""
node = templates.replace(
template,
test_name=self.ctx.namer.new_symbol('loop_test',
reserved_symbols),
test=node.test,
body_name=self.ctx.namer.new_symbol('loop_body',
reserved_symbols),
body=node.body,
state_functions=state_functions,
state_getter_name=state_getter_name,
state_setter_name=state_setter_name,
composite_symbol_names=composite_symbol_names,
opts=opts)
undefined_assigns = self._create_undefined_assigns(possibly_undefs)
return undefined_assigns + node