From 2ccb446ecf92dc01df0a60d965561944d82e405b Mon Sep 17 00:00:00 2001 From: "A. Unique TensorFlower" Date: Wed, 6 Feb 2019 08:00:49 -0800 Subject: [PATCH] Handle nested code-blocks correctly in lowering passes. This fixes issues with the early-return and continue passes which were generating incorrect code when interleaved with "with" and "try" statements. PiperOrigin-RevId: 232675424 --- .../converters/continue_statements.py | 22 +++++++++++++++ .../autograph/converters/return_statements.py | 27 +++++++++++++++++++ 2 files changed, 49 insertions(+) diff --git a/tensorflow/python/autograph/converters/continue_statements.py b/tensorflow/python/autograph/converters/continue_statements.py index c3b6679e383..725e05360d2 100644 --- a/tensorflow/python/autograph/converters/continue_statements.py +++ b/tensorflow/python/autograph/converters/continue_statements.py @@ -121,6 +121,28 @@ class ContinueCanonicalizationTransformer(converter.Base): node.orelse = self.visit_block(node.orelse) return node + def visit_With(self, node): + node.items = self.visit_block(node.items) + node.body = self.visit_block(node.body, + after_visit=self._postprocess_statement) + return node + + def visit_Try(self, node): + node.body = self.visit_block(node.body, + after_visit=self._postprocess_statement) + node.orelse = self.visit_block(node.orelse, + after_visit=self._postprocess_statement) + # In Python 3.8 and later continue is allowed in finally blocks + node.finalbody = self.visit_block(node.finalbody, + after_visit=self._postprocess_statement) + node.handlers = self.visit_block(node.handlers) + return node + + def visit_ExceptHandler(self, node): + node.body = self.visit_block(node.body, + after_visit=self._postprocess_statement) + return node + def transform(node, ctx): transformer = ContinueCanonicalizationTransformer(ctx) diff --git a/tensorflow/python/autograph/converters/return_statements.py b/tensorflow/python/autograph/converters/return_statements.py index 6cf6f266496..723acab908a 100644 --- a/tensorflow/python/autograph/converters/return_statements.py +++ b/tensorflow/python/autograph/converters/return_statements.py @@ -115,6 +115,21 @@ class ConditionalReturnRewriter(converter.Base): anno.setanno(node, STMT_DEFINITELY_RETURNS, True) return node + def visit_Try(self, node): + # We could decide whether a 'try' DEFINITELY_RETURNS based on its components + # It is not clear whether we want to do anything with this given + # a 'try' is likely to throw an exception in some circumstances. + node.body, _ = self._visit_statement_block(node, node.body) + node.orelse, _ = self._visit_statement_block(node, node.orelse) + node.finalbody, _ = self._visit_statement_block(node, node.finalbody) + node.handlers = self.visit_block(node.handlers) + return node + + def visit_ExceptHandler(self, node): + # To determine whether `try` DEFINITELY_RETURNS we need to revisit this. + node.body, _ = self._visit_statement_block(node, node.body) + return node + def visit_If(self, node): node.test = self.visit(node.test) @@ -317,6 +332,17 @@ class ReturnStatementsTransformer(converter.Base): node.body = self._visit_statement_block(node, node.body) return node + def visit_Try(self, node): + node.body = self._visit_statement_block(node, node.body) + node.orelse = self._visit_statement_block(node, node.orelse) + node.finalbody = self._visit_statement_block(node, node.finalbody) + node.handlers = self.visit_block(node.handlers) + return node + + def visit_ExceptHandler(self, node): + node.body = self._visit_statement_block(node, node.body) + return node + def visit_If(self, node): node.test = self.visit(node.test) node.body = self._visit_statement_block(node, node.body) @@ -382,5 +408,6 @@ def transform(node, ctx, default_to_null_return=True): transformer = ReturnStatementsTransformer( ctx, default_to_null_return=default_to_null_return) node = transformer.visit(node) + transformer.debug_print_src(node) return node