ENH: np.angle: Preserve subclasses

eric-wieser · eric-wieser · commit aea9eb208987 · 2018-07-30T23:02:10.000-07:00
diff --git a/numpy/lib/function_base.py b/numpy/lib/function_base.py
@@ -1339,7 +1339,7 @@ def angle(z, deg=False):
     45.0
 
     """
-    z = asarray(z)
+    z = asanyarray(z)
     if issubclass(z.dtype.type, _nx.complexfloating):
         zimag = z.imag
         zreal = z.real
diff --git a/numpy/lib/tests/test_function_base.py b/numpy/lib/tests/test_function_base.py
@@ -1043,6 +1043,16 @@ def test_basic(self):
         assert_array_almost_equal(y, yo, 11)
         assert_array_almost_equal(z, zo, 11)
 
+    def test_subclass(self):
+        x = np.ma.array([1 + 3j, 1, np.sqrt(2)/2 * (1 + 1j)])
+        x[1] = np.ma.masked
+        expected = np.ma.array([np.arctan(3.0 / 1.0), 0, np.arctan(1.0)])
+        expected[1] = np.ma.masked
+        actual = angle(x)
+        assert_equal(type(actual), type(expected))
+        assert_equal(actual.mask, expected.mask)
+        assert_equal(actual, expected)
+
 
 class TestTrimZeros(object):
 
