From 462c2b2eb42e72d95d7ca94cda84db47c789e5f1 Mon Sep 17 00:00:00 2001
From: Anthony Ryan <anthonyryan1@gmail.com>
Date: Thu, 14 Mar 2024 16:25:22 -0400
Subject: [PATCH] Add pickle support to psutil Exceptions

Add __reduce__ method to exceptions commonly thrown, and start
testing that exceptions can be pickled in test_serialization.

Fixes #2272

Signed-off-by: Anthony Ryan <anthonyryan1@gmail.com>
---
 CREDITS                   |  6 +++++-
 HISTORY.rst               |  1 +
 psutil/_common.py         | 12 ++++++++++++
 psutil/tests/test_misc.py | 39 +++++++++++++++++++++++++++++++++++++++
 4 files changed, 57 insertions(+), 1 deletion(-)

diff --git a/CREDITS b/CREDITS
index a3206bb9a..3520723d4 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 ed4223374..1d82aefe7 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 6989feafd..f33f3e035 100644
--- a/psutil/_common.py
+++ b/psutil/_common.py
@@ -331,6 +331,9 @@ def __init__(self, pid, name=None, msg=None):
         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 @@ def __init__(self, pid, name=None, ppid=None, msg=None):
         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 @@ def __init__(self, pid=None, name=None, msg=None):
         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 @@ def __init__(self, seconds, pid=None, name=None):
         self.name = name
         self.msg = "timeout after %s seconds" % seconds
 
+    def __reduce__(self):
+        return (self.__class__, (self.seconds, self.pid, self.name))
+
 
 # ===================================================================
 # --- utils
diff --git a/psutil/tests/test_misc.py b/psutil/tests/test_misc.py
index 700c054f8..bfd1f3535 100755
--- a/psutil/tests/test_misc.py
+++ b/psutil/tests/test_misc.py
@@ -280,6 +280,45 @@ def check(ret):
         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):