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'):