[pypy-commit] pypy default: fix int_only numpy ufunc2 wrt uint64 (issue 1664)

[bdkearns]

Author: Brian Kearns <bdkearns at gmail.com>
Branch: 
Changeset: r68533:fc6ac3a6a746
Date: 2013-12-24 13:08 -0500
http://bitbucket.org/pypy/pypy/changeset/fc6ac3a6a746/

Log:	fix int_only numpy ufunc2 wrt uint64 (issue 1664)

diff --git a/pypy/module/micronumpy/interp_ufuncs.py b/pypy/module/micronumpy/interp_ufuncs.py
--- a/pypy/module/micronumpy/interp_ufuncs.py
+++ b/pypy/module/micronumpy/interp_ufuncs.py
@@ -376,14 +376,19 @@
                 w_rdtype = w_ldtype
             elif w_lhs.is_scalar() and not w_rhs.is_scalar():
                 w_ldtype = w_rdtype
-        if (self.int_only and (not w_ldtype.is_int_type() or not w_rdtype.is_int_type()) or
-                not self.allow_bool and (w_ldtype.is_bool_type() or w_rdtype.is_bool_type()) or
-                not self.allow_complex and (w_ldtype.is_complex_type() or w_rdtype.is_complex_type())):
-            raise OperationError(space.w_TypeError, space.wrap("Unsupported types"))
         calc_dtype = find_binop_result_dtype(space,
             w_ldtype, w_rdtype,
             promote_to_float=self.promote_to_float,
             promote_bools=self.promote_bools)
+        if (self.int_only and (not w_ldtype.is_int_type() or
+                               not w_rdtype.is_int_type() or
+                               not calc_dtype.is_int_type()) or
+                not self.allow_bool and (w_ldtype.is_bool_type() or
+                                         w_rdtype.is_bool_type()) or
+                not self.allow_complex and (w_ldtype.is_complex_type() or
+                                            w_rdtype.is_complex_type())):
+            raise OperationError(space.w_TypeError, space.wrap(
+                "ufunc '%s' not supported for the input types" % self.name))
         if space.is_none(w_out):
             out = None
         elif not isinstance(w_out, W_NDimArray):
diff --git a/pypy/module/micronumpy/test/test_numarray.py b/pypy/module/micronumpy/test/test_numarray.py
--- a/pypy/module/micronumpy/test/test_numarray.py
+++ b/pypy/module/micronumpy/test/test_numarray.py
@@ -1164,14 +1164,20 @@
         assert (2 << a == [2, 4, 8]).all()
 
     def test_rshift(self):
-        from numpypy import arange, array
-
-        a = arange(10)
+        import numpy as np
+        a = np.arange(10)
         assert (a >> 2 == [0, 0, 0, 0, 1, 1, 1, 1, 2, 2]).all()
-        a = array([True, False])
+        a = np.array([True, False])
         assert (a >> 1 == [0, 0]).all()
-        a = arange(3, dtype=float)
+        a = np.arange(3, dtype=float)
         raises(TypeError, lambda: a >> 1)
+        a = np.array([123], dtype='uint64')
+        b = a >> 1
+        assert b == 61
+        assert b.dtype.type is np.uint64
+        a = np.array(123, dtype='uint64')
+        exc = raises(TypeError, "a >> 1")
+        assert 'not supported for the input types' in exc.value.message
 
     def test_rrshift(self):
         from numpypy import arange
diff --git a/pypy/module/micronumpy/test/test_scalar.py b/pypy/module/micronumpy/test/test_scalar.py
--- a/pypy/module/micronumpy/test/test_scalar.py
+++ b/pypy/module/micronumpy/test/test_scalar.py
@@ -113,6 +113,15 @@
         assert a.squeeze() is a
         raises(TypeError, a.squeeze, 2)
 
+    def test_bitshift(self):
+        import numpy as np
+        assert np.int32(123) >> 1 == 61
+        assert type(np.int32(123) >> 1) is np.int64
+        assert np.int64(123) << 1 == 246
+        assert type(np.int64(123) << 1) is np.int64
+        exc = raises(TypeError, "np.uint64(123) >> 1")
+        assert 'not supported for the input types' in exc.value.message
+
     def test_attributes(self):
         import numpy as np
         value = np.dtype('int64').type(12345)