Skip to content

Commit

Permalink
MAINT: Python 3.12 support improve (#4300)
Browse files Browse the repository at this point in the history
* 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]
  • Loading branch information
tylerjereddy authored Oct 12, 2023
1 parent a522325 commit 8c41dc8
Showing 1 changed file with 52 additions and 28 deletions.
80 changes: 52 additions & 28 deletions package/MDAnalysis/lib/picklable_file_io.py
Original file line number Diff line number Diff line change
Expand Up @@ -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):
Expand Down Expand Up @@ -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.
Expand Down Expand Up @@ -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.
Expand Down Expand Up @@ -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):
Expand Down Expand Up @@ -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'):
Expand Down

0 comments on commit 8c41dc8

Please sign in to comment.