Some fixes to returning from generators (mypyc/mypyc#688)

msullivan · web-flow · commit 416732e43e6d · 2019-08-19T17:40:28.000-07:00
* Don't set a traceback frame when raising StopIteration. This seems
   to match CPython and also speeds up generators a *lot*.
 * Use `_PyGen_SetStopIterationValue` to set the iteration value since
   `PyErr_SetObject` doesn't work right if the value is a tuple.
diff --git a/mypyc/genops.py b/mypyc/genops.py
@@ -92,7 +92,7 @@ def f(x: int) -> int:
 from mypyc.ops_exc import (
     raise_exception_op, raise_exception_with_tb_op, reraise_exception_op,
     error_catch_op, restore_exc_info_op, exc_matches_op, get_exc_value_op,
-    get_exc_info_op, keep_propagating_op,
+    get_exc_info_op, keep_propagating_op, set_stop_iteration_value,
 )
 from mypyc.genops_for import ForGenerator, ForRange, ForList, ForIterable, ForEnumerate, ForZip
 from mypyc.rt_subtype import is_runtime_subtype
@@ -901,8 +901,13 @@ def gen_return(self, builder: 'IRBuilder', value: Value, line: int) -> None:
         # StopIteration isn't caught by except blocks inside of the generator function.
         builder.error_handlers.append(None)
         builder.goto_new_block()
-        builder.add(RaiseStandardError(RaiseStandardError.STOP_ITERATION, value,
-                                       line))
+        # Skip creating a traceback frame when we raise here, because
+        # we don't care about the traceback frame and it is kind of
+        # expensive since raising StopIteration is an extremely common case.
+        # Also we call a special internal function to set StopIteration instead of
+        # using RaiseStandardError because the obvious thing doesn't work if the
+        # value is a tuple (???).
+        builder.primitive_op(set_stop_iteration_value, [value], NO_TRACEBACK_LINE_NO)
         builder.add(Unreachable())
         builder.error_handlers.pop()
 
diff --git a/mypyc/lib-rt/pythonsupport.h b/mypyc/lib-rt/pythonsupport.h
@@ -328,6 +328,44 @@ _PyDict_GetItemStringWithError(PyObject *v, const char *key)
 #define CPyUnicode_EqualToASCIIString(x, y) _PyUnicode_EqualToASCIIString(x, y)
 #endif
 
+// Adapted from genobject.c in Python 3.7.2
+// Copied because it wasn't in 3.5.2 and it is undocumented anyways.
+/*
+ * Set StopIteration with specified value.  Value can be arbitrary object
+ * or NULL.
+ *
+ * Returns 0 if StopIteration is set and -1 if any other exception is set.
+ */
+static int
+CPyGen_SetStopIterationValue(PyObject *value)
+{
+    PyObject *e;
+
+    if (value == NULL ||
+        (!PyTuple_Check(value) && !PyExceptionInstance_Check(value)))
+    {
+        /* Delay exception instantiation if we can */
+        PyErr_SetObject(PyExc_StopIteration, value);
+        return 0;
+    }
+    /* Construct an exception instance manually with
+     * PyObject_CallFunctionObjArgs and pass it to PyErr_SetObject.
+     *
+     * We do this to handle a situation when "value" is a tuple, in which
+     * case PyErr_SetObject would set the value of StopIteration to
+     * the first element of the tuple.
+     *
+     * (See PyErr_SetObject/_PyErr_CreateException code for details.)
+     */
+    e = PyObject_CallFunctionObjArgs(PyExc_StopIteration, value, NULL);
+    if (e == NULL) {
+        return -1;
+    }
+    PyErr_SetObject(PyExc_StopIteration, e);
+    Py_DECREF(e);
+    return 0;
+}
+
 #ifdef __cplusplus
 }
 #endif
diff --git a/mypyc/ops_exc.py b/mypyc/ops_exc.py
@@ -17,6 +17,13 @@
     format_str = 'raise_exception({args[0]}); {dest} = 0',
     emit=simple_emit('CPy_Raise({args[0]}); {dest} = 0;'))
 
+set_stop_iteration_value = custom_op(
+    arg_types=[object_rprimitive],
+    result_type=bool_rprimitive,
+    error_kind=ERR_FALSE,
+    format_str = 'set_stop_iteration_value({args[0]}); {dest} = 0',
+    emit=simple_emit('CPyGen_SetStopIterationValue({args[0]}); {dest} = 0;'))
+
 raise_exception_with_tb_op = custom_op(
     arg_types=[object_rprimitive, object_rprimitive, object_rprimitive],
     result_type=bool_rprimitive,
diff --git a/test-data/run.test b/test-data/run.test
@@ -3497,7 +3497,8 @@ else:
     assert False
 
 [case testYield]
-from typing import Generator, Iterable, Union
+from typing import Generator, Iterable, Union, Tuple
+
 def yield_three_times() -> Iterable[int]:
     yield 1
     yield 2
@@ -3567,8 +3568,12 @@ class A(object):
     def generator(self) -> Iterable[int]:
         yield self.x
 
+def return_tuple() -> Generator[int, None, Tuple[int, int]]:
+    yield 0
+    return 1, 2
+
 [file driver.py]
-from native import yield_three_times, yield_twice_and_return, yield_while_loop, yield_for_loop, yield_with_except, complex_yield, yield_with_default, A
+from native import yield_three_times, yield_twice_and_return, yield_while_loop, yield_for_loop, yield_with_except, complex_yield, yield_with_default, A, return_tuple
 from testutil import run_generator
 
 assert run_generator(yield_three_times()) == ((1, 2, 3), None)
@@ -3579,6 +3584,7 @@ assert run_generator(yield_with_except()) == ((10,), None)
 assert run_generator(complex_yield(5, 'foo', 1.0)) == (('2 foo', 3, '4 foo'), 1.0)
 assert run_generator(yield_with_default()) == ((), None)
 assert run_generator(A(0).generator()) == ((0,), None)
+assert run_generator(return_tuple()) == ((0,), (1, 2))
 
 for i in yield_twice_and_return():
     print(i)
