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

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

2 files changed, 66 insertions, 12 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):
 
diff --git a/numpy/lib/type_check.py b/numpy/lib/type_check.py
index a59fe3cc4..5202cebde 100644
--- a/numpy/lib/type_check.py
+++ b/numpy/lib/type_check.py
@@ -9,8 +9,7 @@ __all__ = ['iscomplexobj', 'isrealobj', 'imag', 'iscomplex',
            'common_type']
 
 import numpy.core.numeric as _nx
-from numpy.core.numeric import asarray, asanyarray, array, isnan, \
-                obj2sctype, zeros
+from numpy.core.numeric import asarray, asanyarray, array, isnan, zeros
 from .ufunclike import isneginf, isposinf
 
 _typecodes_by_elsize = 'GDFgdfQqLlIiHhBb?'
@@ -104,9 +103,10 @@ def asfarray(a, dtype=_nx.float_):
         dtype = _nx.float_
     return asarray(a, dtype=dtype)
 
+
 def real(val):
     """
-    Return the real part of the elements of the array.
+    Return the real part of the complex argument.
 
     Parameters
     ----------
@@ -115,9 +115,10 @@ def real(val):
 
     Returns
     -------
-    out : ndarray
-        Output array. If `val` is real, the type of `val` is used for the
-        output.  If `val` has complex elements, the returned type is float.
+    out : ndarray or scalar
+        The real component of the complex argument. If `val` is real, the type
+        of `val` is used for the output.  If `val` has complex elements, the
+        returned type is float.
 
     See Also
     --------
@@ -134,13 +135,19 @@ def real(val):
     >>> a.real = np.array([9, 8, 7])
     >>> a
     array([ 9.+2.j,  8.+4.j,  7.+6.j])
+    >>> np.real(1 + 1j)
+    1.0
 
     """
-    return asanyarray(val).real
+    try:
+        return val.real
+    except AttributeError:
+        return asanyarray(val).real
+
 
 def imag(val):
     """
-    Return the imaginary part of the elements of the array.
+    Return the imaginary part of the complex argument.
 
     Parameters
     ----------
@@ -149,9 +156,10 @@ def imag(val):
 
     Returns
     -------
-    out : ndarray
-        Output array. If `val` is real, the type of `val` is used for the
-        output.  If `val` has complex elements, the returned type is float.
+    out : ndarray or scalar
+        The imaginary component of the complex argument. If `val` is real,
+        the type of `val` is used for the output.  If `val` has complex
+        elements, the returned type is float.
 
     See Also
     --------
@@ -165,9 +173,15 @@ def imag(val):
     >>> a.imag = np.array([8, 10, 12])
     >>> a
     array([ 1. +8.j,  3.+10.j,  5.+12.j])
+    >>> np.imag(1 + 1j)
+    1.0
 
     """
-    return asanyarray(val).imag
+    try:
+        return val.imag
+    except AttributeError:
+        return asanyarray(val).imag
+
 
 def iscomplex(x):
     """