[pypy-commit] pypy default: added rshift and rrshift to ndarray

[alex_gaynor]

Author: Alex Gaynor <alex.gaynor at gmail.com>
Branch: 
Changeset: r52293:ad83678913a0
Date: 2012-02-09 09:12 -0500
http://bitbucket.org/pypy/pypy/changeset/ad83678913a0/

Log:	added rshift and rrshift to ndarray

diff --git a/pypy/module/micronumpy/interp_numarray.py b/pypy/module/micronumpy/interp_numarray.py
--- a/pypy/module/micronumpy/interp_numarray.py
+++ b/pypy/module/micronumpy/interp_numarray.py
@@ -104,6 +104,7 @@
     descr_mod = _binop_impl("mod")
     descr_pow = _binop_impl("power")
     descr_lshift = _binop_impl("left_shift")
+    descr_rshift = _binop_impl("right_shift")
 
     descr_eq = _binop_impl("equal")
     descr_ne = _binop_impl("not_equal")
@@ -136,6 +137,7 @@
     descr_rmod = _binop_right_impl("mod")
     descr_rpow = _binop_right_impl("power")
     descr_rlshift = _binop_right_impl("left_shift")
+    descr_rrshift = _binop_right_impl("right_shift")
 
     descr_rand = _binop_right_impl("bitwise_and")
     descr_ror = _binop_right_impl("bitwise_or")
@@ -1253,6 +1255,7 @@
     __divmod__ = interp2app(BaseArray.descr_divmod),
     __pow__ = interp2app(BaseArray.descr_pow),
     __lshift__ = interp2app(BaseArray.descr_lshift),
+    __rshift__ = interp2app(BaseArray.descr_rshift),
 
     __and__ = interp2app(BaseArray.descr_and),
     __or__ = interp2app(BaseArray.descr_or),
@@ -1265,6 +1268,7 @@
     __rdivmod__ = interp2app(BaseArray.descr_rdivmod),
     __rpow__ = interp2app(BaseArray.descr_rpow),
     __rlshift__ = interp2app(BaseArray.descr_rlshift),
+    __rrshift__ = interp2app(BaseArray.descr_rrshift),
 
     __rand__ = interp2app(BaseArray.descr_rand),
     __ror__ = interp2app(BaseArray.descr_ror),
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
@@ -392,8 +392,8 @@
             ("true_divide", "div", 2, {"promote_to_float": True}),
             ("mod", "mod", 2, {"promote_bools": True}),
             ("power", "pow", 2, {"promote_bools": True}),
-
             ("left_shift", "lshift", 2, {"int_only": True}),
+            ("right_shift", "rshift", 2, {"int_only": True}),
 
             ("equal", "eq", 2, {"comparison_func": True}),
             ("not_equal", "ne", 2, {"comparison_func": True}),
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
@@ -655,6 +655,22 @@
         a = arange(3)
         assert (2 << a == [2, 4, 8]).all()
 
+    def test_rshift(self):
+        from _numpypy import arange, array
+
+        a = arange(10)
+        assert (a >> 2 == [0, 0, 0, 0, 1, 1, 1, 1, 2, 2]).all()
+        a = array([True, False])
+        assert (a >> 1 == [0, 0]).all()
+        a = arange(3, dtype=float)
+        raises(TypeError, lambda: a >> 1)
+
+    def test_rrshift(self):
+        from _numpypy import arange
+
+        a = arange(5)
+        assert (2 >> a == [2, 1, 0, 0, 0]).all()
+
     def test_pow(self):
         from _numpypy import array
         a = array(range(5), float)
diff --git a/pypy/module/micronumpy/types.py b/pypy/module/micronumpy/types.py
--- a/pypy/module/micronumpy/types.py
+++ b/pypy/module/micronumpy/types.py
@@ -299,6 +299,10 @@
     def lshift(self, v1, v2):
         return v1 << v2
 
+    @simple_binary_op
+    def rshift(self, v1, v2):
+        return v1 >> v2
+
     @simple_unary_op
     def sign(self, v):
         if v > 0: