Merge pull request #42 from tomwhite/overflow-bug

Fix overflow bug exposed by Hypothesis failures

rechunker/algorithm.py

@@ -1,13 +1,7 @@
 """Core rechunking algorithm stuff."""
 from typing import Sequence, Optional, List, Tuple
 
-try:
-    from math import prod
-
-    numpy_prod = False
-except ImportError:
-    numpy_prod = True
-    from numpy import prod
+from rechunker.compat import prod
 
 
 def consolidate_chunks(
@@ -119,10 +113,6 @@ def rechunking_plan(
 
     source_chunk_mem = itemsize * prod(source_chunks)
     target_chunk_mem = itemsize * prod(target_chunks)
-    if numpy_prod:
-        # Convert to Python type for JSON serialization
-        source_chunk_mem = source_chunk_mem.item()
-        target_chunk_mem = target_chunk_mem.item()
 
     if source_chunk_mem > max_mem:
         raise ValueError(

rechunker/compat.py

@@ -0,0 +1,13 @@
+from functools import reduce
+import operator
+from typing import Sequence
+
+
+def prod(iterable: Sequence[int]) -> int:
+    """Implementation of `math.prod()` all Python versions."""
+    try:
+        from math import prod as mathprod  # Python 3.8
+
+        return mathprod(iterable)
+    except ImportError:
+        return reduce(operator.mul, iterable, 1)

tests/test_algorithm.py

@@ -1,11 +1,7 @@
 #!/usr/bin/env python
 
 """Tests for `rechunker` package."""
-try:
-    from math import prod
-except ImportError:
-    from numpy import prod
-
+from rechunker.compat import prod
 
 import pytest
 from hypothesis import given, assume
@@ -215,8 +211,6 @@ def shapes_chunks_maxmem_for_ndim(draw):
         shapes_chunks_maxmem(ndim=ndim, itemsize=4, max_len=10_000)
     )
     max_mem = min_mem * 10
-    # needed to handle overflows
-    assume(max_mem > 1)
     return shape, source_chunks, target_chunks, max_mem, itemsize
 
 

tests/test_compat.py

@@ -0,0 +1,10 @@
+import numpy as np
+from rechunker.compat import prod
+
+
+def test_prod():
+    assert prod(()) == 1
+    assert prod((2,)) == 2
+    assert prod((2, 3)) == 6
+    n = np.iinfo(np.int64).max
+    assert prod((n, 2)) == n * 2