No need to redirect to the original frame explicitly during deserialization

package/CHANGELOG

@@ -15,7 +15,8 @@ The rules for this file:
 
 -------------------------------------------------------------------------------
 ??/??/?? IAlibay, ianmkenney, PicoCentauri, pgbarletta, p-j-smith, 
-         richardjgowers, lilyminium, ALescoulie, hmacdope, HeetVekariya
+         richardjgowers, lilyminium, ALescoulie, hmacdope, HeetVekariya,
+         yuxuanzhuang
 
  * 2.7.0
 
@@ -48,6 +49,8 @@ Enhancements
     (Issue #3994, PR #4281)  
   * Add support for reading chainID info from Autodock PDBQT files (Issue #4207,
     PR #4284)
+  * Improve performance in Universe serialization by no explicit redirection to
+    the timestep. (Issue #3721, PR #3722)
 
 Changes
   * Biopython is now an optional dependency (Issue #3820, PR #4332)
@@ -198,6 +201,10 @@ Fixes
   * Fix MSD docs to use the correct error metric in example (Issue #3991)
   * Add 'PairIJ Coeffs' to the list of sections in LAMMPSParser.py
     (Issue #3336)
+  * Fix failure in double-serialization of TextIOPicklable file reader.
+    (Issue #3723, PR #3722)
+  * Fix failure to preserve modification of coordinates after serialization.
+    (Issue #3723, PR #3722)
 
 Enhancements
   * ARM64 (osx and linux) wheels now uploaded on release (Issue #4054)
@@ -364,6 +371,8 @@ Changes
     use `pip install ./package[extra_formats]` instead (PR #3810)
   * `Universe.empty` emmits less warnings (PR #3814)
   * adding element attribute to TXYZParser if all atom names are valid element symbols (PR #3826)
+  * TextIOPicklable serializes the raw class instance instead of class name.
+    (PR #3722)
 
 Deprecations
   * Direct indexing of results in analysis.nucleicacids' NucPairDist

package/MDAnalysis/coordinates/H5MD.py

@@ -828,7 +828,6 @@ def __setstate__(self, state):
         self.__dict__ = state
         self._particle_group = self._file['particles'][
                                list(self._file['particles'])[0]]
-        self[self.ts.frame]
 
 
 class H5MDWriter(base.WriterBase):

package/MDAnalysis/coordinates/base.py

@@ -659,6 +659,9 @@ class ProtoReader(IOBase, metaclass=_Readermeta):
     .. versionchanged:: 2.0.0
        Now supports (un)pickle. Upon unpickling,
        the current timestep is retained by reconstrunction.
+    .. versionchanged:: 2.7.0
+       the modification of coordinates was preserved
+       after serialization.
     """
 
     #: The appropriate Timestep class, e.g.
@@ -1442,10 +1445,6 @@ def _apply_transformations(self, ts):
 
         return ts
 
-    def __setstate__(self, state):
-        self.__dict__ = state
-        self[self.ts.frame]
-
 
 class ReaderBase(ProtoReader):
     """Base class for trajectory readers that extends :class:`ProtoReader` with a

package/MDAnalysis/lib/picklable_file_io.py

@@ -112,12 +112,13 @@
 
     def __setstate__(self, state):
         name = state["name_val"]
-        super().__init__(name, mode='r')
+        self.__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 "
@@ -163,7 +164,7 @@
         raw_class = state["raw_class"]
         name = state["name_val"]
         raw = raw_class(name)
-        super().__init__(raw)
+        self.__init__(raw)
         self.seek(state["tell_val"])
 
     def __reduce_ex__(self, prot):
@@ -175,18 +176,13 @@
                  "name_val": self.name,
                  "tell_val": self.tell()})
 
+
 class TextIOPicklable(io.TextIOWrapper):
     """Character and line based picklable file-like object.
 
     This class provides a file-like :class:`io.TextIOWrapper` object that can
     be pickled. Note that this only works in read mode.
 
-    Note
-    ----
-    After pickling, the current position is reset. `universe.trajectory[i]` has
-    to be used to return to its original frame.
-
-
     Parameters
     ----------
     raw : FileIO object
@@ -205,21 +201,34 @@
 
 
     .. versionadded:: 2.0.0
+    .. versionchanged:: 2.7.0
+       The raw class instance instead of the class name
+       that is wrapped inside will be serialized.
+       After deserialization, the current position is no longer reset
+       so `universe.trajectory[i]` is not needed to seek to the
+       original position.
     """
     def __init__(self, raw):
         super().__init__(raw)
         self.raw_class = raw.__class__
 
-
     def __setstate__(self, args):
         raw_class = args["raw_class"]
         name = args["name_val"]
+        tell = args["tell_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)
+        self.__init__(raw)
+        if tell is not None:
+            self.seek(tell)
 
     def __reduce_ex__(self, prot):
+        try:
+            curr_loc = self.tell()
+        # some readers (e.g. GMS) disable tell() due to using next()
+        except OSError:
+            curr_loc = None
         try:
             name = self.name
         except AttributeError:
@@ -228,7 +237,8 @@
         return (self.__class__.__new__,
                 (self.__class__,),
                 {"raw_class": self.raw_class,
-                 "name_val": name})
+                 "name_val": name,
+                 "tell_val": curr_loc})
 
 
 class BZ2Picklable(bz2.BZ2File):
@@ -289,9 +299,11 @@
         return {"name_val": self._fp.name, "tell_val": self.tell()}
 
     def __setstate__(self, args):
-        super().__init__(args["name_val"])
+        name = args["name_val"]
+        tell = args["tell_val"]
+        self.__init__(name)
         try:
-            self.seek(args["tell_val"])
+            self.seek(tell)
         except KeyError:
             pass
 
@@ -355,9 +367,11 @@
                 "tell_val": self.tell()}
 
     def __setstate__(self, args):
-        super().__init__(args["name_val"])
+        name = args["name_val"]
+        tell = args["tell_val"]
+        self.__init__(name)
         try:
-            self.seek(args["tell_val"])
+            self.seek(tell)
         except KeyError:
             pass
 

testsuite/MDAnalysisTests/coordinates/base.py

@@ -122,9 +122,17 @@ def test_last_slice(self):
     def test_pickle_singleframe_reader(self):
         reader = self.universe.trajectory
         reader_p = pickle.loads(pickle.dumps(reader))
+        reader_p_p = pickle.loads(pickle.dumps(reader_p))
         assert_equal(len(reader), len(reader_p))
         assert_equal(reader.ts, reader_p.ts,
                      "Single-frame timestep is changed after pickling")
+        assert_equal(len(reader), len(reader_p_p))
+        assert_equal(reader.ts, reader_p_p.ts,
+                     "Single-frame timestep is changed after double pickling")
+        reader.ts.positions[0] = np.array([1, 2, 3])
+        reader_p = pickle.loads(pickle.dumps(reader))
+        assert_equal(reader.ts, reader_p.ts,
+                     "Modification of coords is not preserved after serialization")
 
 
 class BaseReference(object):
@@ -443,6 +451,14 @@ def test_pickle_reader(self, reader):
         assert_equal(len(reader), len(reader_p))
         assert_equal(reader.ts, reader_p.ts,
                      "Timestep is changed after pickling")
+        reader_p_p = pickle.loads(pickle.dumps(reader_p))
+        assert_equal(len(reader), len(reader_p_p))
+        assert_equal(reader.ts, reader_p_p.ts,
+                     "Timestep is changed after double pickling")
+        reader.ts.positions[0] = np.array([1, 2, 3])
+        reader_p = pickle.loads(pickle.dumps(reader))
+        assert_equal(reader.ts, reader_p.ts,
+                     "Modification of coords is not preserved after serialization")
 
     def test_frame_collect_all_same(self, reader):
         # check that the timestep resets so that the base reference is the same
@@ -604,6 +620,9 @@ def test_pickle_next_ts_reader(self, reader):
         reader_p = pickle.loads(pickle.dumps(reader))
         assert_equal(next(reader), next(reader_p),
                      "Next timestep is changed after pickling")
+        reader_p_p = pickle.loads(pickle.dumps(reader_p))
+        assert_equal(next(reader), next(reader_p_p),
+                     "Next timestep is changed after double pickling")
 
     #  To make sure pickle works for last frame.
     def test_pickle_last_ts_reader(self, reader):

testsuite/MDAnalysisTests/utils/test_pickleio.py

@@ -90,10 +90,10 @@ def test_iopickle_text(f_text):
     assert_equal(f_text.readlines(), f_text_pickled.readlines())
 
 
-def test_offset_text_to_0(f_text):
+def test_offset_text_same(f_text):
     f_text.readline()
     f_text_pickled = pickle.loads(pickle.dumps(f_text))
-    assert_equal(f_text_pickled.tell(), 0)
+    assert_equal(f_text_pickled.tell(), f_text.tell())
 
 
 @pytest.fixture(params=[