[pypy-commit] pypy default: added left shift to numpy

[alex_gaynor]

Author: Alex Gaynor <alex.gaynor at gmail.com>
Branch: 
Changeset: r52287:d89a66e9dc00
Date: 2012-02-09 08:48 -0500
http://bitbucket.org/pypy/pypy/changeset/d89a66e9dc00/

Log:	added left shift to numpy

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
@@ -103,6 +103,7 @@
     descr_div = _binop_impl("divide")
     descr_pow = _binop_impl("power")
     descr_mod = _binop_impl("mod")
+    descr_lshift = _binop_impl("left_shift")
 
     descr_eq = _binop_impl("equal")
     descr_ne = _binop_impl("not_equal")
@@ -1235,6 +1236,7 @@
     __div__ = interp2app(BaseArray.descr_div),
     __pow__ = interp2app(BaseArray.descr_pow),
     __mod__ = interp2app(BaseArray.descr_mod),
+    __lshift__ = interp2app(BaseArray.descr_lshift),
 
     __radd__ = interp2app(BaseArray.descr_radd),
     __rsub__ = interp2app(BaseArray.descr_rsub),
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
@@ -393,6 +393,8 @@
             ("mod", "mod", 2, {"promote_bools": True}),
             ("power", "pow", 2, {"promote_bools": True}),
 
+            ("left_shift", "lshift", 2, {"int_only": True}),
+
             ("equal", "eq", 2, {"comparison_func": True}),
             ("not_equal", "ne", 2, {"comparison_func": True}),
             ("less", "lt", 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
@@ -625,6 +625,16 @@
         for i in range(5):
             assert b[i] == i / 5.0
 
+    def test_lshift(self):
+        from _numpypy import array
+
+        a = array([0, 1, 2, 3])
+        assert (a << 2 == [0, 4, 8, 12]).all()
+        a = array([True, False])
+        assert (a << 2 == [4, 0]).all()
+        a = array([1.0])
+        raises(TypeError, lambda: a << 2)
+
     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
@@ -295,6 +295,10 @@
             v1 *= v1
         return res
 
+    @simple_binary_op
+    def lshift(self, v1, v2):
+        return v1 << v2
+
     @simple_unary_op
     def sign(self, v):
         if v > 0: