WIP: Reader that does not need open file handles (#239)

package/MDAnalysis/coordinates/XYZ.py

@@ -255,7 +255,7 @@ def write_next_timestep(self, ts=None):
                             "".format(atom, x, y, z))
 
 
-class XYZReader(base.Reader):
+class XYZReader(base.NewReader):
     """Reads from an XYZ file
 
     :Data:
@@ -299,7 +299,7 @@ def __init__(self, filename, **kwargs):
         # coordinates::core.py so the last file extension will tell us if it is
         # bzipped or not
         root, ext = os.path.splitext(self.filename)
-        self.xyzfile = util.anyopen(self.filename, "r")
+
         self.compression = ext[1:] if ext[1:] != "xyz" else None
         self._cache = dict()
 
@@ -308,7 +308,7 @@ def __init__(self, filename, **kwargs):
         # etc.
         # (Also cannot just use seek() or reset() because that would break
         # with urllib2.urlopen() streams)
-        self._read_next_timestep()
+        self.next()
 
     @property
     @cached('n_atoms')
@@ -320,12 +320,13 @@ def n_atoms(self):
         return int(n)
 
     @property
-    @cached('n_frames')
     def n_frames(self):
-        try:
-            return self._read_xyz_n_frames()
-        except IOError:
-            return 0
+        return len(self._offsets)
+
+    @property
+    @cached('offsets')
+    def _offsets(self):
+        return self._read_xyz_n_frames()
 
     def _read_xyz_n_frames(self):
         # the number of lines in the XYZ file will be 2 greater than the
@@ -341,14 +342,12 @@ def _read_xyz_n_frames(self):
                     offsets.append(f.tell())
                 line = f.readline()
                 counter += 1
+        offsets.pop()  # last offset is end of file
 
-        # need to check this is an integer!
-        n_frames = int(counter / linesPerFrame)
-        self._offsets = offsets
-        return n_frames
+        return offsets
 
     def _read_frame(self, frame):
-        self.xyzfile.seek(self._offsets[frame])
+        self._file.seek(self._offsets[frame])
         self.ts.frame = frame - 1  # gets +1'd in next
         return self._read_next_timestep()
 
@@ -357,7 +356,7 @@ def _read_next_timestep(self, ts=None):
         if ts is None:
             ts = self.ts
 
-        f = self.xyzfile
+        f = self._file
 
         try:
             # we assume that there are only two header lines per frame
@@ -371,21 +370,9 @@ def _read_next_timestep(self, ts=None):
             raise EOFError(err)
 
     def _reopen(self):
-        self.close()
-        self.open_trajectory()
-
-    def open_trajectory(self):
-        if self.xyzfile is not None:
-            raise IOError(
-                errno.EALREADY, 'XYZ file already opened', self.filename)
-
-        self.xyzfile = util.anyopen(self.filename, "r")
-
-        # reset ts
-        ts = self.ts
-        ts.frame = -1
-
-        return self.xyzfile
+        """Reposition (the virtual fh) to just before first frame"""
+        self._last_fh_pos = 0
+        self.ts.frame = -1
 
     def Writer(self, filename, n_atoms=None, **kwargs):
         """Returns a XYZWriter for *filename* with the same parameters as this
@@ -412,10 +399,3 @@ def Writer(self, filename, n_atoms=None, **kwargs):
         if n_atoms is None:
             n_atoms = self.n_atoms
         return XYZWriter(filename, n_atoms=n_atoms, **kwargs)
-
-    def close(self):
-        """Close xyz trajectory file if it was open."""
-        if self.xyzfile is None:
-            return
-        self.xyzfile.close()
-        self.xyzfile = None

package/MDAnalysis/coordinates/base.py

@@ -126,11 +126,12 @@
 
 from six.moves import range
 import six
-
-import warnings
-
-import numpy as np
+from contextlib import contextmanager
 import copy
+import itertools
+import numpy as np
+import os.path
+import warnings
 import weakref
 
 from .. import (
@@ -141,6 +142,7 @@
 from ..core import flags
 from .. import units
 from ..lib.util import asiterable, Namespace
+from ..lib import util
 from . import core
 from .. import NoDataError
 
@@ -1630,6 +1632,100 @@ def __del__(self):
         self.close()
 
 
+class NewReader(ProtoReader):
+    """Super top secret magical class that fixes all known (and unknown) bugs"""
+    def __init__(self, filename, convert_units=None, **kwargs):
+        self.filename = filename
+        self._last_fh_pos = 0  # last known position of file handle
+
+        if convert_units is None:
+            convert_units = flags['convert_lengths']
+        self.convert_units = convert_units
+
+        self._auxs = {}
+        ts_kwargs = {}
+        for att in ('dt', 'time_offset'):
+            try:
+                val = kwargs[att]
+            except KeyError:
+                pass
+            else:
+                ts_kwargs[att] = val
+
+        self._ts_kwargs = ts_kwargs
+
+    def _full_iter(self):
+        self._reopen()
+        with util.openany(self.filename, 'r') as self._file:
+            while True:
+                try:
+                    ts = self._read_next_timestep()
+                    for auxname in self.aux_list:
+                        ts = self._auxs[auxname].update_ts(ts)
+                    yield ts
+                except (EOFError, IOError):
+                    # goto start of file
+                    self._read_frame_with_aux(0)
+                    self._last_fh_pos = self._file.tell()
+                    raise StopIteration
+
+    def _sliced_iter(self, start, stop, step):
+        with util.openany(self.filename, 'r') as self._file:
+            for f in range(start, stop, step):
+                yield self._read_frame_with_aux(f)
+
+            self._read_frame_with_aux(0)
+            self._last_fh_pos = self._file.tell()
+        raise StopIteration
+
+    def _goto_frame(self, i):
+        with util.openany(self.filename, 'r') as self._file:
+            ts = self._read_frame_with_aux(i)
+        return ts
+
+    def __iter__(self):
+        return self._full_iter()
+
+    def __getitem__(self, item):
+        def apply_limits(frame):
+            if frame < 0:
+                frame += len(self)
+            if frame < 0 or frame >= len(self):
+                raise IndexError("Index {} exceeds length of trajectory ({})."
+                                 "".format(frame, len(self)))
+            return frame
+
+        if isinstance(item, int):
+            return self._goto_frame(apply_limits(item))
+        elif isinstance(item, (list, np.ndarray)):
+            return self._sliced_iter(item)
+        elif isinstance(item, slice):  # TODO Fix me!
+            start, stop, step = self.check_slice_indices(
+                item.start, item.stop, item.step)
+
+            return self._sliced_iter(start, stop, step)
+
+    def __next__(self):
+        with util.openany(self.filename, 'r') as self._file:
+            try:
+                self._file.seek(self._last_fh_pos)
+                ts = self._read_next_timestep()
+            except (EOFError, IOError):
+                self.rewind()
+                raise StopIteration
+            else:
+                self._last_fh_pos = self._file.tell()
+                for auxname in self.aux_list:
+                    ts = self._auxs[auxname].update_ts(ts)
+                return ts
+
+    next = __next__
+
+    def close(self):
+        # "close" by moving last position to start
+        self._last_fh_pos = 0
+
+
 class _Writermeta(type):
     # Auto register upon class creation
     def __init__(cls, name, bases, classdict):

package/MDAnalysis/lib/formats/libmdaxdr.pyx

@@ -122,20 +122,23 @@ cdef class _XDRFile:
     This class can't be initialized use one of the subclasses XTCFile, TRRFile
     """
     cdef readonly int n_atoms
-    cdef int is_open
+    cdef readonly int is_open
     cdef int reached_eof
     cdef XDRFILE *xfp
     cdef readonly fname
     cdef int current_frame
+    cdef str _mode
     cdef str mode
     cdef np.ndarray box
     cdef np.ndarray _offsets
     cdef int _has_offsets
 
     def __cinit__(self, fname, mode='r'):
         self.fname = fname.encode('utf-8')
+        self._mode = mode
         self.is_open = False
-        self.open(self.fname, mode)
+        self.open()  # to read natoms
+        self.close()
 
     def __dealloc__(self):
         self.close()
@@ -210,6 +213,7 @@ cdef class _XDRFile:
 
     def __enter__(self):
         """Support context manager"""
+        self.open(self.fname, self._mode)
         return self
 
     def __exit__(self, exc_type, exc_val, exc_tb):
@@ -318,7 +322,10 @@ cdef class _XDRFile:
         set_offsets
         """
         if not self._has_offsets:
+            # wrap this in try/except to ensure it gets closed?
+            self.open(self.fname, self._mode)
             self._offsets = self.calc_offsets()
+            self.close()
             self._has_offsets = True
         return self._offsets
 

testsuite/MDAnalysisTests/coordinates/test_xyz.py

@@ -9,10 +9,12 @@
     assert_array_equal,
     assert_,
 )
+import mock
 
+from MDAnalysis.lib.util import openany
 from MDAnalysis.coordinates.XYZ import XYZWriter
 
-from MDAnalysisTests.datafiles import COORDINATES_XYZ, COORDINATES_XYZ_BZ2
+from MDAnalysisTests.datafiles import XYZ, COORDINATES_XYZ, COORDINATES_XYZ_BZ2
 from MDAnalysisTests.coordinates.base import (BaseReaderTest, BaseReference,
                                               BaseWriterTest)
 from MDAnalysisTests.core.groupbase import make_Universe, FakeReader
@@ -148,3 +150,26 @@ def test_list_names(self):
 
         u2 = mda.Universe(self.outfile)
         assert_(all(u2.atoms.names == names))
+
+
+class TestXYZReaderFH(object):
+    # tests for the new filehandle behaviour (issue #239)
+    def setUp(self):
+        self.u = mda.Universe(XYZ)
+
+    def tearDown(self):
+        del self.u
+
+    #@mock.patch('MDAnalysis.lib.util.openany', wraps=openany)
+    def test_iteration(self):
+        # check filehandle is closed
+        assert_(self.u.trajectory._file.closed)
+
+        with mock.patch('MDAnalysis.lib.util.openany', wraps=openany) as mockopen:
+            for ts in self.u.trajectory:
+                print ts.frame
+
+            assert_(mockopen.call_count == 1)
+
+        # check filehandle is closed
+        assert_(self.u.trajectory._file.closed)