Pickle Exception.__cause__ and Exception.__traceback__

AUTHORS.rst

@@ -9,3 +9,4 @@ Authors
 * Jon Dufresne - https://github.com/jdufresne
 * Elliott Sales de Andrade - https://github.com/QuLogic
 * Victor Stinner - https://github.com/vstinner
+* Guido Imperiale - https://github.com/crusaderky

CHANGELOG.rst

@@ -2,6 +2,12 @@
 Changelog
 =========
 
+1.6.0 (unreleased)
+~~~~~~~~~~~~~~~~~~
+* When pickling an Exception, also pickle its traceback and the Exception chain
+  (``raise ... from ...``). Contributed by Guido Imperiale in
+  `#53 <https://github.com/ionelmc/python-tblib/issues/53>`_.
+
 1.5.0 (2019-10-23)
 ~~~~~~~~~~~~~~~~~~
 

README.rst

@@ -57,7 +57,7 @@ Overview
 
 .. end-badges
 
-Traceback serialization library.
+Serialization library for Exceptions and Tracebacks.
 
 * Free software: BSD license
 
@@ -69,6 +69,8 @@ It allows you to:
 * Create traceback objects from strings (the ``from_string`` method). *No pickling is used*.
 * Serialize tracebacks to/from plain dicts (the ``from_dict`` and ``to_dict`` methods). *No pickling is used*.
 * Raise the tracebacks created from the aforementioned sources.
+* Pickle an Exception together with its traceback and exception chain
+  (``raise ... from ...``) *(Python 3 only)*
 
 **Again, note that using the pickle support is completely optional. You are solely responsible for
 security problems should you decide to use the pickle support.**
@@ -133,8 +135,8 @@ those tracebacks or print them - that should cover 99% of the usecases.
     >>> len(s3) > 1
     True
 
-Unpickling
-~~~~~~~~~~
+Unpickling tracebacks
+~~~~~~~~~~~~~~~~~~~~~
 
 ::
 
@@ -196,6 +198,75 @@ Raising
         raise Exception('fail')
     Exception: fail
 
+Pickling Exceptions together with their traceback and chain (Python 3 only)
+~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
+
+::
+
+    >>> try:  # doctest: +SKIP
+    ...     try:
+    ...         1 / 0
+    ...     except Exception as e:
+    ...         raise Exception("foo") from e
+    ... except Exception as e:
+    ...     s = pickle.dumps(e)
+    >>> raise pickle.loads(s)  # doctest: +SKIP
+    Traceback (most recent call last):
+      File "<doctest README.rst[16]>", line 3, in <module>
+        1 / 0
+    ZeroDivisionError: division by zero
+
+    The above exception was the direct cause of the following exception:
+
+    Traceback (most recent call last):
+      File "<doctest README.rst[17]>", line 1, in <module>
+        raise pickle.loads(s)
+      File "<doctest README.rst[16]>", line 5, in <module>
+        raise Exception("foo") from e
+    Exception: foo
+
+BaseException subclasses defined after calling ``pickling_support.install()`` will
+**not** retain their traceback and exception chain pickling.
+To cover custom Exceptions, there are two options:
+
+1. Invoke ``pickling_support.install()`` after all modules have been imported
+
+    .. code-block:: python
+
+        >>> from tblib import pickling_support
+        >>> # Declare all imports of your package's dependencies
+        >>> import numpy  # doctest: +SKIP
+
+        >>> # Declare your own custom Exceptions
+        >>> class CustomError(Exception):
+        ...     pass
+
+        >>> # Finally, install tblib
+        >>> pickling_support.install()
+
+2. Selectively install tblib for Exception instances just before they are pickled
+
+    .. code-block:: python
+
+       pickling_support.install(<Exception instance>, [Exception instance], ...)
+
+   The above will install tblib pickling for all listed exceptions as well as any other
+   exceptions in their exception chains.
+
+   For example, one could write a wrapper to be used with
+   `ProcessPoolExecutor <https://docs.python.org/3/library/concurrent.futures.html>`_,
+   `Dask.distributed <https://distributed.dask.org/>`_, or similar libraries:
+
+::
+
+    >>> from tblib import pickling_support
+    >>> def wrapper(func, *args, **kwargs):
+    ...     try:
+    ...         return func(*args, **kwargs)
+    ...     except Exception as e:
+    ...         pickling_support.install(e)
+    ...         raise
+
 What if we have a local stack, does it show correctly ?
 -------------------------------------------------------
 

src/tblib/pickling_support.py

@@ -1,12 +1,14 @@
-try:
-    import copy_reg
-except ImportError:
-    import copyreg as copy_reg
+import sys
 from types import TracebackType
 
 from . import Frame
 from . import Traceback
 
+if sys.version_info.major >= 3:
+    import copyreg
+else:
+    import copy_reg as copyreg
+
 
 def unpickle_traceback(tb_frame, tb_lineno, tb_next):
     ret = object.__new__(Traceback)
@@ -20,5 +22,50 @@ def pickle_traceback(tb):
     return unpickle_traceback, (Frame(tb.tb_frame), tb.tb_lineno, tb.tb_next and Traceback(tb.tb_next))
 
 
-def install():
-    copy_reg.pickle(TracebackType, pickle_traceback)
+def unpickle_exception(func, args, cause, tb):
+    inst = func(*args)
+    inst.__cause__ = cause
+    inst.__traceback__ = tb
+    return inst
+
+
+def pickle_exception(obj):
+    # All exceptions, unlike generic Python objects, define __reduce_ex__
+    # __reduce_ex__(4) should be no different from __reduce_ex__(3).
+    # __reduce_ex__(5) could bring benefits in the unlikely case the exception
+    # directly contains buffers, but PickleBuffer objects will cause a crash when
+    # running on protocol=4, and there's no clean way to figure out the current
+    # protocol from here. Note that any object returned by __reduce_ex__(3) will
+    # still be pickled with protocol 5 if pickle.dump() is running with it.
+    rv = obj.__reduce_ex__(3)
+    if isinstance(rv, str):
+        raise TypeError("str __reduce__ output is not supported")
+    assert isinstance(rv, tuple) and len(rv) >= 2
+
+    return (unpickle_exception, rv[:2] + (obj.__cause__, obj.__traceback__)) + rv[2:]
+
+
+def _get_subclasses(cls):
+    # Depth-first traversal of all direct and indirect subclasses of cls
+    to_visit = [cls]
+    while to_visit:
+        this = to_visit.pop()
+        yield this
+        to_visit += list(this.__subclasses__())
+
+
+def install(*exception_instances):
+    copyreg.pickle(TracebackType, pickle_traceback)
+
+    if sys.version_info.major < 3:
+        return
+
+    if exception_instances:
+        for exc in exception_instances:
+            while exc is not None:
+                copyreg.pickle(type(exc), pickle_exception)
+                exc = exc.__cause__
+
+    else:
+        for exception_cls in _get_subclasses(BaseException):
+            copyreg.pickle(exception_cls, pickle_exception)

tests/test_pickle_exception.py

@@ -0,0 +1,70 @@
+try:
+    import copyreg
+except ImportError:
+    # Python 2
+    import copy_reg as copyreg
+
+import pickle
+import sys
+
+import pytest
+
+import tblib.pickling_support
+
+has_python3 = sys.version_info.major >= 3
+
+
+def setup_function():
+    copyreg.dispatch_table.clear()
+
+
+def teardown_function():
+    copyreg.dispatch_table.clear()
+
+
+class CustomError(Exception):
+    pass
+
+
+@pytest.mark.parametrize(
+    "protocol", [None] + list(range(1, pickle.HIGHEST_PROTOCOL + 1))
+)
+@pytest.mark.parametrize("global_install", [False, True])
+def test_pickle_exceptions(global_install, protocol):
+    if global_install:
+        tblib.pickling_support.install()
+
+    try:
+        try:
+            1 / 0
+        except Exception as e:
+            # Python 3 only syntax
+            # raise CustomError("foo") from e
+            new_e = CustomError("foo")
+            if has_python3:
+                new_e.__cause__ = e
+            raise new_e
+    except Exception as e:
+        exc = e
+    else:
+        assert False
+
+    # Populate Exception.__dict__, which is used in some cases
+    exc.x = 1
+    if has_python3:
+        exc.__cause__.x = 2
+
+    if not global_install:
+        tblib.pickling_support.install(exc)
+    if protocol:
+        exc = pickle.loads(pickle.dumps(exc, protocol=protocol))
+
+    assert isinstance(exc, CustomError)
+    assert exc.args == ("foo",)
+    assert exc.x == 1
+    if has_python3:
+        assert exc.__traceback__ is not None
+        assert isinstance(exc.__cause__, ZeroDivisionError)
+        assert exc.__cause__.__traceback__ is not None
+        assert exc.__cause__.x == 2
+        assert exc.__cause__.__cause__ is None