def test_basic_break_for_loop(self):
def test_fn(a):
v = []
for x in a:
x -= 1
if x % 2 == 0:
break
v.append(x)
return v
# The break is incompletely canonicalized for for loops. Everything is
# in place except for the condition verification.
def test_equiv_fn(a):
v = []
for x in a:
x -= 1
if x % 2 == 0:
continue
v.append(x)
return v
node = self._parse_and_analyze(test_fn, {})
node = break_canonicalization.transform(node, TestNamer())
result = compiler.ast_to_object(node)
# The break is incompletely canonicalized. Everything is in place, but
# the loop does not break.
self.assertEqual(test_equiv_fn([]), result.test_fn([]))
self.assertEqual(test_equiv_fn([1]), result.test_fn([1]))
self.assertEqual(test_equiv_fn([2]), result.test_fn([2]))
self.assertEqual(test_equiv_fn([1, 2, 3, 4]),
result.test_fn([1, 2, 3, 4]))
def test_continue_deeply_nested(self):
def test_fn(x):
v = []
u = []
w = []
while x > 0:
x -= 1
if x % 2 == 0:
if x % 3 != 0:
u.append(x)
else:
w.append(x)
continue
v.append(x)
return v, u, w
node = self._parse_and_analyze(test_fn, {})
node = break_canonicalization.transform(node, TestNamer())
result = compiler.ast_to_object(node)
self.assertEqual(test_fn(0), result.test_fn(0))
self.assertEqual(test_fn(1), result.test_fn(1))
self.assertEqual(test_fn(2), result.test_fn(2))
self.assertEqual(test_fn(3), result.test_fn(3))
self.assertEqual(test_fn(4), result.test_fn(4))
def test_continue_deeply_nested(self):
def test_fn(x):
v = []
u = []
w = []
while x > 0:
x -= 1
if x % 2 == 0:
if x % 3 != 0:
u.append(x)
else:
w.append(x)
continue
v.append(x)
return v, u, w
node = self._parse_and_analyze(test_fn, {})
node = break_canonicalization.transform(node, TestNamer())
result = compiler.ast_to_object(node)
self.assertEqual(test_fn(0), result.test_fn(0))
self.assertEqual(test_fn(1), result.test_fn(1))
self.assertEqual(test_fn(2), result.test_fn(2))
self.assertEqual(test_fn(3), result.test_fn(3))
self.assertEqual(test_fn(4), result.test_fn(4))
def test_basic_break_for_loop(self):
def test_fn(a):
v = []
for x in a:
x -= 1
if x % 2 == 0:
break
v.append(x)
return v
# The break is incompletely canonicalized for for loops. Everything is
# in place except for the condition verification.
def test_equiv_fn(a):
v = []
for x in a:
x -= 1
if x % 2 == 0:
continue
v.append(x)
return v
node = self._parse_and_analyze(test_fn, {})
node = break_canonicalization.transform(node, TestNamer())
result = compiler.ast_to_object(node)
# The break is incompletely canonicalized. Everything is in place, but
# the loop does not break.
self.assertEqual(test_equiv_fn([]), result.test_fn([]))
self.assertEqual(test_equiv_fn([1]), result.test_fn([1]))
self.assertEqual(test_equiv_fn([2]), result.test_fn([2]))
self.assertEqual(test_equiv_fn([1, 2, 3, 4]), result.test_fn([1, 2, 3, 4]))
def node_to_graph(node, namer, namespace, value_hints, nocompile_decorators):
"""Convert Python code to equivalent TF graph mode code.
Args:
node: A Python AST node representing the code to convert.
namer: A naming.Namer object.
namespace: Dict mapping symbol names to their corresponding live objects.
value_hints: A dict containing value hints for symbols like function
parameters.
nocompile_decorators: A tuple containing decorators to be stripped from
functions during conversion.
Returns:
A tuple (node, deps):
* node: A Python ast node, representing the converted code.
* deps: A set of strings, the fully qualified names of object
dependencies that this node has.
"""
# TODO(mdan): Verify arguments for correctness.
# TODO(mdan): Factor out common elements.
# These include:
# * keeping track of symbols that have been created
# * marking nodes (e.g. py_func wrappers) to suppress further processing
# * code move between blocks
# * insertion of new global references
# * visiting blocks in transformers
# Certain steps, especially canonicalization, insert new symbols into the
# tree, which must be accounted. Although less efficient, it is most robust
# to re-run the analysis.
node = _static_analysis_pass(node, namespace, value_hints)
node = decorators.transform(node, nocompile_decorators)
node = break_canonicalization.transform(node, namer)
# Note: sequencing continue canonicalization before for loop one avoids
# dealing with the extra loop increment operation that the for
# canonicalization creates.
node = continue_canonicalization.transform(node, namer)
namespace['len'] = len
node = _static_analysis_pass(node, namespace, value_hints)
node = for_canonicalization.transform(node, namer)
# for_canonicalization may insert new global references.
node = builtin_functions.transform(node)
# builtin_functions may insert new global references.
namespace['print'] = print
node = _static_analysis_pass(node, namespace, value_hints)
node = print_functions.transform(node)
node = call_trees.transform(node, namer, namespace,
config.DEFAULT_UNCOMPILED_MODULES,
nocompile_decorators)
node = control_flow.transform(node, namer)
node = logical_expressions.transform(node)
node = side_effect_guards.transform(node, namer)
return node
def test_basic_break(self):
def test_fn(x):
v = []
while x > 0:
x -= 1
if x % 2 == 0:
break
v.append(x)
return v
node = self._parse_and_analyze(test_fn, {})
node = break_canonicalization.transform(node, TestNamer())
result = compiler.ast_to_object(node)
self.assertEqual(test_fn(0), result.test_fn(0))
self.assertEqual(test_fn(1), result.test_fn(1))
self.assertEqual(test_fn(2), result.test_fn(2))
self.assertEqual(test_fn(3), result.test_fn(3))
self.assertEqual(test_fn(4), result.test_fn(4))
def test_basic_break(self):
def test_fn(x):
v = []
while x > 0:
x -= 1
if x % 2 == 0:
break
v.append(x)
return v
node = self._parse_and_analyze(test_fn, {})
node = break_canonicalization.transform(node, TestNamer())
result = compiler.ast_to_object(node)
self.assertEqual(test_fn(0), result.test_fn(0))
self.assertEqual(test_fn(1), result.test_fn(1))
self.assertEqual(test_fn(2), result.test_fn(2))
self.assertEqual(test_fn(3), result.test_fn(3))
self.assertEqual(test_fn(4), result.test_fn(4))