fix `print_exc`

docs/modules/traceback.md

@@ -5,8 +5,12 @@ label: traceback
 
 ### `traceback.print_exc() -> None`
 
-Print the last exception and its traceback.
+Print the exception currently being handled and its traceback. This also works
+from helper functions called by an exception handler. If no exception is being
+handled, print nothing.
 
-### `traceback.format_exc() -> str`
+### `traceback.format_exc() -> str | None`
 
-Return the last exception and its traceback as a string.
+Return the exception currently being handled and its traceback as a string.
+This uses the same exception as `print_exc()`, including inside helper functions
+and nested `try` blocks. Return `None` if no exception is being handled.

src/modules/traceback.c

@@ -2,25 +2,37 @@
 #include "pocketpy/objects/exception.h"
 #include "pocketpy/interpreter/vm.h"
 
+static char* traceback_formatexc() {
+    // A nested try body or a helper call must not hide the active handler.
+    for(py_Frame* frame = pk_current_vm->top_frame; frame; frame = frame->f_back) {
+        for(int i = frame->exc_stack.length - 1; i >= 0; i--) {
+            FrameExcInfo* info = c11__at(FrameExcInfo, &frame->exc_stack, i);
+            if(!py_isnil(&info->exc)) return formatexc_internal(&info->exc);
+        }
+    }
+    return NULL;
+}
+
 static bool traceback_format_exc(int argc, py_Ref argv) {
     PY_CHECK_ARGC(0);
-    VM* vm = pk_current_vm;
-    if(vm->top_frame) {
-        FrameExcInfo* info = Frame__top_exc_info(vm->top_frame);
-        if(info && !py_isnil(&info->exc)) {
-            char* res = formatexc_internal(&info->exc);
-            py_newstr(py_retval(), res);
-            PK_FREE(res);
-            return true;
-        }
+    char* res = traceback_formatexc();
+    if(res) {
+        py_newstr(py_retval(), res);
+        PK_FREE(res);
+    } else {
+        py_newnone(py_retval());
     }
-    py_newnone(py_retval());
     return true;
 }
 
 static bool traceback_print_exc(int argc, py_Ref argv) {
     PY_CHECK_ARGC(0);
-    py_printexc();
+    char* res = traceback_formatexc();
+    if(res) {
+        pk_current_vm->callbacks.print(res);
+        pk_current_vm->callbacks.print("\n");
+        PK_FREE(res);
+    }
     py_newnone(py_retval());
     return true;
 }
@@ -30,4 +42,4 @@ void pk__add_module_traceback() {
 
     py_bindfunc(mod, "format_exc", traceback_format_exc);
     py_bindfunc(mod, "print_exc", traceback_print_exc);
-}
+}

tests/802_traceback.py

@@ -9,6 +9,8 @@ try:
     b = a[6]
 except KeyError:
     actual = traceback.format_exc()
+    assert traceback.print_exc() is None
+    assert traceback.format_exc() == actual
 
 expected = '''Traceback (most recent call last):
   File "tests/802_traceback.py", line 9
@@ -20,4 +22,80 @@ if actual != expected:
     print(actual)
     print('--- EXPECTED RESULT ---')
     print(expected)
-    exit(1)
+    exit(1)
+
+
+def format_from_helper():
+    return traceback.format_exc()
+
+
+def format_from_try():
+    try:
+        return format_from_helper()
+    except Exception:
+        assert False
+
+
+assert traceback.format_exc() is None
+assert format_from_try() is None
+assert traceback.print_exc() is None
+
+try:
+    raise ValueError('outer')
+except ValueError:
+    outer = traceback.format_exc()
+    assert outer.endswith('ValueError: outer')
+    assert format_from_helper() == outer
+    assert format_from_try() == outer
+
+    # A try body has no exception of its own and must not hide its handler.
+    try:
+        assert traceback.format_exc() == outer
+        assert format_from_try() == outer
+        raise KeyError('inner')
+    except KeyError:
+        inner = traceback.format_exc()
+        assert inner.endswith("KeyError: 'inner'")
+        assert format_from_try() == inner
+
+    # Leaving the inner handler restores the outer exception.
+    assert traceback.format_exc() == outer
+    assert format_from_try() == outer
+
+assert traceback.format_exc() is None
+assert format_from_try() is None
+
+
+def handle_in_helper():
+    try:
+        raise RuntimeError('helper')
+    except RuntimeError:
+        actual = traceback.format_exc()
+        assert actual.endswith('RuntimeError: helper')
+        assert format_from_try() == actual
+
+
+try:
+    raise ValueError('caller')
+except ValueError:
+    caller = traceback.format_exc()
+    handle_in_helper()
+    assert traceback.format_exc() == caller
+    assert format_from_try() == caller
+
+
+class BrokenStr(Exception):
+    def __str__(self):
+        raise RuntimeError('str failed')
+
+
+try:
+    raise BrokenStr()
+except BrokenStr:
+    actual = traceback.format_exc()
+    assert actual.endswith('BrokenStr: <exception str() failed>')
+    assert format_from_try() == actual
+    assert traceback.print_exc() is None
+    assert traceback.format_exc() == actual
+
+assert traceback.format_exc() is None