Add support for various Python 3.11 things

setup.py

@@ -13,7 +13,7 @@ def read(*names, **kwargs):
 
 setup(
     name='tblib',
-    version='2.0.0',
+    version='2.1.0',
     license='BSD-2-Clause',
     description='Traceback serialization library.',
     long_description='{}\n{}'.format(

src/tblib/pickling_support.py

@@ -22,10 +22,16 @@ def pickle_traceback(tb):
     return unpickle_traceback, (Frame(tb.tb_frame), tb.tb_lineno, tb.tb_next and Traceback(tb.tb_next))
 
 
-def unpickle_exception(func, args, cause, tb):
+# Note: Older versions of tblib will generate pickle archives that call unpickle_exception() with
+# fewer arguments. We assign default values to some of the arguments to support this.
+def unpickle_exception(func, args, cause, tb, context=None, suppress_context=False, notes=None):
     inst = func(*args)
     inst.__cause__ = cause
     inst.__traceback__ = tb
+    inst.__context__ = context
+    inst.__suppress_context__ = suppress_context
+    if notes is not None:
+        inst.__notes__ = notes
     return inst
 
 
@@ -43,7 +49,18 @@ def pickle_exception(obj):
     assert isinstance(rv, tuple)
     assert len(rv) >= 2
 
-    return (unpickle_exception, rv[:2] + (obj.__cause__, obj.__traceback__)) + rv[2:]
+    return (
+        unpickle_exception,
+        rv[:2]
+        + (
+            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),
+        ),
+    ) + rv[2:]
 
 
 def _get_subclasses(cls):
@@ -73,13 +90,36 @@ def install(*exc_classes_or_instances):
 
     for exc in exc_classes_or_instances:
         if isinstance(exc, BaseException):
-            while exc is not None:
-                copyreg.pickle(type(exc), pickle_exception)
-                exc = exc.__cause__
+            _install_for_instance(exc, set())
         elif isinstance(exc, type) and issubclass(exc, BaseException):
             copyreg.pickle(exc, pickle_exception)
             # Allow using @install as a decorator for Exception classes
             if len(exc_classes_or_instances) == 1:
                 return exc
         else:
             raise TypeError('Expected subclasses or instances of BaseException, got %s' % (type(exc)))
+
+
+def _install_for_instance(exc, seen):
+    assert isinstance(exc, BaseException)
+
+    # Prevent infinite recursion if we somehow get a self-referential exception. (Self-referential
+    # exceptions should never normally happen, but if it did somehow happen, we want to pickle the
+    # exception faithfully so the developer can troubleshoot why it happened.)
+    if id(exc) in seen:
+        return
+    seen.add(id(exc))
+
+    copyreg.pickle(type(exc), pickle_exception)
+
+    if exc.__cause__ is not None:
+        _install_for_instance(exc.__cause__, seen)
+    if exc.__context__ is not None:
+        _install_for_instance(exc.__context__, seen)
+
+    # This case is meant to cover BaseExceptionGroup on Python 3.11 as well as backports like the
+    # exceptiongroup module
+    if hasattr(exc, 'exceptions') and isinstance(exc.exceptions, (tuple, list)):
+        for subexc in exc.exceptions:
+            if isinstance(subexc, BaseException):
+                _install_for_instance(subexc, seen)

tests/test_pickle_exception.py

@@ -12,6 +12,7 @@
 import tblib.pickling_support
 
 has_python3 = sys.version_info.major >= 3
+has_python311 = sys.version_info >= (3, 11)
 
 
 @pytest.fixture
@@ -33,17 +34,24 @@ def test_install(clear_dispatch_table, how, protocol):
     if how == 'global':
         tblib.pickling_support.install()
     elif how == 'class':
-        tblib.pickling_support.install(CustomError, ZeroDivisionError)
+        tblib.pickling_support.install(CustomError, ValueError, ZeroDivisionError)
 
     try:
         try:
-            1 / 0  # noqa: B018
+            try:
+                1 / 0  # noqa: B018
+            finally:
+                # The ValueError's __context__ will be the ZeroDivisionError
+                raise ValueError('blah')
         except Exception as e:
             # Python 3 only syntax
             # raise CustomError("foo") from e
             new_e = CustomError('foo')
             if has_python3:
                 new_e.__cause__ = e
+                if has_python311:
+                    new_e.add_note('note 1')
+                    new_e.add_note('note 2')
             raise new_e
     except Exception as e:
         exc = e
@@ -54,6 +62,7 @@ def test_install(clear_dispatch_table, how, protocol):
     exc.x = 1
     if has_python3:
         exc.__cause__.x = 2
+        exc.__cause__.__context__.x = 3
 
     if how == 'instance':
         tblib.pickling_support.install(exc)
@@ -65,11 +74,20 @@ def test_install(clear_dispatch_table, how, protocol):
     assert exc.x == 1
     if has_python3:
         assert exc.__traceback__ is not None
-        assert isinstance(exc.__cause__, ZeroDivisionError)
+
+        assert isinstance(exc.__cause__, ValueError)
         assert exc.__cause__.__traceback__ is not None
         assert exc.__cause__.x == 2
         assert exc.__cause__.__cause__ is None
 
+        assert isinstance(exc.__cause__.__context__, ZeroDivisionError)
+        assert exc.__cause__.__context__.x == 3
+        assert exc.__cause__.__context__.__cause__ is None
+        assert exc.__cause__.__context__.__context__ is None
+
+        if has_python311:
+            assert exc.__notes__ == ['note 1', 'note 2']
+
 
 @tblib.pickling_support.install
 class RegisteredError(Exception):
@@ -90,6 +108,18 @@ def test_install_decorator():
         assert exc.__traceback__ is not None
 
 
+@pytest.mark.skipif(not has_python311, reason='no BaseExceptionGroup before Python 3.11')
+def test_install_instance_recursively(clear_dispatch_table):
+    exc = BaseExceptionGroup('test', [ValueError('foo'), CustomError('bar')])
+    exc.exceptions[0].__cause__ = ZeroDivisionError('baz')
+    exc.exceptions[0].__cause__.__context__ = AttributeError('quux')
+
+    tblib.pickling_support.install(exc)
+
+    installed = {c for c in copyreg.dispatch_table if issubclass(c, BaseException)}
+    assert installed == {ExceptionGroup, ValueError, CustomError, ZeroDivisionError, AttributeError}
+
+
 @pytest.mark.skipif(sys.version_info[0] < 3, reason='No checks done in Python 2')
 def test_install_typeerror():
     with pytest.raises(TypeError):