From 81b6b9af3ce8403c6826f83e641a93979fb74d77 Mon Sep 17 00:00:00 2001
From: Olivier Grisel <olivier.grisel@ensta.org>
Date: Tue, 24 Jan 2023 16:56:13 +0100
Subject: [PATCH] Fix exception __cause__ propagation with
 tblib.pickling_support.install() (#255)
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

Co-authored-by: Olivier Grisel <olivier.grisel@ensta.org>
Co-authored-by: Thomas Moreau <thomas.moreau.2010@gmail.com>
Co-authored-by: Loïc Estève <loic.esteve@ymail.com>
---
 CHANGES.md                      |  3 ++-
 loky/process_executor.py        |  4 +++-
 tests/_test_process_executor.py | 21 +++++++++++++++++++++
 tox.ini                         |  1 +
 4 files changed, 27 insertions(+), 2 deletions(-)

diff --git a/CHANGES.md b/CHANGES.md
index b46879e1..27f1e8e2 100644
--- 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
index 4e3e819c..60a41bc2 100644
--- 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
index f6c00ca9..20c017b5 100644
--- 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
index 297fd294..91d079ca 100644
--- 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 =