Fix exception __cause__ propagation with tblib.pickling_support.insta…

…ll() (#255) Co-authored-by: Olivier Grisel <[email protected]> Co-authored-by: Thomas Moreau <[email protected]> Co-authored-by: Loïc Estève <[email protected]>
joblib · Jan 24, 2023 · 81b6b9a · 81b6b9a
1 parent 1f21215
commit 81b6b9a
Show file tree

Hide file tree

Showing 4 changed files with 27 additions and 2 deletions.
diff --git a/CHANGES.md b/CHANGES.md
@@ -1,6 +1,7 @@
 ### 3.4.0 - in development
 
-- TODO
+- Fix exception `__cause__` not being propagated with
+  `tblib.pickling_support.install()` (#255).
 
 ### 3.3.0 - 2022-09-15
 

diff --git a/loky/process_executor.py b/loky/process_executor.py
@@ -220,7 +220,9 @@ def __str__(self):
         return self.tb
 
 
-class _ExceptionWithTraceback(BaseException):
+# Do not inherit from BaseException to mirror
+# concurrent.futures.process._ExceptionWithTraceback
+class _ExceptionWithTraceback:
 
     def __init__(self, exc):
         tb = getattr(exc, "__traceback__", None)

diff --git a/tests/_test_process_executor.py b/tests/_test_process_executor.py
@@ -1085,6 +1085,26 @@ def check_viztracer_active_and_custom_init():
         finally:
             tracer.stop()
 
+    def test_exception_cause_with_tblib(self):
+        "Ensure tampering with exception pickling do not break __cause__ propagation"
+        tblib_pickling_support = pytest.importorskip('tblib.pickling_support')
+
+        error_message = "This is the error message"
+
+        def raise_value_error():
+            tblib_pickling_support.install()
+            raise ValueError(error_message)
+
+        executor = self.executor_type(max_workers=2)
+        f = executor.submit(raise_value_error)
+        try:
+            f.result()
+        except ValueError as e:
+            assert e.__cause__ is not None
+            assert error_message in str(e.__cause__)
+
+        executor.shutdown(wait=True)
+
 
 def _custom_initializer():
     """_custom_initializer is module function to be picklable
@@ -1093,3 +1113,4 @@ def _custom_initializer():
     use cloudpickle to pickle the initializer.
     """
     loky._custom_global_var = 42
+
diff --git a/tox.ini b/tox.ini
@@ -17,6 +17,7 @@ deps =
      coverage
      viztracer ; python_version >= '3.8' and python_version <= '3.10'
      numpy ; implementation_name == 'cpython'
+     tblib
 whitelist_externals=
      bash
 setenv =