Merge pull request #4152 from pv/fromfile-fix

BUG: core: ensure file handle positions are in sync

8 files changed, 142 insertions, 14 deletions

diff --git a/numpy/core/bscript b/numpy/core/bscript
index ced37a530..a2222eb66 100644
--- a/numpy/core/bscript
+++ b/numpy/core/bscript
@@ -136,6 +136,11 @@ def type_checks(conf):
     NUMPYCONFIG_SYM.append(('SIZEOF_%s' % numpy.build_utils.waf.sanitize_string("Py_intptr_t"),
                            '%d' % size))
 
+    size = conf.check_type_size("off_t", header_name=header_name,
+                                expected_sizes=[4, 8], features=features)
+    NUMPYCONFIG_SYM.append(('SIZEOF_%s' % numpy.build_utils.waf.sanitize_string("off_t"),
+                           '%d' % size))
+
     # We check declaration AND type because that's how distutils does it.
     try:
         conf.check_declaration("PY_LONG_LONG", header_name=header_name,
diff --git a/numpy/core/include/numpy/_numpyconfig.h.in b/numpy/core/include/numpy/_numpyconfig.h.in
index 2cd389d44..2a159f56d 100644
--- a/numpy/core/include/numpy/_numpyconfig.h.in
+++ b/numpy/core/include/numpy/_numpyconfig.h.in
@@ -9,6 +9,7 @@
 #define NPY_SIZEOF_DOUBLE       @SIZEOF_DOUBLE@
 #define NPY_SIZEOF_LONGDOUBLE  	@SIZEOF_LONG_DOUBLE@
 #define NPY_SIZEOF_PY_INTPTR_T  @SIZEOF_PY_INTPTR_T@
+#define NPY_SIZEOF_OFF_T        @SIZEOF_OFF_T@
 
 #define NPY_SIZEOF_COMPLEX_FLOAT        @SIZEOF_COMPLEX_FLOAT@
 #define NPY_SIZEOF_COMPLEX_DOUBLE       @SIZEOF_COMPLEX_DOUBLE@
diff --git a/numpy/core/include/numpy/npy_3kcompat.h b/numpy/core/include/numpy/npy_3kcompat.h
index d0cd9ac1a..f4078dad2 100644
--- a/numpy/core/include/numpy/npy_3kcompat.h
+++ b/numpy/core/include/numpy/npy_3kcompat.h
@@ -146,12 +146,13 @@ PyUnicode_Concat2(PyObject **left, PyObject *right)
  * Get a FILE* handle to the file represented by the Python object
  */
 static NPY_INLINE FILE*
-npy_PyFile_Dup(PyObject *file, char *mode)
+npy_PyFile_Dup(PyObject *file, char *mode, npy_off_t *orig_pos)
 {
     int fd, fd2;
     PyObject *ret, *os;
-    Py_ssize_t pos;
+    npy_off_t pos;
     FILE *handle;
+
     /* Flush first to ensure things end up in the file in the correct order */
     ret = PyObject_CallMethod(file, "flush", "");
     if (ret == NULL) {
@@ -162,6 +163,9 @@ npy_PyFile_Dup(PyObject *file, char *mode)
     if (fd == -1) {
         return NULL;
     }
+
+    /* The handle needs to be dup'd because we have to call fclose
+       at the end */
     os = PyImport_ImportModule("os");
     if (os == NULL) {
         return NULL;
@@ -173,6 +177,8 @@ npy_PyFile_Dup(PyObject *file, char *mode)
     }
     fd2 = PyNumber_AsSsize_t(ret, NULL);
     Py_DECREF(ret);
+
+    /* Convert to FILE* handle */
 #ifdef _WIN32
     handle = _fdopen(fd2, mode);
 #else
@@ -182,6 +188,15 @@ npy_PyFile_Dup(PyObject *file, char *mode)
         PyErr_SetString(PyExc_IOError,
                         "Getting a FILE* from a Python file object failed");
     }
+
+    /* Record the original raw file handle position */
+    *orig_pos = npy_ftell(handle);
+    if (*orig_pos == -1) {
+        PyErr_SetString(PyExc_IOError, "obtaining file position failed");
+        return -1;
+    }
+
+    /* Seek raw handle to the Python-side position */
     ret = PyObject_CallMethod(file, "tell", "");
     if (ret == NULL) {
         fclose(handle);
@@ -193,7 +208,10 @@ npy_PyFile_Dup(PyObject *file, char *mode)
         fclose(handle);
         return NULL;
     }
-    npy_fseek(handle, pos, SEEK_SET);
+    if (npy_fseek(handle, pos, SEEK_SET) == -1) {
+        PyErr_SetString(PyExc_IOError, "seeking file failed");
+        return -1;
+    }
     return handle;
 }
 
@@ -201,14 +219,35 @@ npy_PyFile_Dup(PyObject *file, char *mode)
  * Close the dup-ed file handle, and seek the Python one to the current position
  */
 static NPY_INLINE int
-npy_PyFile_DupClose(PyObject *file, FILE* handle)
+npy_PyFile_DupClose(PyObject *file, FILE* handle, npy_off_t orig_pos)
 {
+    int fd;
     PyObject *ret;
-    Py_ssize_t position;
+    npy_off_t position;
+
     position = npy_ftell(handle);
+
+    /* Close the FILE* handle */
     fclose(handle);
 
-    ret = PyObject_CallMethod(file, "seek", NPY_SSIZE_T_PYFMT "i", position, 0);
+    /* Restore original file handle position, in order to not confuse
+       Python-side data structures */
+    fd = PyObject_AsFileDescriptor(file);
+    if (fd == -1) {
+        return -1;
+    }
+    if (npy_lseek(fd, orig_pos, SEEK_SET) == -1) {
+        PyErr_SetString(PyExc_IOError, "seeking file failed");
+        return -1;
+    }
+
+    if (position == -1) {
+        PyErr_SetString(PyExc_IOError, "obtaining file position failed");
+        return -1;
+    }
+
+    /* Seek Python-side handle to the FILE* handle position */
+    ret = PyObject_CallMethod(file, "seek", NPY_OFF_T_PYFMT "i", position, 0);
     if (ret == NULL) {
         return -1;
     }
@@ -230,8 +269,8 @@ npy_PyFile_Check(PyObject *file)
 
 #else
 
-#define npy_PyFile_Dup(file, mode) PyFile_AsFile(file)
-#define npy_PyFile_DupClose(file, handle) (0)
+#define npy_PyFile_Dup(file, mode, orig_pos_p) PyFile_AsFile(file)
+#define npy_PyFile_DupClose(file, handle, orig_pos) (0)
 #define npy_PyFile_Check PyFile_Check
 
 #endif
diff --git a/numpy/core/include/numpy/npy_common.h b/numpy/core/include/numpy/npy_common.h
index 08582bf79..15961e853 100644
--- a/numpy/core/include/numpy/npy_common.h
+++ b/numpy/core/include/numpy/npy_common.h
@@ -56,13 +56,40 @@
         #define NPY_INLINE
 #endif
 
-/* Enable 64 bit file position support on win-amd64. Ticket #1660 */
+/* 64 bit file position support, also on win-amd64. Ticket #1660 */
 #if defined(_MSC_VER) && defined(_WIN64) && (_MSC_VER > 1400)
+    #include <io.h>
     #define npy_fseek _fseeki64
     #define npy_ftell _ftelli64
+    #define npy_lseek _lseeki64
+    #define npy_off_t npy_int64
+
+    #if NPY_SIZEOF_INT == 8
+        #define NPY_OFF_T_PYFMT "i"
+    #elif NPY_SIZEOF_LONG == 8
+        #define NPY_OFF_T_PYFMT "l"
+    #elif NPY_SIZEOF_LONGLONG == 8
+        #define NPY_OFF_T_PYFMT "L"
+    #else
+        #error Unsupported size for type off_t
+    #endif
 #else
     #define npy_fseek fseek
     #define npy_ftell ftell
+    #define npy_lseek lseek
+    #define npy_off_t off_t
+
+    #if NPY_SIZEOF_OFF_T == NPY_SIZEOF_SHORT
+        #define NPY_OFF_T_PYFMT "h"
+    #elif NPY_SIZEOF_OFF_T == NPY_SIZEOF_INT
+        #define NPY_OFF_T_PYFMT "i"
+    #elif NPY_SIZEOF_OFF_T == NPY_SIZEOF_LONG
+        #define NPY_OFF_T_PYFMT "l"
+    #elif NPY_SIZEOF_OFF_T == NPY_SIZEOF_LONGLONG
+        #define NPY_OFF_T_PYFMT "L"
+    #else
+        #error Unsupported size for type off_t
+    #endif
 #endif
 
 /* enums for detected endianness */
diff --git a/numpy/core/setup.py b/numpy/core/setup.py
index 465752a40..2cb2f3a8b 100644
--- a/numpy/core/setup.py
+++ b/numpy/core/setup.py
@@ -280,6 +280,7 @@ def check_types(config_cmd, ext, build_dir):
     expected['Py_intptr_t'] = [4, 8]
     expected['PY_LONG_LONG'] = [8]
     expected['long long'] = [8]
+    expected['off_t'] = [4, 8]
 
     # Check we have the python header (-dev* packages on Linux)
     result = config_cmd.check_header('Python.h')
@@ -326,7 +327,7 @@ def check_types(config_cmd, ext, build_dir):
             raise SystemError("Checking sizeof (%s) failed !" % complex_def)
 
 
-    for type in ('Py_intptr_t',):
+    for type in ('Py_intptr_t', 'off_t'):
         res = config_cmd.check_type_size(type, headers=["Python.h"],
                 library_dirs=[pythonlib_dir()],
                 expected=expected[type])
diff --git a/numpy/core/src/multiarray/methods.c b/numpy/core/src/multiarray/methods.c
index 539059fec..37b7d3c18 100644
--- a/numpy/core/src/multiarray/methods.c
+++ b/numpy/core/src/multiarray/methods.c
@@ -566,6 +566,7 @@ array_tofile(PyArrayObject *self, PyObject *args, PyObject *kwds)
     FILE *fd;
     char *sep = "";
     char *format = "";
+    npy_off_t orig_pos;
     static char *kwlist[] = {"file", "sep", "format", NULL};
 
     if (!PyArg_ParseTupleAndKeywords(args, kwds, "O|ss", kwlist,
@@ -587,7 +588,7 @@ array_tofile(PyArrayObject *self, PyObject *args, PyObject *kwds)
         own = 0;
     }
 
-    fd = npy_PyFile_Dup(file, "wb");
+    fd = npy_PyFile_Dup(file, "wb", &orig_pos);
     if (fd == NULL) {
         PyErr_SetString(PyExc_IOError,
                 "first argument must be a string or open file");
@@ -596,7 +597,7 @@ array_tofile(PyArrayObject *self, PyObject *args, PyObject *kwds)
     if (PyArray_ToFile(self, fd, sep, format) < 0) {
         goto fail;
     }
-    if (npy_PyFile_DupClose(file, fd) < 0) {
+    if (npy_PyFile_DupClose(file, fd, orig_pos) < 0) {
         goto fail;
     }
     if (own && npy_PyFile_CloseFile(file) < 0) {
diff --git a/numpy/core/src/multiarray/multiarraymodule.c b/numpy/core/src/multiarray/multiarraymodule.c
index 240ee4233..df08495f8 100644
--- a/numpy/core/src/multiarray/multiarraymodule.c
+++ b/numpy/core/src/multiarray/multiarraymodule.c
@@ -1972,6 +1972,7 @@ array_fromfile(PyObject *NPY_UNUSED(ignored), PyObject *args, PyObject *keywds)
     static char *kwlist[] = {"file", "dtype", "count", "sep", NULL};
     PyArray_Descr *type = NULL;
     int own;
+    npy_off_t orig_pos;
     FILE *fp;
 
     if (!PyArg_ParseTupleAndKeywords(args, keywds,
@@ -1991,7 +1992,7 @@ array_fromfile(PyObject *NPY_UNUSED(ignored), PyObject *args, PyObject *keywds)
         Py_INCREF(file);
         own = 0;
     }
-    fp = npy_PyFile_Dup(file, "rb");
+    fp = npy_PyFile_Dup(file, "rb", &orig_pos);
     if (fp == NULL) {
         PyErr_SetString(PyExc_IOError,
                 "first argument must be an open file");
@@ -2003,7 +2004,7 @@ array_fromfile(PyObject *NPY_UNUSED(ignored), PyObject *args, PyObject *keywds)
     }
     ret = PyArray_FromFile(fp, type, (npy_intp) nin, sep);
 
-    if (npy_PyFile_DupClose(file, fp) < 0) {
+    if (npy_PyFile_DupClose(file, fp, orig_pos) < 0) {
         goto fail;
     }
     if (own && npy_PyFile_CloseFile(file) < 0) {
diff --git a/numpy/core/tests/test_multiarray.py b/numpy/core/tests/test_multiarray.py
index fa698f1ac..6336d41c7 100644
--- a/numpy/core/tests/test_multiarray.py
+++ b/numpy/core/tests/test_multiarray.py
@@ -5,6 +5,7 @@ import sys
 import os
 import warnings
 import operator
+import io
 if sys.version_info[0] >= 3:
     import builtins
 else:
@@ -2369,6 +2370,58 @@ class TestIO(object):
         y = np.fromstring(s, sep="@")
         assert_array_equal(x, y)
 
+    def test_file_position_after_fromfile(self):
+        # gh-4118
+        sizes = [io.DEFAULT_BUFFER_SIZE//8,
+                 io.DEFAULT_BUFFER_SIZE,
+                 io.DEFAULT_BUFFER_SIZE*8]
+
+        for size in sizes:
+            f = open(self.filename, 'wb')
+            f.seek(size-1)
+            f.write(b'\0')
+            f.close()
+
+            for mode in ['rb', 'r+b']:
+                err_msg = "%d %s" % (size, mode)
+
+                f = open(self.filename, mode)
+                f.read(2)
+                np.fromfile(f, dtype=np.float64, count=1)
+                pos = f.tell()
+                f.close()
+                assert_equal(pos, 10, err_msg=err_msg)
+
+        os.unlink(self.filename)
+
+    def test_file_position_after_tofile(self):
+        # gh-4118
+        sizes = [io.DEFAULT_BUFFER_SIZE//8,
+                 io.DEFAULT_BUFFER_SIZE,
+                 io.DEFAULT_BUFFER_SIZE*8]
+
+        for size in sizes:
+            err_msg = "%d" % (size,)
+
+            f = open(self.filename, 'wb')
+            f.seek(size-1)
+            f.write(b'\0')
+            f.seek(10)
+            f.write(b'12')
+            np.array([0], dtype=np.float64).tofile(f)
+            pos = f.tell()
+            f.close()
+            assert_equal(pos, 10 + 2 + 8, err_msg=err_msg)
+
+            f = open(self.filename, 'r+b')
+            f.read(2)
+            np.array([0], dtype=np.float64).tofile(f)
+            pos = f.tell()
+            f.close()
+            assert_equal(pos, 10, err_msg=err_msg)
+
+        os.unlink(self.filename)
+
     def _check_from(self, s, value, **kw):
         y = np.fromstring(asbytes(s), **kw)
         assert_array_equal(y, value)