Repository: pauldmccarthy/indexed_gzip (language: python)

Issue:
Can't create indexes from un-seekable fileobj's
It would be good to be able to create indexes from a fileobj that is un-seekable at the time of index creation. @pauldmccarthy do you think that is possible given the code, or does the code rely fundamentally on seeking somewhere to create the index?
Support index creation for unseekable file objects
Fixes https://github.com/pauldmccarthy/indexed_gzip/issues/102.


Discussion / hints:

Fix patch:
diff --git a/CHANGELOG.md b/CHANGELOG.md
index 0c14b6dc..a497b06a 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -1,6 +1,13 @@
 # `indexed_gzip` changelog
 
 
+## 1.7.0 (September 12th 2022)
+
+
+* Changes to allow an index to be built from file-likes which don't support
+  `seek()` or `tell()` operations (#103, #105).
+
+
 ## 1.6.13 (April 14th 2022)
 
 
diff --git a/indexed_gzip/__init__.py b/indexed_gzip/__init__.py
index 51dc4f10..7f272828 100644
--- a/indexed_gzip/__init__.py
+++ b/indexed_gzip/__init__.py
@@ -19,4 +19,4 @@
 """
 
 
-__version__ = '1.6.13'
+__version__ = '1.7.0'
diff --git a/indexed_gzip/zran.c b/indexed_gzip/zran.c
index fde122f8..73e749f5 100644
--- a/indexed_gzip/zran.c
+++ b/indexed_gzip/zran.c
@@ -274,8 +274,9 @@ int ZRAN_EXPAND_INDEX_CRC_ERROR = -2;
  * ZRAN_EXPAND_INDEX_FAIL.
  */
 static int _zran_expand_index(
-    zran_index_t *index, /* The index                      */
-    uint64_t      until  /* Expand the index to this point */
+    zran_index_t *index, /* The index */
+    uint64_t      until  /* Expand the index to this point. If 0,
+                            expand the index until EOF is reached. */
 );
 
 
@@ -635,7 +636,7 @@ int zran_init(zran_index_t *index,
      * window size.
      */
     if (spacing <= window_size)
-      goto fail;
+        goto fail;
 
     /* The file must be opened in read-only mode */
     if (!is_readonly(fd, f))
@@ -644,16 +645,26 @@ int zran_init(zran_index_t *index,
     /*
      * Calculate the size of the compressed file
      */
-    if (fseek_(fd, f, 0, SEEK_END) != 0)
-        goto fail;
+    if (seekable_(fd, f)) {
+        if (fseek_(fd, f, 0, SEEK_END) != 0)
+            goto fail;
 
-    compressed_size = ftell_(fd, f);
+        compressed_size = ftell_(fd, f);
 
-    if (compressed_size < 0)
-        goto fail;
+        if (compressed_size < 0)
+            goto fail;
 
-    if (fseek_(fd, f, 0, SEEK_SET) != 0)
-        goto fail;
+        if (fseek_(fd, f, 0, SEEK_SET) != 0)
+            goto fail;
+    } else {
+        /*
+         * File is not seekable, so don't calculate
+         * compressed_size.  It will be updated in
+         * _zran_read_data_from_file when the EOF
+         * is reached
+         */
+        compressed_size = 0;
+    }
 
     /*
      * Allocate some initial space
@@ -824,9 +835,6 @@ int zran_build_index(zran_index_t *index, uint64_t from, uint64_t until)
     if (_zran_invalidate_index(index, from) != 0)
         return ZRAN_BUILD_INDEX_FAIL;
 
-    if (until == 0)
-      until = index->compressed_size;
-
     return _zran_expand_index(index, until);
 }
 
@@ -850,9 +858,11 @@ int _zran_get_point_at(
 
     /*
      * Bad input - past the end of the compressed or
-     * uncompressed streams (if the latter is known).
+     * uncompressed streams (if their sizes are known).
      */
-    if (compressed && offset >= index->compressed_size)
+    if (compressed                 &&
+        index->compressed_size > 0 &&
+        offset >= index->compressed_size)
         goto eof;
 
     if (!compressed                  &&
@@ -1361,9 +1371,8 @@ static int _zran_read_data_from_file(zran_index_t *index,
         memmove(index->readbuf, stream->next_in, stream->avail_in);
     }
 
-    zran_log("Reading from file %llu [== %llu?] "
+    zran_log("Reading from file %llu "
              "[into readbuf offset %u]\n",
-             ftell_(index->fd, index->f),
              cmp_offset + stream->avail_in,
              stream->avail_in);
 
@@ -1411,6 +1420,11 @@ static int _zran_read_data_from_file(zran_index_t *index,
                          "size: %llu\n", uncmp_offset);
                 index->uncompressed_size = uncmp_offset;
             }
+            if (index->compressed_size == 0) {
+                zran_log("Updating compressed data "
+                         "size: %llu\n", cmp_offset);
+                index->compressed_size = cmp_offset + 8;
+            }
             goto eof;
         }
         /*
@@ -1784,8 +1798,15 @@ static int _zran_inflate(zran_index_t *index,
          * to the correct location in the file for us.
          */
         if (start == NULL) {
-            if (fseek_(index->fd, index->f, 0, SEEK_SET) != 0) {
-                goto fail;
+            /*
+             * If file is not seekable, assume that
+             * it is positioned at the beginning of
+             * the stream.
+             */
+            if (seekable_(index->fd, index->f)) {
+                if (fseek_(index->fd, index->f, 0, SEEK_SET) != 0) {
+                    goto fail;
+                }
             }
 
             /*
@@ -2298,7 +2319,7 @@ int _zran_expand_index(zran_index_t *index, uint64_t until) {
          * The index already covers the requested
          * offset. Nothing needs to be done.
          */
-        if (until <= start->cmp_offset)
+        if (until != 0 && until <= start->cmp_offset)
             return 0;
     }
 
@@ -2315,10 +2336,10 @@ int _zran_expand_index(zran_index_t *index, uint64_t until) {
 
     /*
      * If the caller passed until == 0,
-     * we force some data to be read.
+     * we expand to EOF.
      */
     if (until == 0) {
-      until = index->spacing;
+        until = UINT64_MAX;
     }
 
     /*
@@ -2328,7 +2349,6 @@ int _zran_expand_index(zran_index_t *index, uint64_t until) {
      * points in the index.
      */
     if (start != NULL) {
-
         cmp_offset        = start->cmp_offset;
         uncmp_offset      = start->uncmp_offset;
         last_uncmp_offset = uncmp_offset;
@@ -2341,14 +2361,13 @@ int _zran_expand_index(zran_index_t *index, uint64_t until) {
 
     /*
      * Don't finish until we're at the end of the
-     * file, or we've expanded the index past
-     * the requested offset (and have created at
-     * least one new index point -
-     * last_created == NULL tells us whether a
-     * point has been created).
-     */
-    while ((cmp_offset < index->compressed_size) &&
-           (last_created == NULL || last_created->cmp_offset < until)) {
+     * file (break at bottom of loop), or we've
+     * expanded the index past the requested
+     * offset (and have created at least one new
+     * index point - last_created == NULL tells
+     * us whether a point has been created).
+     */
+    while (last_created == NULL || last_created->cmp_offset < until) {
 
         /*
          * On the first call to _zran_inflate, we
diff --git a/indexed_gzip/zran_file_util.c b/indexed_gzip/zran_file_util.c
index f711dc98..cb01dde1 100644
--- a/indexed_gzip/zran_file_util.c
+++ b/indexed_gzip/zran_file_util.c
@@ -231,6 +231,48 @@ int _getc_python(PyObject *f) {
     return buf;
 }
 
+/*
+ * Calls the .seekable() method on Python file-like objects.
+ */
+int _seekable_python(PyObject *f) {
+    PyObject *data = NULL;
+    int64_t   result;
+
+    _ZRAN_FILE_UTIL_ACQUIRE_GIL
+
+    data = PyObject_CallMethod(f, "seekable", NULL);
+    if (data == NULL)
+        goto fail;
+
+    result = PyLong_AsLong(data);
+    if (result == -1 && PyErr_Occurred())
+        goto fail;
+
+    Py_DECREF(data);
+    _ZRAN_FILE_UTIL_RELEASE_GIL
+    return result;
+
+fail:
+    Py_XDECREF(data);
+    _ZRAN_FILE_UTIL_RELEASE_GIL
+    return -1;
+}
+
+
+/*
+ * Hacky implementation of seekable() for Python 2.
+ * If f.tell() returns an error, assume that the
+ * object is unseekable.
+ */
+int _seekable_python2(PyObject *f) {
+    int64_t ret;
+    ret = _ftell_python(f);
+    if (ret < 0) {
+        PyErr_Clear();
+    }
+    return ret >= 0;
+}
+
 /*
  * Calls ferror on fd if specified, otherwise the Python-specific method on f.
  */
@@ -299,3 +341,15 @@ size_t fwrite_(const void *ptr,
 int getc_(FILE *fd, PyObject *f) {
     return fd != NULL ? getc(fd): _getc_python(f);
 }
+
+/*
+ * Returns whether the given file is seekable. If fd is specified, assumes it's always seekable.
+ * If f is specified, calls f.seekable() to see if the Python file object is seekable.
+ */
+int seekable_(FILE *fd, PyObject *f) {
+    #if PY_MAJOR_VERSION > 2
+    return fd != NULL ? 1: _seekable_python(f);
+    #else
+    return fd != NULL ? 1: _seekable_python2(f);
+    #endif
+}
diff --git a/indexed_gzip/zran_file_util.h b/indexed_gzip/zran_file_util.h
index 5bd8bd8a..28747096 100644
--- a/indexed_gzip/zran_file_util.h
+++ b/indexed_gzip/zran_file_util.h
@@ -60,6 +60,11 @@ size_t _fwrite_python(const void *ptr, size_t size, size_t nmemb, PyObject *f);
  */
 int _getc_python(PyObject *f);
 
+/*
+ * Calls the .seekable() method on Python file-like objects.
+ */
+int _seekable_python(PyObject *f);
+
 /*
  * Calls ferror on fd if specified, otherwise the Python-specific method on f.
  */
@@ -101,5 +106,10 @@ size_t fwrite_(
  */
 int getc_(FILE *fd, PyObject *f);
 
+/*
+ * Returns whether the given file is seekable. If fd is specified, assumes it's always seekable.
+ * If f is specified, calls f.seekable() to see if the Python file object is seekable.
+ */
+int seekable_(FILE *fd, PyObject *f);
 
 #endif /* __ZRAN_FILE_UTIL_H__ */
diff --git a/indexed_gzip/zran_file_util.pxd b/indexed_gzip/zran_file_util.pxd
index 8c04337d..30ef41af 100644
--- a/indexed_gzip/zran_file_util.pxd
+++ b/indexed_gzip/zran_file_util.pxd
@@ -26,6 +26,8 @@ cdef extern from "zran_file_util.h":
     
     int _getc_python(PyObject *f)
 
+    int _seekable_python(PyObject *f)
+
     int ferror_(FILE *fd, PyObject *f)
 
     int fseek_(FILE *fd, PyObject *f, int64_t offset, int whence)
@@ -41,3 +43,5 @@ cdef extern from "zran_file_util.h":
     size_t fwrite_(const void *ptr, size_t size, size_t nmemb, FILE *fd, PyObject *f)
 
     int getc_(FILE *fd, PyObject *f)
+
+    int seekable_(FILE *fd, PyObject *f)


Test patch:
diff --git a/indexed_gzip/tests/__init__.py b/indexed_gzip/tests/__init__.py
index aac834c2..25d90966 100644
--- a/indexed_gzip/tests/__init__.py
+++ b/indexed_gzip/tests/__init__.py
@@ -39,6 +39,11 @@ def __exit__(self, *a, **kwa):
     return ctx()
 
 
+def touch(path):
+    """Create an empty file."""
+    with open(path, 'wt') as f:
+        pass
+
 
 def poll(until):
     """Waits until ``until`` returns ``True``, printing out a message every
diff --git a/indexed_gzip/tests/ctest_indexed_gzip.pyx b/indexed_gzip/tests/ctest_indexed_gzip.pyx
index 39da7bf9..1c22848c 100644
--- a/indexed_gzip/tests/ctest_indexed_gzip.pyx
+++ b/indexed_gzip/tests/ctest_indexed_gzip.pyx
@@ -33,6 +33,7 @@ import indexed_gzip as igzip
 from . import gen_test_data
 from . import check_data_valid
 from . import tempdir
+from . import touch
 from . import compress
 
 from libc.stdio cimport (SEEK_SET,
@@ -804,7 +805,7 @@ def test_get_index_seek_points():
         # make a test file
         data = np.arange(spacing, dtype=np.uint64)
         with gzip.open(fname, 'wb') as f:
-            f.write(data.tostring())
+            f.write(data.tobytes())
 
         # check points before and after index creation
         with igzip._IndexedGzipFile(fname, spacing=spacing) as f:
@@ -828,13 +829,14 @@ def test_import_export_index():
         idxfname = op.join(td, 'test.gzidx')
 
         # make a test file
-        data = np.arange(65536, dtype=np.uint64)
+        data = np.arange(524288, dtype=np.uint64)
         with gzip.open(fname, 'wb') as f:
-            f.write(data.tostring())
+            f.write(data.tobytes())
 
         # generate an index file
-        with igzip._IndexedGzipFile(fname) as f:
+        with igzip._IndexedGzipFile(fname, spacing=131072) as f:
             f.build_full_index()
+            points = list(f.seek_points())
             f.export_index(idxfname)
 
         # Check that index file works via __init__
@@ -842,6 +844,7 @@ def test_import_export_index():
             f.seek(65535 * 8)
             val = np.frombuffer(f.read(8), dtype=np.uint64)
             assert val[0] == 65535
+            assert points == list(f.seek_points())
 
         # Check that index file works via import_index
         with igzip._IndexedGzipFile(fname) as f:
@@ -849,13 +852,28 @@ def test_import_export_index():
             f.seek(65535 * 8)
             val = np.frombuffer(f.read(8), dtype=np.uint64)
             assert val[0] == 65535
+            assert points == list(f.seek_points())
+
+
+def test_import_export_index_open_file():
+
+    with tempdir() as td:
+        fname    = op.join(td, 'test.gz')
+        idxfname = op.join(td, 'test.gzidx')
+
+        # make a test file
+        data = np.arange(524288, dtype=np.uint64)
+        with gzip.open(fname, 'wb') as f:
+            f.write(data.tobytes())
 
         # generate an index file from open file handle
-        with igzip._IndexedGzipFile(fname) as f:
+        with igzip._IndexedGzipFile(fname, spacing=131072) as f:
             f.build_full_index()
+            points = list(f.seek_points())
 
             # Should raise if wrong permissions
             with pytest.raises(ValueError):
+                touch(idxfname)
                 with open(idxfname, 'rb') as idxf:
                     f.export_index(fileobj=idxf)
 
@@ -877,17 +895,63 @@ def test_import_export_index():
             f.seek(65535 * 8)
             val = np.frombuffer(f.read(8), dtype=np.uint64)
             assert val[0] == 65535
+            assert points == list(f.seek_points())
 
         # Test exporting to / importing from a file-like object
         idxf = BytesIO()
-        with igzip._IndexedGzipFile(fname) as f:
+        with igzip._IndexedGzipFile(fname, spacing=131072) as f:
+            f.build_full_index()
             f.export_index(fileobj=idxf)
+            points = list(f.seek_points())
         idxf.seek(0)
         with igzip._IndexedGzipFile(fname) as f:
             f.import_index(fileobj=idxf)
             f.seek(65535 * 8)
             val = np.frombuffer(f.read(8), dtype=np.uint64)
             assert val[0] == 65535
+            assert points == list(f.seek_points())
+
+
+def test_build_index_from_unseekable():
+    with tempdir() as td:
+        fname    = op.join(td, 'test.gz')
+        idxfname = op.join(td, 'test.gzidx')
+
+        # make a test file
+        data = np.arange(524288, dtype=np.uint64)
+        with gzip.open(fname, 'wb') as f:
+            f.write(data.tostring())
+
+        # Test creating the index when file is unseekable,
+        # then using the index when file is seekable.
+        with open(fname, 'rb') as f:
+            b = f.read()
+            fileobj = BytesIO(b)
+
+        def new_seek(*args, **kwargs):
+            raise OSError()
+        def new_tell(*args, **kwargs):
+            raise OSError()
+        old_seek = fileobj.seek
+        old_tell = fileobj.tell
+        fileobj.seekable = lambda: False
+        fileobj.seek = new_seek
+        fileobj.tell = new_tell
+        # generate an index file
+        with igzip._IndexedGzipFile(fileobj, spacing=131072) as f:
+            f.build_full_index()
+            f.export_index(idxfname)
+            points = list(f.seek_points())
+        fileobj.seek = old_seek
+        fileobj.tell = old_tell
+        fileobj.seekable = lambda: True
+        fileobj.seek(0)
+        # Check that index file works via __init__
+        with igzip._IndexedGzipFile(fileobj, index_file=idxfname) as f:
+            f.seek(65535 * 8)
+            val = np.frombuffer(f.read(8), dtype=np.uint64)
+            assert val[0] == 65535
+            assert points == list(f.seek_points())
 
 
 def test_wrapper_class():
diff --git a/indexed_gzip/tests/ctest_zran.pyx b/indexed_gzip/tests/ctest_zran.pyx
index 913846f1..cac4f12d 100644
--- a/indexed_gzip/tests/ctest_zran.pyx
+++ b/indexed_gzip/tests/ctest_zran.pyx
@@ -369,6 +369,14 @@ def test_getc():
         PyErr_Clear()
 
 
+def test_seekable():
+
+    f = BytesIO(b"abc")
+    assert zran_file_util._seekable_python(<PyObject*>f) == 1
+    f.seekable = lambda: False
+    assert zran_file_util._seekable_python(<PyObject*>f) == 0
+
+
 def test_init(testfile, no_fds):
     """Tests a bunch of permutations of the parameters to zran_init. """
 
diff --git a/indexed_gzip/tests/test_indexed_gzip.py b/indexed_gzip/tests/test_indexed_gzip.py
index fde13505..c1c85a28 100644
--- a/indexed_gzip/tests/test_indexed_gzip.py
+++ b/indexed_gzip/tests/test_indexed_gzip.py
@@ -160,6 +160,12 @@ def test_get_index_seek_points():
 def test_import_export_index():
     ctest_indexed_gzip.test_import_export_index()
 
+def test_import_export_index_open_file():
+    ctest_indexed_gzip.test_import_export_index_open_file()
+
+def test_build_index_from_unseekable():
+    ctest_indexed_gzip.test_build_index_from_unseekable()
+
 def test_wrapper_class():
     ctest_indexed_gzip.test_wrapper_class()
 
diff --git a/indexed_gzip/tests/test_zran.py b/indexed_gzip/tests/test_zran.py
index ca7aae2f..2c4a56eb 100644
--- a/indexed_gzip/tests/test_zran.py
+++ b/indexed_gzip/tests/test_zran.py
@@ -43,6 +43,9 @@ def test_fwrite():
     def test_getc():
         ctest_zran.test_getc()
 
+    def test_seekable():
+        ctest_zran.test_seekable()
+
     def test_init(testfile):
         for no_fds in (True, False):
             ctest_zran.test_init(testfile, no_fds)