Fix off-by-one error in the statement directive preprocessor.

PiperOrigin-RevId: 242669425
This commit is contained in:
Dan Moldovan 2019-04-09 08:17:08 -07:00 committed by TensorFlower Gardener
parent 3d69fd003d
commit 374d665e0f
2 changed files with 10 additions and 1 deletions

View File

@ -92,7 +92,7 @@ class DirectivesTransformer(converter.Base):
return call_node
def _process_statement_directive(self, call_node, directive):
if self.local_scope_level < 1:
if self.local_scope_level < 2:
raise ValueError(
'"%s" must be used inside a statement' % directive.__name__)
target = self.get_local(ENCLOSING_LOOP)

View File

@ -74,6 +74,15 @@ class DirectivesTest(converter_testing.TestCase):
self.assertEqual(d['back_prop'].id, 'a')
self.assertNotIn('swap_memory', d)
def test_loop_target_with_no_loop(self):
def test_fn():
directives.set_loop_options()
node, ctx = self.prepare(test_fn, {'directives': directives})
with self.assertRaisesRegexp(ValueError, 'must be used inside a statement'):
node = directives_converter.transform(node, ctx)
def test_invalid_default(self):
def invalid_directive(valid_arg, invalid_default=object()):