Bugfix/exception kwargs

.pre-commit-config.yaml

@@ -23,13 +23,10 @@ repos:
       - id: pytest-coverage
         name: pytest-coverage
         entry: >
-          powershell -Command "
-          $env:UV_PROJECT_ENVIRONMENT = '.venv-3.12';
-          uv run --python 3.12 pytest --cov=offline_debug;
-          if ($LASTEXITCODE -ne 0) { exit $LASTEXITCODE };
-          $env:UV_PROJECT_ENVIRONMENT = '.venv-3.13';
-          uv run --python 3.13 pytest --cov=offline_debug --cov-append --cov-report=term --cov-fail-under=100
-          "
+          bash -c '
+          UV_PROJECT_ENVIRONMENT=.venv-3.12 uv run --python 3.12 pytest --cov=offline_debug &&
+          UV_PROJECT_ENVIRONMENT=.venv-3.13 uv run --python 3.13 pytest --cov=offline_debug --cov-append --cov-report=term --cov-fail-under=100
+          '
         language: system
         types: [python]
         pass_filenames: false

README.md

@@ -71,7 +71,7 @@ if isinstance(data, ExceptionGroupData):
 - **Python 3.13 Compatibility**: Leverages PEP 667 features where `f_locals` is a write-through proxy, allowing for accurate local
   variable restoration.
 - **Support python 3.12 as well**
-- **Robust Serialization**:
+- **Resilient Serialization**:
     - `pickle` is used for exceptions and variables.
     - `marshal` is used for code objects.
     - Non-picklable items are gracefully handled by storing their `repr`.

offline_debug/_inner/_pickle_helpers.py

@@ -0,0 +1,144 @@
+"""
+Pickling for exceptions that vanilla pickle cannot reconstruct.
+
+``BaseException`` reduces to ``(type(self), self.args, state)``, so pickle
+rebuilds an exception by calling ``Cls(*self.args)`` and then applying state.
+Exceptions whose ``__init__``/``__new__`` takes required keyword-only arguments
+(and which do not populate ``self.args`` via ``super().__init__``) pickle fine
+but fail on load with a ``TypeError``. Since a traceback-capture library must
+serialize whatever the user raised, we take over reconstruction for such
+exceptions.
+
+We keep the takeover as narrow as possible so well-behaved exceptions keep
+their vanilla pickle behavior:
+
+- Classes that customize the pickle protocol (``__reduce__``, ``__setstate__``,
+  ...) are left alone; their hooks know how to rebuild them (e.g. ``OSError``).
+- Classes whose constructors are all C-level (every builtin exception) are left
+  alone; ``Cls(*args)`` is guaranteed to work and restores C-level state such
+  as ``StopIteration.value`` that lives outside ``__dict__``.
+- Only classes with a Python-defined ``__init__``/``__new__`` — the only kind
+  that can demand required keyword-only arguments — are reconstructed by us,
+  and even then reconstruction first retries normal construction and only
+  falls back to ``__new__`` (bypassing the constructor) when it fails.
+
+Exception groups get a dedicated reconstructor because a bare
+``cls.__new__(cls)`` is not enough for them: ``BaseExceptionGroup.__new__``
+requires ``message``/``exceptions`` to build a valid group for any subclass.
+
+.. warning::
+    The path of this module and the names of :func:`reconstruct_exception` and
+    :func:`reconstruct_exception_group` are embedded in every dump that takes
+    the takeover path. Renaming or moving them makes previously saved dumps
+    unloadable, so treat them as part of the on-disk format.
+"""
+
+import contextlib
+import io
+import pickle
+import types
+from typing import IO, Any
+
+# Pickle protocol hooks that let a class control its own serialization.
+# If an exception class customizes any of these, we must not override its
+# reduction: the class (e.g. OSError, ImportError) knows how to rebuild itself.
+_PICKLE_PROTOCOL_HOOKS = (
+    "__reduce_ex__",
+    "__reduce__",
+    "__getstate__",
+    "__setstate__",
+    "__getnewargs__",
+    "__getnewargs_ex__",
+)
+
+
+def _customizes_pickling(cls: type[BaseException]) -> bool:
+    """Whether ``cls`` overrides any pickle protocol hook of ``BaseException``."""
+    return any(
+        getattr(cls, hook, None) is not getattr(BaseException, hook, None)
+        for hook in _PICKLE_PROTOCOL_HOOKS
+    )
+
+
+def _has_python_constructor(cls: type[BaseException]) -> bool:
+    """
+    Whether ``cls`` has a Python-defined ``__init__`` or ``__new__``.
+
+    Builtin exceptions expose C slot wrappers here, and for them default pickle
+    reconstruction always works. Only a Python-level constructor can introduce
+    required keyword-only arguments that break ``Cls(*args)`` at load time.
+    """
+    return isinstance(cls.__init__, types.FunctionType) or isinstance(
+        cls.__new__, types.FunctionType
+    )
+
+
+def reconstruct_exception(
+    cls: type[BaseException], args: tuple[Any, ...], state: dict[str, Any] | None
+) -> BaseException:
+    """Rebuild an exception, bypassing its constructor only if it rejects ``args``."""
+    try:
+        exc = cls(*args)
+    except Exception:  # noqa: BLE001 - a rejecting constructor is exactly the case we handle
+        exc = BaseException.__new__(cls)
+        # Suppressed failures happen when e.g. `args` is shadowed by a read-only property.
+        with contextlib.suppress(Exception):
+            exc.args = args
+    if state:
+        exc.__dict__.update(state)
+    return exc
+
+
+def reconstruct_exception_group(
+    cls: type[BaseExceptionGroup[Any]],
+    message: str,
+    exceptions: tuple[BaseException, ...],
+    state: dict[str, Any] | None,
+) -> BaseExceptionGroup[Any]:
+    """
+    Rebuild an exception group, bypassing its constructor only when necessary.
+
+    ``BaseExceptionGroup.__new__`` establishes ``message``/``exceptions`` for any
+    subclass, so when normal construction fails (e.g. a subclass demanding extra
+    required keyword-only arguments) we call it directly.
+    """
+    try:
+        exc = cls(message, exceptions)
+    except Exception:  # noqa: BLE001 - a rejecting constructor is exactly the case we handle
+        exc = BaseExceptionGroup.__new__(cls, message, exceptions)
+    if state:
+        exc.__dict__.update(state)
+    return exc
+
+
+class CustomExceptionPickler(pickle.Pickler):
+    """Pickler that takes over reconstruction of constructor-rejecting exceptions."""
+
+    # The inline suppression below is needed because this returns NotImplemented to
+    # fall back to default reduction, which the stdlib stub's return type omits.
+    def reducer_override(self, obj: object, /) -> object:  # ty: ignore[invalid-method-override]
+        if not isinstance(obj, BaseException):
+            return NotImplemented
+        cls = type(obj)
+        if _customizes_pickling(cls) or not _has_python_constructor(cls):
+            # Default reduction is guaranteed to round-trip; taking over would
+            # lose state the class restores itself (e.g. OSError.errno).
+            return NotImplemented
+        state = obj.__dict__.copy() or None
+        # Groups need their dedicated reconstructor: message/exceptions live
+        # behind __new__ rather than a settable `args`.
+        if isinstance(obj, BaseExceptionGroup):
+            return reconstruct_exception_group, (cls, obj.message, obj.exceptions, state)
+        return reconstruct_exception, (cls, obj.args, state)
+
+
+def exception_safe_dump(obj: object, file: IO[bytes]) -> None:
+    """Serialize ``obj`` to an open binary ``file`` using :class:`CustomExceptionPickler`."""
+    CustomExceptionPickler(file).dump(obj)
+
+
+def exception_safe_dumps(obj: object) -> bytes:
+    """Serialize ``obj`` to bytes using :class:`CustomExceptionPickler`."""
+    buf = io.BytesIO()
+    exception_safe_dump(obj, buf)
+    return buf.getvalue()

offline_debug/_inner/load_traceback.py

@@ -8,6 +8,7 @@
 from pathlib import Path
 from types import CodeType
 
+from offline_debug._inner._pickle_helpers import reconstruct_exception_group
 from offline_debug._inner.c_api import (
     create_frame,
     link_frame,
@@ -39,9 +40,13 @@ def _reconstruct_exc_data(data: ExceptionData) -> BaseException:
 
     if isinstance(data, ExceptionGroupData) and isinstance(exc, BaseExceptionGroup):
         inner_excs = [_reconstruct_exc_data(e) for e in data.exceptions]
-        # We must use derive to create a new ExceptionGroup with reconstructed inner exceptions.
-        # The exceptions that are inside the unpickled exc object are have incomplete data.
-        exc = exc.derive(inner_excs)
+        # The exceptions inside the unpickled exc object have incomplete data, so
+        # rebuild the group around the fully reconstructed ones. We must not use
+        # derive() for this: its default implementation returns a plain
+        # ExceptionGroup, dropping the subclass type and its custom state.
+        exc = reconstruct_exception_group(
+            type(exc), exc.message, tuple(inner_excs), exc.__dict__.copy() or None
+        )
 
     reconstructed_frames: list[tuple[types.FrameType, FrameData]] = []
     for f_data in data.tb_frames:

offline_debug/_inner/save_traceback.py

@@ -6,6 +6,7 @@
 from io import BytesIO
 from pathlib import Path
 
+from offline_debug._inner._pickle_helpers import exception_safe_dump, exception_safe_dumps
 from offline_debug._inner.models import ExceptionData, ExceptionGroupData, FrameData
 
 # Internal attributes that are either unpicklable or redundant in a new process.
@@ -24,24 +25,42 @@ def _get_stack_depth(frame: types.FrameType) -> int:
     return depth
 
 
-def _filter_dict(d: dict) -> dict:
-    """Filter dictionary to include only picklable items."""
+def _filter_dict(d: dict, roundtrip_cache: dict[int, str | None]) -> dict:
+    """
+    Filter dictionary to include only items that survive a pickle round-trip.
+
+    ``roundtrip_cache`` maps ``id(value)`` to ``None`` (survives) or a placeholder
+    string, so a value shared across frames (e.g. module globals) is only checked
+    once per save. The cached objects stay alive for the whole save because the
+    frames still reference them, so the ids are stable.
+    """
     result = {}
     for k, v in d.items():
         if k in _INTERNAL_ATTRIBUTES_TO_SKIP:
             continue
-        try:
-            # We must verify if the value is picklable because many globals
-            # (like open file handles, database connections, or modules)
-            # cannot be saved to disk.
-            pickle.dumps(v)
-            result[k] = v
-        except Exception:  # noqa: BLE001
-            result[k] = f"<unpicklable {type(v).__name__}: {v!r}>"
+        cache_key = id(v)
+        if cache_key not in roundtrip_cache:
+            try:
+                # We must verify that the value survives a full pickle round-trip
+                # because many globals (like open file handles, database connections,
+                # or modules) cannot be saved to disk, and some values pickle but fail
+                # to unpickle (e.g. a custom __reduce__ whose callable raises on load).
+                # Such values would otherwise break the entire load, so we replace
+                # them with a placeholder. We use the same pickler that serializes
+                # these dicts so the check reflects what will actually be written.
+                pickle.loads(exception_safe_dumps(v))  # noqa: S301
+                roundtrip_cache[cache_key] = None
+            except BaseException:  # noqa: BLE001 - even a KeyboardInterrupt raised by a
+                # value's reconstruction must not abort capturing the traceback.
+                roundtrip_cache[cache_key] = f"<unpicklable {type(v).__name__}: {v!r}>"
+        placeholder = roundtrip_cache[cache_key]
+        result[k] = v if placeholder is None else placeholder
     return result
 
 
-def _serialize_exc_data(exc: BaseException) -> ExceptionData:
+def _serialize_exc_data(
+    exc: BaseException, roundtrip_cache: dict[int, str | None]
+) -> ExceptionData:
     """Recursively serialize exception data into dataclasses."""
     tb_frames: list[FrameData] = []
     curr_tb = exc.__traceback__
@@ -60,8 +79,8 @@ def _serialize_exc_data(exc: BaseException) -> ExceptionData:
         tb_frames.append(
             FrameData(
                 code=marshal.dumps(f.f_code),
-                globals=_filter_dict(f.f_globals),
-                locals=_filter_dict(f.f_locals),
+                globals=_filter_dict(f.f_globals, roundtrip_cache),
+                locals=_filter_dict(f.f_locals, roundtrip_cache),
                 lasti=curr_tb.tb_lasti,
                 lineno=curr_tb.tb_lineno,
                 stack_depth=_get_stack_depth(f),
@@ -71,22 +90,26 @@ def _serialize_exc_data(exc: BaseException) -> ExceptionData:
         curr_tb = curr_tb.tb_next
 
     try:
-        exc_pickle = pickle.dumps(exc)
+        exc_pickle = exception_safe_dumps(exc)
+        # A dump that cannot be loaded later is worse than a placeholder, so also
+        # verify the exception survives loading (e.g. a custom __reduce__ whose
+        # reconstruction fails only at load time).
+        pickle.loads(exc_pickle)  # noqa: S301
     except Exception:  # noqa: BLE001
-        exc_pickle = pickle.dumps(
+        exc_pickle = exception_safe_dumps(
             RuntimeError(f"Unpicklable exception {type(exc).__name__}: {exc!s}")
         )
 
-    cause = _serialize_exc_data(exc.__cause__) if exc.__cause__ else None
-    context = _serialize_exc_data(exc.__context__) if exc.__context__ else None
+    cause = _serialize_exc_data(exc.__cause__, roundtrip_cache) if exc.__cause__ else None
+    context = _serialize_exc_data(exc.__context__, roundtrip_cache) if exc.__context__ else None
 
     if isinstance(exc, BaseExceptionGroup):
         return ExceptionGroupData(
             exc_pickle=exc_pickle,
             tb_frames=tb_frames,
             cause=cause,
             context=context,
-            exceptions=[_serialize_exc_data(e) for e in exc.exceptions],
+            exceptions=[_serialize_exc_data(e, roundtrip_cache) for e in exc.exceptions],
         )
 
     return ExceptionData(
@@ -99,15 +122,15 @@ def _serialize_exc_data(exc: BaseException) -> ExceptionData:
 
 def save_traceback(exc: BaseException, file: Path | BytesIO | None) -> ExceptionData:
     """Serialize an exception and its traceback to a file."""
-    data = _serialize_exc_data(exc)
+    data = _serialize_exc_data(exc, roundtrip_cache={})
     if file is None:
         return data
 
     if isinstance(file, Path):
         with file.open("wb") as f:
-            pickle.dump(data, f)
+            exception_safe_dump(data, f)
     elif isinstance(file, BytesIO):
-        pickle.dump(data, file)
+        exception_safe_dump(data, file)
     else:
         msg = f"Unexpected type for file {type(file).__name__}"
         raise TypeError(msg)

pyproject.toml

@@ -1,6 +1,6 @@
 [project]
 name = "offline-debug"
-version = "0.3.0"
+version = "0.3.1"
 description = "Debug exceptions offline by saving them to a dump and raising them at a later point."
 readme = "README.md"
 authors = [