BUG: take casting logic with an out= argument

royJackman · royJackman · commit 024b453c46d5 · 2020-05-26T00:14:17.000-04:00
Closes numpy#16319
diff --git a/numpy/ma/core.py b/numpy/ma/core.py
@@ -5972,7 +5972,17 @@ def take(self, indices, axis=None, out=None, mode='raise'):
         if out is None:
             out = _data.take(indices, axis=axis, mode=mode)[...].view(cls)
         else:
-            np.take(_data, indices, axis=axis, mode=mode, out=out)
+            # Check if the numeric input values are within the range of the
+            # output dtype, if so, convert type and output, else raise error
+            in_dtype = self.dtype
+            out_dtype = out.dtype
+            if ntypes.issubdtype(in_dtype, np.number) and ntypes.issubdtype(out_dtype, np.number):
+                if np.logical_and(_data >= np.iinfo(out_dtype).min, _data <= np.iinfo(out_dtype).max).all():
+                    np.take(_data.astype(out_dtype), indices, axis=axis, mode=mode, out=out)
+                else:
+                    raise TypeError('Output format does not cover input range')
+            else:
+                np.take(_data, indices, axis=axis, mode=mode, out=out)
         # Get the mask
         if isinstance(out, MaskedArray):
             if _mask is nomask:
diff --git a/numpy/ma/tests/test_core.py b/numpy/ma/tests/test_core.py
@@ -3485,6 +3485,27 @@ def test_take(self):
         assert_equal(take(x, [0, 2], axis=1),
                      array([[10, 30], [40, 60]], mask=[[0, 1], [1, 0]]))
 
+    def test_take_dtype_conversion(self):
+        # Take to smaller dtype within range
+        a = arange(3, dtype=np.int16)
+        b = empty((3,), dtype=np.int8)
+        a.take([0,1,2], out=b)
+        assert_equal(a,b)
+        # Smaller dtype out of range
+        a = a + 128
+        with assert_raises(TypeError):
+            a.take([0,1,2], out=b)
+        # Larger dtype
+        b = empty((3,), dtype=np.int32)
+        a.take([0,1,2], out=b)
+        assert_equal(a, b)
+        
+        # Take to smaller dtype within range float to int
+        a = np.random.rand(3).astype(np.float32)
+        b = empty((3,), dtype=np.int8)
+        a.take([0,1,2], out=b)
+        assert_equal(a.astype(np.int8), b)
+
     def test_take_masked_indices(self):
         # Test take w/ masked indices
         a = np.array((40, 18, 37, 9, 22))
