Merge pull request #8388 from gfyoung/real-imag-scalar

API: Return scalars for scalar inputs to np.real/imag

1 files changed, 40 insertions, 0 deletions

diff --git a/numpy/lib/tests/test_type_check.py b/numpy/lib/tests/test_type_check.py
index 473b558be..383ffa55c 100644
--- a/numpy/lib/tests/test_type_check.py
+++ b/numpy/lib/tests/test_type_check.py
@@ -98,10 +98,30 @@ class TestReal(TestCase):
         y = np.random.rand(10,)
         assert_array_equal(y, np.real(y))
 
+        y = np.array(1)
+        out = np.real(y)
+        assert_array_equal(y, out)
+        assert_(isinstance(out, np.ndarray))
+
+        y = 1
+        out = np.real(y)
+        assert_equal(y, out)
+        assert_(not isinstance(out, np.ndarray))
+
     def test_cmplx(self):
         y = np.random.rand(10,)+1j*np.random.rand(10,)
         assert_array_equal(y.real, np.real(y))
 
+        y = np.array(1 + 1j)
+        out = np.real(y)
+        assert_array_equal(y.real, out)
+        assert_(isinstance(out, np.ndarray))
+
+        y = 1 + 1j
+        out = np.real(y)
+        assert_equal(1.0, out)
+        assert_(not isinstance(out, np.ndarray))
+
 
 class TestImag(TestCase):
 
@@ -109,10 +129,30 @@ class TestImag(TestCase):
         y = np.random.rand(10,)
         assert_array_equal(0, np.imag(y))
 
+        y = np.array(1)
+        out = np.imag(y)
+        assert_array_equal(0, out)
+        assert_(isinstance(out, np.ndarray))
+
+        y = 1
+        out = np.imag(y)
+        assert_equal(0, out)
+        assert_(not isinstance(out, np.ndarray))
+
     def test_cmplx(self):
         y = np.random.rand(10,)+1j*np.random.rand(10,)
         assert_array_equal(y.imag, np.imag(y))
 
+        y = np.array(1 + 1j)
+        out = np.imag(y)
+        assert_array_equal(y.imag, out)
+        assert_(isinstance(out, np.ndarray))
+
+        y = 1 + 1j
+        out = np.imag(y)
+        assert_equal(1.0, out)
+        assert_(not isinstance(out, np.ndarray))
+
 
 class TestIscomplex(TestCase):