Merge pull request #14459 from eric-wieser/MemoryError-size

ENH: Print the amount of memory that would be used by a failed allocation

2 files changed, 92 insertions, 2 deletions

diff --git a/numpy/core/_exceptions.py b/numpy/core/_exceptions.py
index b3805af04..88a45561f 100644
--- a/numpy/core/_exceptions.py
+++ b/numpy/core/_exceptions.py
@@ -147,6 +147,54 @@ class _ArrayMemoryError(MemoryError):
         self.shape = shape
         self.dtype = dtype
 
-    def __str__(self):
-        return "Unable to allocate array with shape {} and data type {}".format(self.shape, self.dtype)
+    @property
+    def _total_size(self):
+        num_bytes = self.dtype.itemsize
+        for dim in self.shape:
+            num_bytes *= dim
+        return num_bytes
+
+    @staticmethod
+    def _size_to_string(num_bytes):
+        """ Convert a number of bytes into a binary size string """
+        import math
+
+        # https://en.wikipedia.org/wiki/Binary_prefix
+        LOG2_STEP = 10
+        STEP = 1024
+        units = ['bytes', 'KiB', 'MiB', 'GiB', 'TiB', 'PiB', 'EiB']
+
+        unit_i = max(num_bytes.bit_length() - 1, 1) // LOG2_STEP
+        unit_val = 1 << (unit_i * LOG2_STEP)
+        n_units = num_bytes / unit_val
+        del unit_val
+
+        # ensure we pick a unit that is correct after rounding
+        if round(n_units) == STEP:
+            unit_i += 1
+            n_units /= STEP
+
+        # deal with sizes so large that we don't have units for them
+        if unit_i >= len(units):
+            new_unit_i = len(units) - 1
+            n_units *= 1 << ((unit_i - new_unit_i) * LOG2_STEP)
+            unit_i = new_unit_i
+
+        unit_name = units[unit_i]
+        # format with a sensible number of digits
+        if unit_i == 0:
+            # no decimal point on bytes
+            return '{:.0f} {}'.format(n_units, unit_name)
+        elif round(n_units) < 1000:
+            # 3 significant figures, if none are dropped to the left of the .
+            return '{:#.3g} {}'.format(n_units, unit_name)
+        else:
+            # just give all the digits otherwise
+            return '{:#.0f} {}'.format(n_units, unit_name)
 
+    def __str__(self):
+        size_str = self._size_to_string(self._total_size)
+        return (
+            "Unable to allocate {} for an array with shape {} and data type {}"
+            .format(size_str, self.shape, self.dtype)
+        )
diff --git a/numpy/core/tests/test__exceptions.py b/numpy/core/tests/test__exceptions.py
new file mode 100644
index 000000000..494b51f34
--- /dev/null
+++ b/numpy/core/tests/test__exceptions.py
@@ -0,0 +1,42 @@
+"""
+Tests of the ._exceptions module. Primarily for exercising the __str__ methods.
+"""
+import numpy as np
+
+_ArrayMemoryError = np.core._exceptions._ArrayMemoryError
+
+class TestArrayMemoryError:
+    def test_str(self):
+        e = _ArrayMemoryError((1023,), np.dtype(np.uint8))
+        str(e)  # not crashing is enough
+
+    # testing these properties is easier than testing the full string repr
+    def test__size_to_string(self):
+        """ Test e._size_to_string """
+        f = _ArrayMemoryError._size_to_string
+        Ki = 1024
+        assert f(0) == '0 bytes'
+        assert f(1) == '1 bytes'
+        assert f(1023) == '1023 bytes'
+        assert f(Ki) == '1.00 KiB'
+        assert f(Ki+1) == '1.00 KiB'
+        assert f(10*Ki) == '10.0 KiB'
+        assert f(int(999.4*Ki)) == '999. KiB'
+        assert f(int(1023.4*Ki)) == '1023. KiB'
+        assert f(int(1023.5*Ki)) == '1.00 MiB'
+        assert f(Ki*Ki) == '1.00 MiB'
+
+        # 1023.9999 Mib should round to 1 GiB
+        assert f(int(Ki*Ki*Ki*0.9999)) == '1.00 GiB'
+        assert f(Ki*Ki*Ki*Ki*Ki*Ki) == '1.00 EiB'
+        # larger than sys.maxsize, adding larger prefices isn't going to help
+        # anyway.
+        assert f(Ki*Ki*Ki*Ki*Ki*Ki*123456) == '123456. EiB'
+
+    def test__total_size(self):
+        """ Test e._total_size """
+        e = _ArrayMemoryError((1,), np.dtype(np.uint8))
+        assert e._total_size == 1
+
+        e = _ArrayMemoryError((2, 4), np.dtype((np.uint64, 16)))
+        assert e._total_size == 1024