[pypy-commit] pypy default: test/fix searchsorted return type for scalars

[bdkearns]

Author: Brian Kearns <bdkearns at gmail.com>
Branch: 
Changeset: r71279:3581f7a906c9
Date: 2014-05-05 15:02 -0400
http://bitbucket.org/pypy/pypy/changeset/3581f7a906c9/

Log:	test/fix searchsorted return type for scalars

diff --git a/pypy/module/micronumpy/ndarray.py b/pypy/module/micronumpy/ndarray.py
--- a/pypy/module/micronumpy/ndarray.py
+++ b/pypy/module/micronumpy/ndarray.py
@@ -738,6 +738,8 @@
         ret = W_NDimArray.from_shape(
             space, v.get_shape(), descriptor.get_dtype_cache(space).w_longdtype)
         app_searchsort(space, self, v, space.wrap(side), ret)
+        if ret.is_scalar():
+            return ret.get_scalar_value()
         return ret
 
     def descr_setasflat(self, space, w_v):
diff --git a/pypy/module/micronumpy/test/test_sorting.py b/pypy/module/micronumpy/test/test_sorting.py
--- a/pypy/module/micronumpy/test/test_sorting.py
+++ b/pypy/module/micronumpy/test/test_sorting.py
@@ -351,13 +351,21 @@
             assert (x.argsort(kind='m') == np.arange(32)).all()
 
     def test_searchsort(self):
-        from numpy import arange
+        import numpy as np
         import sys
-        a = arange(1, 6)
+        a = np.arange(1, 6)
         ret = a.searchsorted(3)
         assert ret == 2
+        assert isinstance(ret, np.generic)
+        ret = a.searchsorted(np.array(3))
+        assert ret == 2
+        assert isinstance(ret, np.generic)
+        ret = a.searchsorted(np.array([3]))
+        assert ret == 2
+        assert isinstance(ret, np.ndarray)
         ret = a.searchsorted(3, side='right')
         assert ret == 3
+        assert isinstance(ret, np.generic)
         ret = a.searchsorted([-10, 10, 2, 3])
         assert (ret == [0, 5, 1, 2]).all()
         if '__pypy__' in sys.builtin_module_names: