From 8c41dc876149282183d23bc1adb86b830ab401bc Mon Sep 17 00:00:00 2001
From: Tyler Reddy <tyler.je.reddy@gmail.com>
Date: Wed, 11 Oct 2023 20:45:28 -0600
Subject: [PATCH] MAINT: Python 3.12 support improve (#4300)

* Fixes gh-4299, and allows the testsuite to have 12 failures
instead of 56 with Python `3.12.0rc3` for me locally on x86_64 Linux;
full suite still passing for Python `3.11.x`

* the main idea is to circumvent more aggressive CPython
blocks on serialization by overriding the highest-priority
pickling method available, `__reduce_ex__`, which CPython
had locked down, making it such that our old `__getstate__`
pickling shims were ignored as indicators of pickling safety

* there is probably a slightly simpler diff that allows
this to work, but this took a few hours to draft so feel
free to push in simplifications if you're sure you've tested
on 3.11.x and 3.12 RC

[skip cirrus]
---
 package/MDAnalysis/lib/picklable_file_io.py | 80 +++++++++++++--------
 1 file changed, 52 insertions(+), 28 deletions(-)

diff --git a/package/MDAnalysis/lib/picklable_file_io.py b/package/MDAnalysis/lib/picklable_file_io.py
index dd5bef44c14..dd0dad7ad26 100644
--- a/package/MDAnalysis/lib/picklable_file_io.py
+++ b/package/MDAnalysis/lib/picklable_file_io.py
@@ -109,16 +109,23 @@ def __init__(self, name, mode='r'):
         self._mode = mode
         super().__init__(name, mode)
 
-    def __getstate__(self):
+
+    def __setstate__(self, state):
+        name = state["name_val"]
+        super().__init__(name, mode='r')
+        try:
+            self.seek(state["tell_val"])
+        except KeyError:
+            pass
+
+    def __reduce_ex__(self, prot):
         if self._mode != 'r':
             raise RuntimeError("Can only pickle files that were opened "
                                "in read mode, not {}".format(self._mode))
-        return self.name, self.tell()
-
-    def __setstate__(self, args):
-        name = args[0]
-        super().__init__(name, mode='r')
-        self.seek(args[1])
+        return (self.__class__,
+                (self.name, self._mode),
+                {"name_val": self.name,
+                 "tell_val": self.tell()})
 
 
 class BufferIOPicklable(io.BufferedReader):
@@ -151,16 +158,22 @@ def __init__(self, raw):
         super().__init__(raw)
         self.raw_class = raw.__class__
 
-    def __getstate__(self):
-        return self.raw_class, self.name, self.tell()
 
-    def __setstate__(self, args):
-        raw_class = args[0]
-        name = args[1]
+    def __setstate__(self, state):
+        raw_class = state["raw_class"]
+        name = state["name_val"]
         raw = raw_class(name)
         super().__init__(raw)
-        self.seek(args[2])
+        self.seek(state["tell_val"])
 
+    def __reduce_ex__(self, prot):
+        # don't ask, for Python 3.12+ see:
+        # https://github.com/python/cpython/pull/104370
+        return (self.raw_class,
+                (self.name,),
+                {"raw_class": self.raw_class,
+                 "name_val": self.name,
+                 "tell_val": self.tell()})
 
 class TextIOPicklable(io.TextIOWrapper):
     """Character and line based picklable file-like object.
@@ -197,22 +210,26 @@ def __init__(self, raw):
         super().__init__(raw)
         self.raw_class = raw.__class__
 
-    def __getstate__(self):
-        try:
-            name = self.name
-        except AttributeError:
-            # This is kind of ugly--BZ2File does not save its name.
-            name = self.buffer._fp.name
-        return self.raw_class, name
 
     def __setstate__(self, args):
-        raw_class = args[0]
-        name = args[1]
+        raw_class = args["raw_class"]
+        name = args["name_val"]
         # raw_class is used for further expansion this functionality to
         # Gzip files, which also requires a text wrapper.
         raw = raw_class(name)
         super().__init__(raw)
 
+    def __reduce_ex__(self, prot):
+        try:
+            name = self.name
+        except AttributeError:
+            # This is kind of ugly--BZ2File does not save its name.
+            name = self.buffer._fp.name
+        return (self.__class__.__new__,
+                (self.__class__,),
+                {"raw_class": self.raw_class,
+                 "name_val": name})
+
 
 class BZ2Picklable(bz2.BZ2File):
     """File object (read-only) for bzip2 (de)compression that can be pickled.
@@ -269,11 +286,14 @@ def __getstate__(self):
         if not self._bz_mode.startswith('r'):
             raise RuntimeError("Can only pickle files that were opened "
                                "in read mode, not {}".format(self._bz_mode))
-        return self._fp.name, self.tell()
+        return {"name_val": self._fp.name, "tell_val": self.tell()}
 
     def __setstate__(self, args):
-        super().__init__(args[0])
-        self.seek(args[1])
+        super().__init__(args["name_val"])
+        try:
+            self.seek(args["tell_val"])
+        except KeyError:
+            pass
 
 
 class GzipPicklable(gzip.GzipFile):
@@ -331,11 +351,15 @@ def __getstate__(self):
         if not self._gz_mode.startswith('r'):
             raise RuntimeError("Can only pickle files that were opened "
                                "in read mode, not {}".format(self._gz_mode))
-        return self.name, self.tell()
+        return {"name_val": self.name,
+                "tell_val": self.tell()}
 
     def __setstate__(self, args):
-        super().__init__(args[0])
-        self.seek(args[1])
+        super().__init__(args["name_val"])
+        try:
+            self.seek(args["tell_val"])
+        except KeyError:
+            pass
 
 
 def pickle_open(name, mode='rt'):