Repository: giampaolo/psutil
Pull request #2380 (2024-03-14 20:38:13)

Add pickle support to psutil Exceptions
## Summary

* OS: N/A
* Bug fix: yes
* Type: core
* Fixes: 2272

## Description

Add `__reduce__` method to exceptions commonly thrown, and start testing that exceptions can be pickled in test_serialization.

Prior to this, it was not possible to pickle a psutil exception (for error reporting), and this also made it impossible to send them through multiprocessing pipes. Code demonstrating the problem is available in #2272 

giampaolo: Nice change. I like it. 

Patch:
diff --git a/CREDITS b/CREDITS
index a3206bb9..3520723d 100644
--- a/CREDITS
+++ b/CREDITS
@@ -824,4 +824,8 @@ I: 2361
 
 N: Shade Gladden
 W: https://github.com/shadeyg56
-I: 2376
\ No newline at end of file
+I: 2376
+
+N: Anthony Ryan
+W: https://github.com/anthonyryan1
+I: 2272
diff --git a/HISTORY.rst b/HISTORY.rst
index ed422337..1d82aefe 100644
--- a/HISTORY.rst
+++ b/HISTORY.rst
@@ -11,6 +11,7 @@
 
 - 2360_, [macOS]: can't compile on macOS < 10.13.  (patch by Ryan Schmidt)
 - 2254_, [Linux]: offline cpus raise NotImplementedError in cpu_freq() (patch by Shade Gladden)
+- 2272_, Add pickle support to psutil Exceptions
 
 5.9.8
 =====
diff --git a/psutil/_common.py b/psutil/_common.py
index 6989feaf..f33f3e03 100644
--- a/psutil/_common.py
+++ b/psutil/_common.py
@@ -331,6 +331,9 @@ class NoSuchProcess(Error):
         self.name = name
         self.msg = msg or "process no longer exists"
 
+    def __reduce__(self):
+        return (self.__class__, (self.pid, self.name, self.msg))
+
 
 class ZombieProcess(NoSuchProcess):
     """Exception raised when querying a zombie process. This is
@@ -347,6 +350,9 @@ class ZombieProcess(NoSuchProcess):
         self.ppid = ppid
         self.msg = msg or "PID still exists but it's a zombie"
 
+    def __reduce__(self):
+        return (self.__class__, (self.pid, self.name, self.ppid, self.msg))
+
 
 class AccessDenied(Error):
     """Exception raised when permission to perform an action is denied."""
@@ -359,6 +365,9 @@ class AccessDenied(Error):
         self.name = name
         self.msg = msg or ""
 
+    def __reduce__(self):
+        return (self.__class__, (self.pid, self.name, self.msg))
+
 
 class TimeoutExpired(Error):
     """Raised on Process.wait(timeout) if timeout expires and process
@@ -374,6 +383,9 @@ class TimeoutExpired(Error):
         self.name = name
         self.msg = "timeout after %s seconds" % seconds
 
+    def __reduce__(self):
+        return (self.__class__, (self.seconds, self.pid, self.name))
+
 
 # ===================================================================
 # --- utils


Test patch:
diff --git a/psutil/tests/test_misc.py b/psutil/tests/test_misc.py
index 700c054f..bfd1f353 100755
--- a/psutil/tests/test_misc.py
+++ b/psutil/tests/test_misc.py
@@ -280,6 +280,45 @@ class TestMisc(PsutilTestCase):
         check(psutil.disk_usage(os.getcwd()))
         check(psutil.users())
 
+        b = pickle.loads(
+            pickle.dumps(
+                psutil.NoSuchProcess(
+                    pid=4567, name='name test', msg='msg test'
+                )
+            )
+        )
+        self.assertEqual(b.pid, 4567)
+        self.assertEqual(b.name, 'name test')
+        self.assertEqual(b.msg, 'msg test')
+
+        b = pickle.loads(
+            pickle.dumps(
+                psutil.ZombieProcess(
+                    pid=4567, name='name test', ppid=42, msg='msg test'
+                )
+            )
+        )
+        self.assertEqual(b.pid, 4567)
+        self.assertEqual(b.ppid, 42)
+        self.assertEqual(b.name, 'name test')
+        self.assertEqual(b.msg, 'msg test')
+
+        b = pickle.loads(
+            pickle.dumps(psutil.AccessDenied(pid=123, name='name', msg='msg'))
+        )
+        self.assertEqual(b.pid, 123)
+        self.assertEqual(b.name, 'name')
+        self.assertEqual(b.msg, 'msg')
+
+        b = pickle.loads(
+            pickle.dumps(
+                psutil.TimeoutExpired(seconds=33, pid=4567, name='name')
+            )
+        )
+        self.assertEqual(b.seconds, 33)
+        self.assertEqual(b.pid, 4567)
+        self.assertEqual(b.name, 'name')
+
     # # XXX: https://github.com/pypa/setuptools/pull/2896
     # @unittest.skipIf(APPVEYOR, "temporarily disabled due to setuptools bug")
     # def test_setup_script(self):