Ensure return statements don't bypass finally

Instead of translating a Python return directly into a Go return
statement, set a reserved variable πR to the return value and continue
up the checkpoint stack. When a function is done and it's not a
propagating an exception it will return πR.

This fixes #364

Author: trotterdylan

compiler/stmt.py

@@ -139,6 +139,7 @@ def visit_ClassDef(self, node):
         self.writer.write_temp_decls(body_visitor.block)
         self.writer.write_block(body_visitor.block,
                                 body_visitor.writer.getvalue())
+        self.writer.write('return nil, nil')
       tmpl = textwrap.dedent("""\
           }).Eval(πF, πF.Globals(), nil, nil)
           if πE != nil {
@@ -355,9 +356,8 @@ def visit_Return(self, node):
       raise util.ParseError(node, 'returning a value in a generator function')
     if node.value:
       with self.visit_expr(node.value) as value:
-        self.writer.write('return {}, nil'.format(value.expr))
-    else:
-      self.writer.write('return nil, nil')
+        self.writer.write('πR = {}'.format(value.expr))
+    self.writer.write('continue')
 
   def visit_Try(self, node):
     # The general structure generated by this method is shown below:
@@ -404,39 +404,38 @@ def visit_Try(self, node):
 
     with self.block.alloc_temp('*πg.BaseException') as exc:
       if except_label:
-        if (len(node.handlers) == 1 and not node.handlers[0].type and
-            not node.orelse):
-          # When there's just a bare except, no dispatch is required.
-          self._write_except_block(except_label, exc.expr, node.handlers[0])
+        with self.block.alloc_temp('*πg.Traceback') as tb:
+          self.writer.write_label(except_label)
+          self.writer.write_tmpl(textwrap.dedent("""\
+              if πE == nil {
+                continue
+              }
+              πE = nil
+              $exc, $tb = πF.ExcInfo()"""), exc=exc.expr, tb=tb.expr)
+          handler_labels = self._write_except_dispatcher(
+              exc.expr, tb.expr, node.handlers)
+
+        # Write the bodies of each of the except handlers.
+        for handler_label, except_node in zip(handler_labels, node.handlers):
+          self._write_except_block(handler_label, exc.expr, except_node)
           if node.finalbody:
             self.writer.write('πF.PopCheckpoint()')  # finally_label
           self.writer.write('goto Label{}'.format(finally_label))
-        else:
-          with self.block.alloc_temp('*πg.Traceback') as tb:
-            self.writer.write_label(except_label)
-            self.writer.write('{}, {} = πF.ExcInfo()'.format(exc.expr, tb.expr))
-            handler_labels = self._write_except_dispatcher(
-                exc.expr, tb.expr, node.handlers)
-
-          # Write the bodies of each of the except handlers.
-          for handler_label, except_node in zip(handler_labels, node.handlers):
-            self._write_except_block(handler_label, exc.expr, except_node)
-            if node.finalbody:
-              self.writer.write('πF.PopCheckpoint()')  # finally_label
-            self.writer.write('goto Label{}'.format(finally_label))
 
       # Write the finally body.
       self.writer.write_label(finally_label)
       if node.finalbody:
         with self.block.alloc_temp('*πg.Traceback') as tb:
-          self.writer.write('πE = nil')
           self.writer.write('{}, {} = πF.RestoreExc(nil, nil)'.format(
               exc.expr, tb.expr))
           self._visit_each(node.finalbody)
           self.writer.write_tmpl(textwrap.dedent("""\
               if $exc != nil {
               \tπE = πF.Raise($exc.ToObject(), nil, $tb.ToObject())
               \tcontinue
+              }
+              if πR != nil {
+              \tcontinue
               }"""), exc=exc.expr, tb=tb.expr)
 
   def visit_While(self, node):
@@ -532,6 +531,9 @@ def visit_With(self, node):
             if $exc != nil && $swallow_exc != true {
             \tπE = πF.Raise(nil, nil, nil)
             \tcontinue
+            }
+            if πR != nil {
+            \tcontinue
             }"""), exc=exc.expr, swallow_exc=swallow_exc_bool.expr)
 
   def visit_function_inline(self, node):
@@ -584,14 +586,25 @@ def visit_function_inline(self, node):
             self.writer.write(fmt.format(
                 util.adjust_local_name(var.name), var.init_expr))
         self.writer.write_temp_decls(func_block)
+        self.writer.write('var πR *πg.Object; _ = πR')
+        self.writer.write('var πE *πg.BaseException; _ = πE')
         if func_block.is_generator:
-          self.writer.write('return πg.NewGenerator(πF, func(πSent *πg.Object) '
-                            '(*πg.Object, *πg.BaseException) {')
+          self.writer.write(
+              'return πg.NewGenerator(πF, func(πSent *πg.Object) '
+              '(*πg.Object, *πg.BaseException) {')
           with self.writer.indent_block():
             self.writer.write_block(func_block, visitor.writer.getvalue())
+            self.writer.write('return nil, πE')
           self.writer.write('}).ToObject(), nil')
         else:
           self.writer.write_block(func_block, visitor.writer.getvalue())
+          self.writer.write(textwrap.dedent("""\
+              if πE != nil {
+              \tπR = nil
+              } else if πR == nil {
+              \tπR = πg.None
+              }
+              return πR, πE"""))
       self.writer.write('}), πF.Globals()).ToObject()')
     return result
 
@@ -733,7 +746,6 @@ def _write_except_block(self, label, exc, except_node):
       self.block.bind_var(self.writer, except_node.name.id,
                           '{}.ToObject()'.format(exc))
     self._visit_each(except_node.body)
-    self.writer.write('πE = nil')
     self.writer.write('πF.RestoreExc(nil, nil)')
 
   def _write_except_dispatcher(self, exc, tb, handlers):

compiler/util.py

@@ -87,7 +87,6 @@ def write_block(self, block_, body):
       block_: The Block object representing the code block.
       body: String containing Go code making up the body of the code block.
     """
-    self.write('var πE *πg.BaseException; _ = πE')
     self.write('for ; πF.State() >= 0; πF.PopCheckpoint() {')
     with self.indent_block():
       self.write('switch πF.State() {')
@@ -99,9 +98,7 @@ def write_block(self, block_, body):
       # Assume that body is aligned with goto labels.
       with self.indent_block(-1):
         self.write(body)
-      self.write('return nil, nil')
     self.write('}')
-    self.write('return nil, πE')
 
   def write_import_block(self, imports):
     if not imports:

compiler/util_test.py

@@ -43,7 +43,6 @@ def testWriteBlock(self):
     output = writer.getvalue()
     dispatch = 'switch πF.State() {\n\tcase 0:\n\tdefault: panic'
     self.assertIn(dispatch, output)
-    self.assertIn('return nil, nil\n}', output)
 
   def testWriteImportBlockEmptyImports(self):
     writer = util.Writer()

testing/try_test.py

@@ -220,3 +220,15 @@ def f():
 
 
 foo()
+
+
+# Return statement should not bypass the finally.
+def foo():
+  try:
+    return 1
+  finally:
+    return 2
+  return 3
+
+
+assert foo() == 2

testing/with_test.py

@@ -182,16 +182,30 @@ def __exit__(self, *args):
 assert j == 3
 
 
-# This checks for a bug where a with clause inside an except body raises an
-# exception because it was checking ExcInfo() to determine whether an exception
-# occurred.
 class Foo(object):
+  exited = False
   def __enter__(self):
     pass
   def __exit__(self, *args):
-    pass
+    self.exited = True
+
+
+# This checks for a bug where a with clause inside an except body raises an
+# exception because it was checking ExcInfo() to determine whether an exception
+# occurred.
 try:
   raise AssertionError
 except:
-  with Foo():
+  foo = Foo()
+  with foo:
     pass
+  assert foo.exited
+
+
+# Return statement should not bypass the with exit handler.
+foo = Foo()
+def bar():
+  with foo:
+    return
+bar()
+assert foo.exited

tools/grumpc

@@ -83,22 +83,22 @@ def main(args):
   package_name = args.modname.split('.')[-1]
   writer.write('package {}'.format(package_name))
   writer.write_import_block(mod_block.imports)
-
-  writer.write('func initModule(πF *πg.Frame, '
-               '_ []*πg.Object) (*πg.Object, *πg.BaseException) {')
-  with writer.indent_block():
+  writer.write_tmpl(textwrap.dedent("""\
+    var Code *πg.Code
+    func init() {
+    \tCode = πg.NewCode("<module>", $script, nil, 0, func(πF *πg.Frame, _ []*πg.Object) (*πg.Object, *πg.BaseException) {
+    \t\tvar πR *πg.Object; _ = πR
+    \t\tvar πE *πg.BaseException; _ = πE"""), script=util.go_str(args.script))
+  with writer.indent_block(2):
     for s in sorted(mod_block.strings):
       writer.write('ß{} := πg.InternStr({})'.format(s, util.go_str(s)))
     writer.write_temp_decls(mod_block)
     writer.write_block(mod_block, visitor.writer.getvalue())
-  writer.write('}')
-  writer.write('var Code *πg.Code')
-
   writer.write_tmpl(textwrap.dedent("""\
-    func init() {
-    \tCode = πg.NewCode("<module>", $script, nil, 0, initModule)
+    \t\treturn nil, πE
+    \t})
     \tπg.RegisterModule($modname, Code)
-    }"""), script=util.go_str(args.script), modname=util.go_str(args.modname))
+    }"""), modname=util.go_str(args.modname))
   return 0