ENH: Save to ZIP files without using temporary files.

Since Python 3.6 it is possible to write directly to a ZIP file,
without creating temporary files.

1 files changed, 33 insertions, 22 deletions

diff --git a/numpy/lib/npyio.py b/numpy/lib/npyio.py
index 7598b2c6b..96355ebc8 100644
--- a/numpy/lib/npyio.py
+++ b/numpy/lib/npyio.py
@@ -661,8 +661,6 @@ def _savez(file, args, kwds, compress, allow_pickle=True, pickle_kwargs=None):
     # Import is postponed to here since zipfile depends on gzip, an optional
     # component of the so-called standard library.
     import zipfile
-    # Import deferred for startup time improvement
-    import tempfile
 
     if isinstance(file, basestring):
         if not file.endswith('.npz'):
@@ -686,31 +684,44 @@ def _savez(file, args, kwds, compress, allow_pickle=True, pickle_kwargs=None):
 
     zipf = zipfile_factory(file, mode="w", compression=compression)
 
-    # Stage arrays in a temporary file on disk, before writing to zip.
-
-    # Since target file might be big enough to exceed capacity of a global
-    # temporary directory, create temp file side-by-side with the target file.
-    file_dir, file_prefix = os.path.split(file) if _is_string_like(file) else (None, 'tmp')
-    fd, tmpfile = tempfile.mkstemp(prefix=file_prefix, dir=file_dir, suffix='-numpy.npy')
-    os.close(fd)
-    try:
+    if sys.version_info >= (3, 6):
+        # Since Python 3.6 it is possible to write directly to a ZIP file.
         for key, val in namedict.items():
             fname = key + '.npy'
-            fid = open(tmpfile, 'wb')
-            try:
-                format.write_array(fid, np.asanyarray(val),
+            val = np.asanyarray(val)
+            force_zip64 = val.nbytes >= 2**30
+            with zipf.open(fname, 'w', force_zip64=force_zip64) as fid:
+                format.write_array(fid, val,
                                    allow_pickle=allow_pickle,
                                    pickle_kwargs=pickle_kwargs)
-                fid.close()
-                fid = None
-                zipf.write(tmpfile, arcname=fname)
-            except IOError as exc:
-                raise IOError("Failed to write to %s: %s" % (tmpfile, exc))
-            finally:
-                if fid:
+    else:
+        # Stage arrays in a temporary file on disk, before writing to zip.
+
+        # Import deferred for startup time improvement
+        import tempfile
+        # Since target file might be big enough to exceed capacity of a global
+        # temporary directory, create temp file side-by-side with the target file.
+        file_dir, file_prefix = os.path.split(file) if _is_string_like(file) else (None, 'tmp')
+        fd, tmpfile = tempfile.mkstemp(prefix=file_prefix, dir=file_dir, suffix='-numpy.npy')
+        os.close(fd)
+        try:
+            for key, val in namedict.items():
+                fname = key + '.npy'
+                fid = open(tmpfile, 'wb')
+                try:
+                    format.write_array(fid, np.asanyarray(val),
+                                       allow_pickle=allow_pickle,
+                                       pickle_kwargs=pickle_kwargs)
                     fid.close()
-    finally:
-        os.remove(tmpfile)
+                    fid = None
+                    zipf.write(tmpfile, arcname=fname)
+                except IOError as exc:
+                    raise IOError("Failed to write to %s: %s" % (tmpfile, exc))
+                finally:
+                    if fid:
+                        fid.close()
+        finally:
+            os.remove(tmpfile)
 
     zipf.close()