BUG: Cast size to int64 when loading from archive

drasmuss · drasmuss · commit da668fc74653 · 2016-05-09T09:45:42.000-04:00
Prevents overflow errors for large arrays on systems
where the default int type is int32.
diff --git a/numpy/lib/format.py b/numpy/lib/format.py
@@ -623,7 +623,7 @@ def read_array(fp, allow_pickle=True, pickle_kwargs=None):
     if len(shape) == 0:
         count = 1
     else:
-        count = numpy.multiply.reduce(shape)
+        count = numpy.multiply.reduce(shape, dtype=numpy.int64)
 
     # Now read the actual data.
     if dtype.hasobject:
diff --git a/numpy/lib/tests/test_format.py b/numpy/lib/tests/test_format.py
@@ -836,5 +836,19 @@ def test_large_file_support():
     assert_array_equal(r, d)
 
 
+@dec.slow
+def test_large_archive():
+    a = np.empty((2 ** 30, 2), dtype=np.uint8)
+    fname = os.path.join(tempdir, "large_archive")
+
+    with open(fname, "wb") as f:
+        np.savez(f, arr=a)
+
+    with open(fname, "rb") as f:
+        new_a = np.load(f)["arr"]
+
+    assert a.shape == new_a.shape
+
+
 if __name__ == "__main__":
     run_module_suite()
