[mypyc] Fix tracebacks of exceptions passed to throw() (#22287)

When calling single-argument `throw(ex)` on a coroutine or generator,
the `traceback` argument is `None`, so passing it to `PyErr_Restore`
clears the existing exception traceback. Recover the traceback from the
exception instead to match interpreted Python.

Don't add a traceback entry when forwarding an exception injected
through `throw()` at `await`/`yield from`. Otherwise, the child coro/gen
would see the parent's frame in its traceback. The parent's entry is
still added when propagating the exception, so this also avoids a
duplicated entry.

Author: p-sawicki

mypyc/irbuild/statement.py

@@ -1439,8 +1439,12 @@ def native_step(sent: Value) -> Value:
     builder.goto_and_activate(loop_block)
 
     def try_body() -> None:
+        # Traceback will be added if yield_from_except_op propagates an error.
+        # Skip it here to not have a duplicated entry.
         builder.assign(
-            received_reg, emit_yield(builder, builder.read(to_yield_reg, line), line), line
+            received_reg,
+            emit_yield(builder, builder.read(to_yield_reg, line), NO_TRACEBACK_LINE_NO),
+            line,
         )
 
     def except_body() -> None:

mypyc/lib-rt/exc_ops.c

@@ -83,13 +83,16 @@ void CPyErr_SetObjectAndTraceback(PyObject *type, PyObject *value, PyObject *tra
         type = (PyObject *)Py_TYPE(value);
     }
 
-    // Set the value and traceback of an error. Because calling
-    // PyErr_Restore takes away a reference to each object passed in
-    // as an argument, we manually increase the reference count of
-    // each argument before calling it.
+    // PyErr_Restore steals a reference to each argument.
     Py_INCREF(type);
     Py_INCREF(value);
-    Py_INCREF(traceback);
+    // An omitted traceback (or None) preserves an exception instance's
+    // existing traceback, as generator.throw() and coroutine.throw() do.
+    if (Py_IsNone(traceback) && PyExceptionInstance_Check(value)) {
+        traceback = PyException_GetTraceback(value);
+    } else {
+        Py_INCREF(traceback);
+    }
     PyErr_Restore(type, value, traceback);
 }
 

mypyc/test-data/run-async.test

@@ -2733,6 +2733,87 @@ def test_close_suspended_native_child() -> None:
     coro.send(None)
     coro.close()
 
+async def throw_child(explicit: bool) -> None:
+    try:
+        await Suspend(1)
+    except ValueError as exc:
+        if explicit:
+            raise exc
+        raise
+
+async def throw_parent(explicit: bool) -> None:
+    await throw_child(explicit)
+
+def traceback_frames(tb: Any) -> list[str]:
+    result = []
+    while tb is not None:
+        result.append(tb.tb_frame.f_code.co_name)
+        tb = tb.tb_next
+    return result
+
+def throw_origin() -> None:
+    raise ValueError('injected')
+
+def check_throw_traceback(nested: bool, explicit: bool, kind: str) -> None:
+    error: Any
+    try:
+        throw_origin()
+    except ValueError as exc:
+        error = exc
+    if kind == 'fresh':
+        error = ValueError('fresh')
+    original = traceback_frames(error.__traceback__)
+    coro = (throw_parent if nested else throw_child)(explicit)
+    assert coro.send(None) == 0
+    with assertRaises(ValueError):
+        if kind == 'triple':
+            coro.throw(type(error), error, error.__traceback__)
+        else:
+            coro.throw(error)
+    expected = ['throw_parent'] if nested else []
+    if explicit:
+        expected += ['throw_child']
+    expected += ['throw_child', '__await__'] + original
+    # Skip assertRaises and this helper's throw() call.
+    actual = traceback_frames(error.__traceback__)[2:]
+    assert actual == expected, (nested, explicit, kind, actual, expected)
+
+def test_throw_fresh_reraise() -> None:
+    check_throw_traceback(False, False, 'fresh')
+
+def test_throw_existing_reraise() -> None:
+    check_throw_traceback(False, False, 'existing')
+
+def test_throw_triple_reraise() -> None:
+    check_throw_traceback(False, False, 'triple')
+
+def test_throw_fresh_raise_exc() -> None:
+    check_throw_traceback(False, True, 'fresh')
+
+def test_throw_existing_raise_exc() -> None:
+    check_throw_traceback(False, True, 'existing')
+
+def test_throw_triple_raise_exc() -> None:
+    check_throw_traceback(False, True, 'triple')
+
+def test_throw_nested_fresh_reraise() -> None:
+    check_throw_traceback(True, False, 'fresh')
+
+def test_throw_nested_existing_reraise() -> None:
+    check_throw_traceback(True, False, 'existing')
+
+def test_throw_nested_triple_reraise() -> None:
+    check_throw_traceback(True, False, 'triple')
+
+def test_throw_nested_fresh_raise_exc() -> None:
+    check_throw_traceback(True, True, 'fresh')
+
+def test_throw_nested_existing_raise_exc() -> None:
+    check_throw_traceback(True, True, 'existing')
+
+def test_throw_nested_triple_raise_exc() -> None:
+    check_throw_traceback(True, True, 'triple')
+
 [typing fixtures/typing-full.pyi]
 
 [case testRunAsyncNativeCoroutineDrivenByInterpretedCode]

mypyc/test-data/run-generators.test

@@ -410,12 +410,8 @@ with ctx_manager() as c:
     raise Exception
   File "native.py", line 10, in generator
     yield 3
-  File "native.py", line 31, in wrapper
-    return (yield from x)
   File "native.py", line 9, in generator
     yield 2
-  File "native.py", line 31, in wrapper
-    return (yield from x)
 caught exception without value
 caught exception with value some string
 caught exception with value some other string