Spill the old exception in iterators with excepts (mypyc/mypyc#683)

msullivan · web-flow · commit 148a1d0e9613 · 2019-08-17T20:05:24.000-07:00
diff --git a/mypyc/genops.py b/mypyc/genops.py
@@ -953,12 +953,12 @@ class ExceptNonlocalControl(CleanupNonlocalControl):
     Just makes sure that sys.exc_info always gets restored when we leave.
     This is super annoying.
     """
-    def __init__(self, outer: NonlocalControl, saved: Value) -> None:
+    def __init__(self, outer: NonlocalControl, saved: Union[Value, AssignmentTarget]) -> None:
         super().__init__(outer)
         self.saved = saved
 
     def gen_cleanup(self, builder: 'IRBuilder', line: int) -> None:
-        builder.primitive_op(restore_exc_info_op, [self.saved], line)
+        builder.primitive_op(restore_exc_info_op, [builder.read(self.saved)], line)
 
 
 class FinallyNonlocalControl(CleanupNonlocalControl):
@@ -3550,7 +3550,7 @@ def visit_try_except(self,
         # exception is raised, based on the exception in exc_info.
         self.error_handlers.append(double_except_block)
         self.activate_block(except_entry)
-        old_exc = self.primitive_op(error_catch_op, [], line)
+        old_exc = self.maybe_spill(self.primitive_op(error_catch_op, [], line))
         # Compile the except blocks with the nonlocal control flow overridden to clear exc_info
         self.nonlocal_control.append(
             ExceptNonlocalControl(self.nonlocal_control[-1], old_exc))
@@ -3583,14 +3583,14 @@ def visit_try_except(self,
         # restore the saved exc_info information and continue propagating
         # the exception if it exists.
         self.activate_block(cleanup_block)
-        self.primitive_op(restore_exc_info_op, [old_exc], line)
+        self.primitive_op(restore_exc_info_op, [self.read(old_exc)], line)
         self.goto(exit_block)
 
         # Cleanup for if we leave except through a raised exception:
         # restore the saved exc_info information and continue propagating
         # the exception.
         self.activate_block(double_except_block)
-        self.primitive_op(restore_exc_info_op, [old_exc], line)
+        self.primitive_op(restore_exc_info_op, [self.read(old_exc)], line)
         self.primitive_op(keep_propagating_op, [], NO_TRACEBACK_LINE_NO)
         self.add(Unreachable())
 
diff --git a/test-data/run.test b/test-data/run.test
@@ -1172,7 +1172,7 @@ Traceback (most recent call last):
 Exception
 
 [case testTryExcept]
-from typing import Any
+from typing import Any, Iterator
 import wrapsys
 def g(b: bool) -> None:
     try:
@@ -1250,6 +1250,12 @@ def m(x: object) -> int:
         return -1
     return st + 1
 
+def iter_exception() -> Iterator[str]:
+    try:
+        r(0)
+    except KeyError as e:
+        yield 'lol'
+
 [file wrapsys.py]
 # This is a gross hack around some limitations of the test system/mypyc.
 from typing import Any
@@ -1258,9 +1264,9 @@ def exc_info() -> Any:
     return sys.exc_info()  # type: ignore
 
 [file driver.py]
-import traceback
-import sys
-from native import g, f, h, i, j, k, l, m
+import sys, traceback
+from native import g, f, h, i, j, k, l, m, iter_exception
+from testutil import assertRaises
 print("== i ==")
 try:
     i()
@@ -1297,6 +1303,10 @@ print("== l ==")
 l()
 
 m('lol')
+
+with assertRaises(IndexError):
+    list(iter_exception())
+
 [out]
 == i ==
 <class 'IndexError'>
