def visit_Subscript(self, node):
node = self.generic_visit(node)
if not isinstance(node.slice, gast.Index):
return node
if not isinstance(node.ctx, gast.Load):
# Index writes are handled at a higher level, one at which the rvalue is
# also available.
return node
dtype = self.get_definition_directive(
node.value,
directives.set_element_type,
'dtype',
default=templates.replace_as_expression('None'))
template = """
ag__.get_item(
target,
key,
opts=ag__.GetItemOpts(element_dtype=dtype))
"""
return templates.replace_as_expression(template,
target=node.value,
key=node.slice,
dtype=dtype)
def _generate_pop_operation(self, original_call_node, pop_var_name):
assert isinstance(original_call_node.func, gast.Attribute)
if original_call_node.args:
pop_element = original_call_node.args[0]
else:
pop_element = parser.parse_expression('None')
# The call will be something like "target.pop()", and the dtype is hooked to
# target, hence the func.value.
dtype = anno.getanno(
original_call_node.func.value,
'element_type',
default=templates.replace_as_expression('None'))
shape = anno.getanno(
original_call_node.func.value,
'element_shape',
default=templates.replace_as_expression('None'))
template = """
target, pop_var_name = ag__.list_pop(
target, element,
opts=ag__.ListPopOpts(element_dtype=dtype, element_shape=shape))
"""
return templates.replace(
template,
target=original_call_node.func.value,
pop_var_name=pop_var_name,
element=pop_element,
dtype=dtype,
shape=shape)
def test_replace_as_expression_restrictions(self):
template = """
foo(a)
bar(b)
"""
with self.assertRaises(ValueError):
templates.replace_as_expression(template)
def _generate_pop_operation(self, original_call_node, pop_var_name):
assert isinstance(original_call_node.func, gast.Attribute)
if original_call_node.args:
pop_element = original_call_node.args[0]
else:
pop_element = parser.parse_expression('None')
# The call will be something like "target.pop()", and the dtype is hooked to
# target, hence the func.value.
dtype = anno.getanno(
original_call_node.func.value,
'element_type',
default=templates.replace_as_expression('None'))
shape = anno.getanno(
original_call_node.func.value,
'element_shape',
default=templates.replace_as_expression('None'))
template = """
target, pop_var_name = ag__.list_pop(
target, element,
opts=ag__.ListPopOpts(element_dtype=dtype, element_shape=shape))
"""
return templates.replace(
template,
target=original_call_node.func.value,
pop_var_name=pop_var_name,
element=pop_element,
dtype=dtype,
shape=shape)
def test_replace_as_expression_restrictions(self):
template = """
foo(a)
bar(b)
"""
with self.assertRaises(ValueError):
templates.replace_as_expression(template)
def visit_Subscript(self, node):
node = self.generic_visit(node)
if not isinstance(node.slice, gast.Index):
# TODO(mdan): It might make more sense to wave them through.
raise NotImplementedError('non-index slice')
if not isinstance(node.ctx, gast.Load):
# Index writes are handled at a higher level, one at which the rvalue is
# also available.
return node
dtype = anno.getanno(node.value,
'element_type',
default=templates.replace_as_expression('None'))
template = """
ag__.get_item(
target,
key,
opts=ag__.GetItemOpts(element_dtype=dtype))
"""
return templates.replace_as_expression(template,
target=node.value,
key=node.slice,
dtype=dtype)
def visit_For(self, node):
scope = anno.getanno(node, NodeAnno.BODY_SCOPE)
break_var = self.context.namer.new_symbol('break__', scope.referenced)
node.target = self.visit(node.target)
node.iter = self.visit(node.iter)
node.body, break_used = self._track_body(node.body, break_var)
# A break in the else clause applies to the containing scope.
node.orelse = self.visit_block(node.orelse)
if break_used:
node.orelse = self._guard_if_present(node.orelse, break_var)
template = """
var_name = False
for_stmt
"""
# Python's else clause only triggers if the loop exited cleanly (e.g.
# break did not trigger).
node = templates.replace(
template,
var_name=break_var,
for_stmt=node)
extra_test = templates.replace_as_expression(
'not var_name', var_name=break_var)
anno.setanno(node[1], 'extra_test', extra_test)
return node
def _replace_stack_call(self, node):
assert len(node.args) == 1
dtype = anno.getanno(node.args[0],
'element_type',
default=templates.replace_as_expression('None'))
template = """
ag__.list_stack(
target,
opts=ag__.ListStackOpts(
element_dtype=dtype,
original_call=orig_call))
"""
return templates.replace_as_expression(template,
dtype=dtype,
target=node.args[0],
orig_call=node.func)
def visit_IfExp(self, node):
template = """
ag__.utils.run_cond(test, lambda: (body,), lambda: (orelse,))
"""
desugared_ifexp = templates.replace_as_expression(
template, test=node.test, body=node.body, orelse=node.orelse)
return desugared_ifexp
def _as_function(self, func_name, args):
template = """
func_name(args)
"""
replacement = templates.replace_as_expression(
template, func_name=parser.parse_expression(func_name), args=args)
anno.setanno(replacement, SAFE_BOOLEAN_OPERAND, True)
return replacement
def _as_function(self, func_name, args):
template = """
func_name(args)
"""
replacement = templates.replace_as_expression(
template, func_name=parser.parse_expression(func_name), args=args)
anno.setanno(replacement, SAFE_BOOLEAN_OPERAND, True)
return replacement
def _replace_stack_call(self, node):
assert len(node.args) == 1
dtype = anno.getanno(
node.args[0],
'element_type',
default=templates.replace_as_expression('None'))
template = """
ag__.list_stack(
target,
opts=ag__.ListStackOpts(
element_dtype=dtype,
original_call=orig_call))
"""
return templates.replace_as_expression(
template,
dtype=dtype,
target=node.args[0],
orig_call=node.func)
def _convert_builtin(self, f, args, as_expression):
template = """
ag__.func(args)
"""
if as_expression:
return templates.replace_as_expression(
template, func=py_builtins.overload_of(f).__name__, args=args)
else:
return templates.replace(
template, func=py_builtins.overload_of(f).__name__, args=args)
def _convert_builtin(self, f, args, as_expression):
template = """
ag__.func(args)
"""
if as_expression:
return templates.replace_as_expression(
template, func=py_builtins.overload_of(f).__name__, args=args)
else:
return templates.replace(template,
func=py_builtins.overload_of(f).__name__,
args=args)
def _wrap_to_py_func_single_return(self, node, dtype):
# TODO(mdan): Properly handle varargs, etc.
template = """
ag__.utils.wrap_py_func(func, dtype, (args,), kwargs, False)
"""
return templates.replace_as_expression(
template,
func=node.func,
dtype=parser.parse_expression(dtype),
args=node.args,
kwargs=ast_util.keywords_to_dict(node.keywords))
def _wrap_to_py_func_single_return(self, node, dtype):
# TODO(mdan): Properly handle varargs, etc.
template = """
autograph_utils.wrap_py_func(func, dtype, (args,), kwargs, False)
"""
return templates.replace_as_expression(
template,
func=node.func,
dtype=parser.parse_expression(dtype),
args=node.args,
kwargs=ast_util.keywords_to_dict(node.keywords))
def _wrap_to_py_func_single_return(self, node, dtype):
# TODO (mdan): Properly handle varargs, etc. id:492
# https://github.com/imdone/tensorflow/issues/493
template = """
ag__.utils.wrap_py_func(func, dtype, (args,), kwargs, False)
"""
return templates.replace_as_expression(
template,
func=node.func,
dtype=parser.parse_expression(dtype),
args=node.args,
kwargs=ast_util.keywords_to_dict(node.keywords))
def visit_IfExp(self, node):
if anno.hasanno(node.test, anno.Basic.QN):
name_root = anno.getanno(node.test, anno.Basic.QN).ssf()
else:
name_root = 'ifexp'
true_fn_name = self._create_branch(node.body, '%s_true' % name_root)
false_fn_name = self._create_branch(node.orelse, '%s_false' % name_root)
return templates.replace_as_expression(
'ag__.utils.run_cond(test, true_fn_name, false_fn_name)',
test=node.test,
true_fn_name=true_fn_name,
false_fn_name=false_fn_name)
def _generate_pop_operation(self, original_call_node, pop_var_name):
assert isinstance(original_call_node.func, gast.Attribute)
if original_call_node.args:
pop_element = original_call_node.args[0]
else:
pop_element = parser.parse_expression('None')
# The call will be something like "target.pop()", and the dtype is hooked to
# target, hence the func.value.
# TODO(mdan): For lists of lists, this won't work.
# The reason why it won't work is because it's unclear how to annotate
# the list as a "list of lists with a certain element type" when using
# operations like `l.pop().pop()`.
dtype = self.get_definition_directive(
original_call_node.func.value,
directives.set_element_type,
'dtype',
default=templates.replace_as_expression('None'))
shape = self.get_definition_directive(
original_call_node.func.value,
directives.set_element_type,
'shape',
default=templates.replace_as_expression('None'))
template = """
target, pop_var_name = ag__.list_pop(
target, element,
opts=ag__.ListPopOpts(element_dtype=dtype, element_shape=shape))
"""
return templates.replace(
template,
target=original_call_node.func.value,
pop_var_name=pop_var_name,
element=pop_element,
dtype=dtype,
shape=shape)
def visit_Subscript(self, node):
node = self.generic_visit(node)
if not isinstance(node.slice, gast.Index):
return node
if not isinstance(node.ctx, gast.Load):
# Index writes are handled at a higher level, one at which the rvalue is
# also available.
return node
dtype = self.get_definition_directive(
node.value,
directives.set_element_type,
'dtype',
default=templates.replace_as_expression('None'))
template = """
ag__.get_item(
target,
key,
opts=ag__.GetItemOpts(element_dtype=dtype))
"""
return templates.replace_as_expression(
template, target=node.value, key=node.slice.value, dtype=dtype)
def _generate_pop_operation(self, original_call_node, pop_var_name):
assert isinstance(original_call_node.func, gast.Attribute)
if original_call_node.args:
pop_element = original_call_node.args[0]
else:
pop_element = parser.parse_expression('None')
# The call will be something like "target.pop()", and the dtype is hooked to
# target, hence the func.value.
# TODO(mdan): For lists of lists, this won't work.
# The reason why it won't work is because it's unclear how to annotate
# the list as a "list of lists with a certain element type" when using
# operations like `l.pop().pop()`.
dtype = self.get_definition_directive(
original_call_node.func.value,
directives.set_element_type,
'dtype',
default=templates.replace_as_expression('None'))
shape = self.get_definition_directive(
original_call_node.func.value,
directives.set_element_type,
'shape',
default=templates.replace_as_expression('None'))
template = """
target, pop_var_name = ag__.list_pop(
target, element,
opts=ag__.ListPopOpts(element_dtype=dtype, element_shape=shape))
"""
return templates.replace(template,
target=original_call_node.func.value,
pop_var_name=pop_var_name,
element=pop_element,
dtype=dtype,
shape=shape)
def visit_Subscript(self, node):
node = self.generic_visit(node)
if not isinstance(node.slice, gast.Index):
# TODO(mdan): It might make more sense to wave them through.
raise NotImplementedError('non-index slice')
if not isinstance(node.ctx, gast.Load):
# Index writes are handled at a higher level, one at which the rvalue is
# also available.
return node
dtype = anno.getanno(
node.value,
'element_type',
default=templates.replace_as_expression('None'))
template = """
ag__.get_item(
target,
key,
opts=ag__.GetItemOpts(element_dtype=dtype))
"""
return templates.replace_as_expression(
template, target=node.value, key=node.slice, dtype=dtype)
def visit_IfExp(self, node):
if anno.hasanno(node.test, anno.Basic.QN):
name_root = anno.getanno(node.test, anno.Basic.QN).ssf()
else:
name_root = 'ifexp'
true_fn_name = self._create_branch(node.body, '%s_true' % name_root)
false_fn_name = self._create_branch(node.orelse,
'%s_false' % name_root)
return templates.replace_as_expression(
'ag__.utils.run_cond(test, true_fn_name, false_fn_name)',
test=node.test,
true_fn_name=true_fn_name,
false_fn_name=false_fn_name)
def _empty_list(self, node):
if not anno.hasanno(node, 'element_type'):
raise NotImplementedError(
'type inference for empty lists is not yet supported; '
'use set_element_type(<list>, <dtype>) to continue')
dtype = anno.getanno(node, 'element_type')
if not isinstance(dtype, dtypes.DType):
# TODO(mdan): Allow non-TF dtypes?
# That would be consistent with the dynamic dispatch pattern, but
# we must make sure that doesn't become confusing.
raise NotImplementedError('element type "%s" not yet supported' % dtype)
dtype_name = dtype.name
# TODO(mdan): Does it ever make sense not to use tensor lists?
template = """
tf.TensorArray(tf.dtype_name, size=0, dynamic_size=True)
"""
return templates.replace_as_expression(template, dtype_name=dtype_name)
def _empty_list(self, node):
if not anno.hasanno(node, 'element_type'):
raise NotImplementedError(
'type inference for empty lists is not yet supported; '
'use utils.set_element_type(<list>, <dtype>) to continue')
dtype = anno.getanno(node, 'element_type')
if not isinstance(dtype, dtypes.DType):
# TODO(mdan): Allow non-TF dtypes?
# That would be consistent with the dynamic dispatch pattern, but
# we must make sure that doesn't become confusing.
raise NotImplementedError('element type "%s" not yet supported' %
dtype)
dtype_name = dtype.name
# TODO(mdan): Does it ever make sense not to use tensor lists?
template = """
tf.TensorArray(tf.dtype_name, size=0, dynamic_size=True)
"""
return templates.replace_as_expression(template, dtype_name=dtype_name)
def visit_For(self, node):
self.generic_visit(node.target)
self.generic_visit(node.iter)
scope = anno.getanno(node, NodeAnno.BODY_SCOPE)
break_var = self.context.namer.new_symbol('break_requested',
scope.referenced)
self.break_uses.append([False, break_var])
node.body = self._manual_visit_list(node.body)
if self.break_uses[-1][0]:
extra_cond = templates.replace_as_expression('not var_name',
var_name=break_var)
anno.setanno(node, 'extra_cond', extra_cond)
final_nodes = [self._create_break_init(), node]
else:
final_nodes = node
self.break_uses.pop()
for n in node.orelse:
self.generic_visit(n)
return final_nodes
def visit_For(self, node):
self.generic_visit(node.target)
self.generic_visit(node.iter)
scope = anno.getanno(node, NodeAnno.BODY_SCOPE)
break_var = self.context.namer.new_symbol('break_requested',
scope.referenced)
self.break_uses.append([False, break_var])
node.body = self._manual_visit_list(node.body)
if self.break_uses[-1][0]:
extra_cond = templates.replace_as_expression(
'not var_name', var_name=break_var)
anno.setanno(node, 'extra_cond', extra_cond)
final_nodes = [self._create_break_init(), node]
else:
final_nodes = node
self.break_uses.pop()
for n in node.orelse:
self.generic_visit(n)
return final_nodes
def visit_For(self, node):
scope = anno.getanno(node, NodeAnno.BODY_SCOPE)
break_var = self.ctx.namer.new_symbol('break_', scope.referenced)
node.target = self.visit(node.target)
node.iter = self.visit(node.iter)
node.body, break_used = self._track_body(node.body, break_var)
# A break in the else clause applies to the containing scope.
node.orelse = self.visit_block(node.orelse)
if break_used:
# Python's else clause only triggers if the loop exited cleanly (e.g.
# break did not trigger).
guarded_orelse = self._guard_if_present(node.orelse, break_var)
extra_test = templates.replace_as_expression(
'not var_name', var_name=break_var)
# The extra test is hidden in the AST, which will confuse the static
# analysis. To mitigate that, we insert a no-op statement that ensures
# the control variable is marked as used.
# TODO(mdan): Use a marker instead, e.g. ag__.condition_loop_on(var_name)
template = """
var_name = False
for target in iter_:
(var_name,)
body
else:
orelse
"""
node = templates.replace(
template,
var_name=break_var,
iter_=node.iter,
target=node.target,
body=node.body,
orelse=guarded_orelse)
anno.setanno(node[1], 'extra_test', extra_test)
return node
def visit_For(self, node):
scope = anno.getanno(node, NodeAnno.BODY_SCOPE)
break_var = self.context.namer.new_symbol('break_requested',
scope.referenced)
self.break_uses.append([False, break_var])
node = self.generic_visit(node)
if self.break_uses[-1][0]:
template = """
var_name = False
original_for
"""
node = templates.replace(template,
var_name=break_var,
original_for=node)
extra_cond = templates.replace_as_expression('not var_name',
var_name=break_var)
new_for_node = node[1]
anno.setanno(new_for_node, 'extra_cond', extra_cond)
self.break_uses.pop()
return node
def visit_For(self, node):
scope = anno.getanno(node, NodeAnno.BODY_SCOPE)
break_var = self.ctx.namer.new_symbol('break_', scope.referenced)
node.target = self.visit(node.target)
node.iter = self.visit(node.iter)
node.body, break_used = self._process_body(node.body, break_var)
# A break in the else clause applies to the containing scope.
node.orelse = self.visit_block(node.orelse)
if break_used:
# Python's else clause only triggers if the loop exited cleanly (e.g.
# break did not trigger).
guarded_orelse = self._guard_if_present(node.orelse, break_var)
extra_test = templates.replace_as_expression('not var_name',
var_name=break_var)
# The extra code is hidden in the AST, which will confuse the static
# analysis. To mitigate that, we insert a no-op statement that ensures
# the control variable is marked as used.
# TODO(mdan): Use a marker instead, e.g. ag__.condition_loop_on(var_name)
template = """
var_name = tf.constant(False)
for target in iter_:
(var_name,)
body
else:
orelse
"""
node = templates.replace(template,
var_name=break_var,
iter_=node.iter,
target=node.target,
body=node.body,
orelse=guarded_orelse)
anno.setanno(node[1], 'extra_test', extra_test)
return node
def _replace_pop_call(self, node):
# Expressions that use pop() are converted to a statement + expression.
#
# For example:
#
# print(target.pop())
#
# ... is converted to:
#
# target, target_pop = ag__.list_pop(target)
# print(target_pop)
#
# Here, we just generate the variable name and swap it in,
# and _generate_pop_operation will handle the rest.
#
# Multiple uses of pop() are allowed:
#
# print(tartget.pop(), target.pop())
# print(tartget.pop().pop())
#
assert isinstance(node.func, gast.Attribute)
scope = anno.getanno(node, NodeAnno.ARGS_SCOPE)
target_node = node.func.value
# Attempt to use a related name if one exists. Otherwise use something
# generic.
if anno.hasanno(target_node, anno.Basic.QN):
target_name = anno.getanno(target_node, anno.Basic.QN).ssf()
else:
target_name = 'list_'
pop_var_name = self.ctx.namer.new_symbol(target_name, scope.referenced)
pop_uses = self.get_local(POP_USES, [])
pop_uses.append((node, pop_var_name))
self.set_local(POP_USES, pop_uses)
return templates.replace_as_expression('var_name',
var_name=pop_var_name)
def visit_For(self, node):
scope = anno.getanno(node, NodeAnno.BODY_SCOPE)
break_var = self.context.namer.new_symbol('break_requested',
scope.referenced)
self.break_uses.append([False, break_var])
node = self.generic_visit(node)
if self.break_uses[-1][0]:
template = """
var_name = False
original_for
"""
node = templates.replace(
template,
var_name=break_var,
original_for=node)
extra_cond = templates.replace_as_expression(
'not var_name', var_name=break_var)
new_for_node = node[1]
anno.setanno(new_for_node, 'extra_cond', extra_cond)
self.break_uses.pop()
return node
def _replace_pop_call(self, node):
# Expressions that use pop() are converted to a statement + expression.
#
# For example:
#
# print(target.pop())
#
# ... is converted to:
#
# target, target_pop = ag__.list_pop(target)
# print(target_pop)
#
# Here, we just generate the variable name and swap it in,
# and _generate_pop_operation will handle the rest.
#
# Multiple uses of pop() are allowed:
#
# print(tartget.pop(), target.pop())
# print(tartget.pop().pop())
#
assert isinstance(node.func, gast.Attribute)
scope = anno.getanno(node, NodeAnno.ARGS_SCOPE)
target_node = node.func.value
# Attempt to use a related name if one exists. Otherwise use something
# generic.
if anno.hasanno(target_node, anno.Basic.QN):
target_name = anno.getanno(target_node, anno.Basic.QN).ssf()
else:
target_name = 'list_'
pop_var_name = self.ctx.namer.new_symbol(target_name, scope.referenced)
pop_uses = self.get_local(POP_USES, [])
pop_uses.append((node, pop_var_name))
self.set_local(POP_USES, pop_uses)
return templates.replace_as_expression('var_name', var_name=pop_var_name)
def visit_If(self, node):
node = self.generic_visit(node)
body_scope = anno.getanno(node, annos.NodeAnno.BODY_SCOPE)
orelse_scope = anno.getanno(node, annos.NodeAnno.ORELSE_SCOPE)
defined_in = anno.getanno(node, anno.Static.DEFINED_VARS_IN)
live_out = anno.getanno(node, anno.Static.LIVE_VARS_OUT)
modified_in_cond = body_scope.modified | orelse_scope.modified
returned_from_cond = set()
for s in modified_in_cond:
if s in live_out:
returned_from_cond.add(s)
elif s.is_composite():
# Special treatment for compound objects: if any of their owner entities
# are live, then they are outputs as well.
if any(owner in live_out for owner in s.owner_set):
returned_from_cond.add(s)
need_alias_in_body = body_scope.modified & defined_in
need_alias_in_orelse = orelse_scope.modified & defined_in
created_in_body = body_scope.modified & returned_from_cond - defined_in
created_in_orelse = orelse_scope.modified & returned_from_cond - defined_in
if created_in_body != created_in_orelse:
raise ValueError(
'if statement may not initialize all variables: the true branch'
' creates %s, while the false branch creates %s. Make sure all'
' these variables are initialized either in both'
' branches or before the if statement.' %
(self._fmt_symbol_list(created_in_body),
self._fmt_symbol_list(created_in_orelse)))
# Alias the closure variables inside the conditional functions, to allow
# the functions access to the respective variables.
# We will alias variables independently for body and orelse scope,
# because different branches might write different variables.
aliased_body_orig_names = tuple(need_alias_in_body)
aliased_orelse_orig_names = tuple(need_alias_in_orelse)
aliased_body_new_names = tuple(
self.ctx.namer.new_symbol(s.ssf(), body_scope.referenced)
for s in aliased_body_orig_names)
aliased_orelse_new_names = tuple(
self.ctx.namer.new_symbol(s.ssf(), orelse_scope.referenced)
for s in aliased_orelse_orig_names)
alias_body_map = dict(
zip(aliased_body_orig_names, aliased_body_new_names))
alias_orelse_map = dict(
zip(aliased_orelse_orig_names, aliased_orelse_new_names))
node_body = ast_util.rename_symbols(node.body, alias_body_map)
node_orelse = ast_util.rename_symbols(node.orelse, alias_orelse_map)
returned_from_cond = tuple(returned_from_cond)
if returned_from_cond:
if len(returned_from_cond) == 1:
# TODO(mdan): Move this quirk into the operator implementation.
cond_results = returned_from_cond[0]
else:
cond_results = gast.Tuple(
[s.ast() for s in returned_from_cond], None)
returned_from_body = tuple(
alias_body_map[s] if s in need_alias_in_body else s
for s in returned_from_cond)
returned_from_orelse = tuple(
alias_orelse_map[s] if s in need_alias_in_orelse else s
for s in returned_from_cond)
else:
# When the cond would return no value, we leave the cond called without
# results. That in turn should trigger the side effect guards. The
# branch functions will return a dummy value that ensures cond
# actually has some return value as well.
cond_results = None
# TODO(mdan): This doesn't belong here; it's specific to the operator.
returned_from_body = templates.replace_as_expression(
'tf.constant(1)')
returned_from_orelse = templates.replace_as_expression(
'tf.constant(1)')
body_name = self.ctx.namer.new_symbol('if_true', body_scope.referenced)
orelse_name = self.ctx.namer.new_symbol('if_false',
orelse_scope.referenced)
body_def = self._create_cond_branch(
body_name,
aliased_orig_names=aliased_body_orig_names,
aliased_new_names=aliased_body_new_names,
body=node_body,
returns=returned_from_body)
orelse_def = self._create_cond_branch(
orelse_name,
aliased_orig_names=aliased_orelse_orig_names,
aliased_new_names=aliased_orelse_new_names,
body=node_orelse,
returns=returned_from_orelse)
cond_expr = self._create_cond_expr(cond_results, node.test, body_name,
orelse_name)
return body_def + orelse_def + cond_expr
def visit_List(self, node):
node = self.generic_visit(node)
template = """
ag__.new_list(elements)
"""
return templates.replace_as_expression(template, elements=node)
def visit_If(self, node):
node = self.generic_visit(node)
body_scope = anno.getanno(node, annos.NodeAnno.BODY_SCOPE)
orelse_scope = anno.getanno(node, annos.NodeAnno.ORELSE_SCOPE)
defined_in = anno.getanno(node, anno.Static.DEFINED_VARS_IN)
live_out = anno.getanno(node, anno.Static.LIVE_VARS_OUT)
modified_in_cond = body_scope.modified | orelse_scope.modified
returned_from_cond = set()
for s in modified_in_cond:
if s in live_out:
returned_from_cond.add(s)
elif s.is_composite():
# Special treatment for compound objects: if any of their owner entities
# are live, then they are outputs as well.
if any(owner in live_out for owner in s.owner_set):
returned_from_cond.add(s)
need_alias_in_body = body_scope.modified & defined_in
need_alias_in_orelse = orelse_scope.modified & defined_in
created_in_body = body_scope.modified & returned_from_cond - defined_in
created_in_orelse = orelse_scope.modified & returned_from_cond - defined_in
if created_in_body != created_in_orelse:
raise ValueError(
'if statement may not initialize all variables: the true branch'
' creates %s, while the false branch creates %s. Make sure all'
' these variables are initialized either in both'
' branches or before the if statement.' %
(self._fmt_symbol_list(created_in_body),
self._fmt_symbol_list(created_in_orelse)))
# Alias the closure variables inside the conditional functions, to allow
# the functions access to the respective variables.
# We will alias variables independently for body and orelse scope,
# because different branches might write different variables.
aliased_body_orig_names = tuple(need_alias_in_body)
aliased_orelse_orig_names = tuple(need_alias_in_orelse)
aliased_body_new_names = tuple(
self.ctx.namer.new_symbol(s.ssf(), body_scope.referenced)
for s in aliased_body_orig_names)
aliased_orelse_new_names = tuple(
self.ctx.namer.new_symbol(s.ssf(), orelse_scope.referenced)
for s in aliased_orelse_orig_names)
alias_body_map = dict(zip(aliased_body_orig_names, aliased_body_new_names))
alias_orelse_map = dict(
zip(aliased_orelse_orig_names, aliased_orelse_new_names))
node_body = ast_util.rename_symbols(node.body, alias_body_map)
node_orelse = ast_util.rename_symbols(node.orelse, alias_orelse_map)
returned_from_cond = tuple(returned_from_cond)
if returned_from_cond:
if len(returned_from_cond) == 1:
# TODO(mdan): Move this quirk into the operator implementation.
cond_results = returned_from_cond[0]
else:
cond_results = gast.Tuple([s.ast() for s in returned_from_cond], None)
returned_from_body = tuple(
alias_body_map[s] if s in need_alias_in_body else s
for s in returned_from_cond)
returned_from_orelse = tuple(
alias_orelse_map[s] if s in need_alias_in_orelse else s
for s in returned_from_cond)
else:
# When the cond would return no value, we leave the cond called without
# results. That in turn should trigger the side effect guards. The
# branch functions will return a dummy value that ensures cond
# actually has some return value as well.
cond_results = None
# TODO(mdan): This doesn't belong here; it's specific to the operator.
returned_from_body = templates.replace_as_expression('1')
returned_from_orelse = templates.replace_as_expression('1')
body_name = self.ctx.namer.new_symbol('if_true', body_scope.referenced)
orelse_name = self.ctx.namer.new_symbol('if_false', orelse_scope.referenced)
body_def = self._create_cond_branch(
body_name,
aliased_orig_names=aliased_body_orig_names,
aliased_new_names=aliased_body_new_names,
body=node_body,
returns=returned_from_body)
orelse_def = self._create_cond_branch(
orelse_name,
aliased_orig_names=aliased_orelse_orig_names,
aliased_new_names=aliased_orelse_new_names,
body=node_orelse,
returns=returned_from_orelse)
cond_expr = self._create_cond_expr(cond_results, node.test, body_name,
orelse_name)
return body_def + orelse_def + cond_expr
def visit_List(self, node):
node = self.generic_visit(node)
template = """
ag__.new_list(elements)
"""
return templates.replace_as_expression(template, elements=node)
