From 26bfc36fa8f9290110cd508a443834a1a6135c32 Mon Sep 17 00:00:00 2001 From: Vladimir Gerts Date: Fri, 14 Nov 2025 16:02:03 +0100 Subject: [PATCH 1/2] Preserve special attributes of built-in exceptions --- src/tblib/pickling_support.py | 99 ++++++++++++++-------------------- tests/test_pickle_exception.py | 53 +++++++++++++++++- 2 files changed, 92 insertions(+), 60 deletions(-) diff --git a/src/tblib/pickling_support.py b/src/tblib/pickling_support.py index 2f57449..5a61140 100644 --- a/src/tblib/pickling_support.py +++ b/src/tblib/pickling_support.py @@ -52,71 +52,52 @@ def unpickle_exception(func, args, cause, tb, context=None, suppress_context=Fal return inst -def pickle_exception( - obj, builtin_reducers=(OSError.__reduce__, BaseException.__reduce__), builtin_inits=(OSError.__init__, BaseException.__init__) -): +def _get_public_class_attributes(cls: type) -> set[str]: + return { + attr + for mro_cls in cls.mro() + for attr in mro_cls.__dict__.keys() + if not attr.startswith('_') and not callable(getattr(mro_cls, attr)) + } + + +def pickle_exception(obj): reduced_value = obj.__reduce__() if isinstance(reduced_value, str): raise TypeError('Did not expect {repr(obj)}.__reduce__() to return a string!') func = type(obj) - # Detect busted objects: they have a custom __init__ but no __reduce__. - # This also means the resulting exceptions may be a bit "dulled" down - the args from __reduce__ are discarded. - if func.__reduce__ in builtin_reducers and func.__init__ not in builtin_inits: - _, args, *optionals = reduced_value - attrs = { - '__dict__': obj.__dict__, - 'args': obj.args, - } - args = () - if isinstance(obj, OSError): - # Only set OSError-specific attributes if they are not None - # Setting them to None explicitly breaks the string representation - if obj.errno is not None: - attrs['errno'] = obj.errno - if obj.strerror is not None: - attrs['strerror'] = obj.strerror - if (winerror := getattr(obj, 'winerror', None)) is not None: - attrs['winerror'] = winerror - if obj.filename is not None: - attrs['filename'] = obj.filename - if obj.filename2 is not None: - attrs['filename2'] = obj.filename2 - if ExceptionGroup is not None and isinstance(obj, ExceptionGroup): - args = (obj.message, obj.exceptions) - - return ( - unpickle_exception_with_attrs, - ( - func, - attrs, - obj.__cause__, - obj.__traceback__, - obj.__context__, - obj.__suppress_context__, - # __notes__ doesn't exist prior to Python 3.11; and even on Python 3.11 it may be absent - getattr(obj, '__notes__', None), - args, - ), - *optionals, - ) + + _, args, *optionals = reduced_value + attrs = { + '__dict__': obj.__dict__, + 'args': obj.args, + } + + if ExceptionGroup is not None and isinstance(obj, ExceptionGroup): + args = (obj.message, obj.exceptions) + else: - func, args, *optionals = reduced_value - - return ( - unpickle_exception, - ( - func, - args, - obj.__cause__, - obj.__traceback__, - obj.__context__, - obj.__suppress_context__, - # __notes__ doesn't exist prior to Python 3.11; and even on Python 3.11 it may be absent - getattr(obj, '__notes__', None), - ), - *optionals, - ) + public_class_attributes = _get_public_class_attributes(type(obj)) + additional_obj_attributes = {attr: value for attr in public_class_attributes if (value := getattr(obj, attr, None)) is not None} + attrs.update(additional_obj_attributes) + args = () + + return ( + unpickle_exception_with_attrs, + ( + func, + attrs, + obj.__cause__, + obj.__traceback__, + obj.__context__, + obj.__suppress_context__, + # __notes__ doesn't exist prior to Python 3.11; and even on Python 3.11 it may be absent + getattr(obj, '__notes__', None), + args, + ), + *optionals, + ) def _get_subclasses(cls): diff --git a/tests/test_pickle_exception.py b/tests/test_pickle_exception.py index 0990d9d..a7232c5 100644 --- a/tests/test_pickle_exception.py +++ b/tests/test_pickle_exception.py @@ -15,7 +15,7 @@ import tblib.pickling_support has_python311 = sys.version_info >= (3, 11) - +has_python310 = sys.version_info >= (3, 10) @pytest.fixture def clear_dispatch_table(): @@ -543,3 +543,54 @@ def test_exception_group(): assert isinstance(exc.exceptions[8], OSError) assert exc.exceptions[8].errno == 2 assert str(exc.exceptions[8]) == real_oserror_str + + +def test_systemexit_error(clear_dispatch_table): + try: + raise SystemExit(42) + except SystemExit as e: + exc = e + + tblib.pickling_support.install(exc) + exc = pickle.loads(pickle.dumps(exc)) + + assert isinstance(exc, SystemExit) + assert exc.code == 42 + assert exc.__traceback__ is not None + + +@pytest.mark.parametrize( + ('builtin_exception', 'args', 'kw_args'), + [ + ( + AttributeError, + ('msg',), + {'name': 'name', 'obj': None} if has_python310 else {}, # takes no keywords before v3.10 + ), + (NameError, ('msg',), {'name': 'name'} if has_python310 else {}), # takes no keywords before v3.10 + (ImportError, ('msg',), {'name': 'name', 'path': 'some/path'}), + (OSError, (2, 'err', 3, None, 5), {}), + (StopIteration, ('value',), {}), + ( + SyntaxError, + ( + 'msg', + ('fname', 42, 21, 'src', 123, 456) if has_python310 else ('fname', 42, 21, 'src'), + ), + {}, + ), # last two attributes added in v3.10 + (SystemExit, (42,), {}), + (UnicodeDecodeError, ('encoding', bytearray(), 1, 2, 'reason'), {}), + (UnicodeEncodeError, ('encoding', 'object', 1, 2, 'reason'), {}), + (UnicodeTranslateError, ('object', 1, 2, 'reason'), {}), + ], +) +def test_builtin_exceptions_roundtrip(builtin_exception, args, kw_args, clear_dispatch_table): + tblib.pickling_support.install() + exc_orig = builtin_exception(*args, **kw_args) + exc_unpickled = pickle.loads(pickle.dumps(exc_orig)) + + public_class_attributes = tblib.pickling_support._get_public_class_attributes(type(exc_orig)) + + for attr in public_class_attributes: + assert getattr(exc_unpickled, attr, None) == getattr(exc_orig, attr, None) From 0d8b4f595c8364e75e3b480389dd5efc724febd7 Mon Sep 17 00:00:00 2001 From: Vladimir Gerts Date: Fri, 14 Nov 2025 16:17:51 +0100 Subject: [PATCH 2/2] fix formatting --- tests/test_pickle_exception.py | 1 + 1 file changed, 1 insertion(+) diff --git a/tests/test_pickle_exception.py b/tests/test_pickle_exception.py index a7232c5..3e8dc96 100644 --- a/tests/test_pickle_exception.py +++ b/tests/test_pickle_exception.py @@ -17,6 +17,7 @@ has_python311 = sys.version_info >= (3, 11) has_python310 = sys.version_info >= (3, 10) + @pytest.fixture def clear_dispatch_table(): bak = copyreg.dispatch_table.copy()