1 files changed, 50 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)
+        )