BUG: ensure np.median does not drop subclass for NaN result.

Currently, np.median is almost completely safe for subclasses, except
if the result is NaN.  In that case, it assumes the result is a scalar
and substitutes a NaN with the right dtype.  This PR fixes that, since
subclasses like astropy's Quantity generally use array scalars to
preserve subclass information such as the unit.

Author: mhvk

numpy/lib/tests/test_function_base.py

@@ -3432,6 +3432,16 @@ def mean(self, axis=None, dtype=None, out=None):
         a = MySubClass([1, 2, 3])
         assert_equal(np.median(a), -7)
 
+    @pytest.mark.parametrize('arr',
+                             ([1., 2., 3.], [1., np.nan, 3.], np.nan, 0.))
+    def test_subclass2(self, arr):
+        """Check that we return subclasses, even if a NaN scalar."""
+        class MySubclass(np.ndarray):
+            pass
+
+        m = np.median(np.array(arr).view(MySubclass))
+        assert isinstance(m, MySubclass)
+
     def test_out(self):
         o = np.zeros((4,))
         d = np.ones((3, 4))

numpy/lib/utils.py

@@ -1029,14 +1029,12 @@ def _median_nancheck(data, result, axis, out):
     # masked NaN values are ok
     if np.ma.isMaskedArray(n):
         n = n.filled(False)
-    if result.ndim == 0:
-        if n == True:
-            if out is not None:
-                out[...] = data.dtype.type(np.nan)
-                result = out
-            else:
-                result = data.dtype.type(np.nan)
-    elif np.count_nonzero(n.ravel()) > 0:
+    if np.count_nonzero(n.ravel()) > 0:
+        # Without given output, it is possible that the current result is a
+        # numpy scalar, which is not writeable.  If so, just return nan.
+        if isinstance(result, np.generic):
+            return data.dtype.type(np.nan)
+
         result[n] = np.nan
     return result