[pypy-commit] pypy default: merge numpy-indexing-by-arrays-bool, this adds some basic indexing by bool

fijal noreply at buildbot.pypy.org
Fri Jan 20 15:21:49 CET 2012


Author: Maciej Fijalkowski <fijall at gmail.com>
Branch: 
Changeset: r51524:ab36137540d5
Date: 2012-01-20 16:21 +0200
http://bitbucket.org/pypy/pypy/changeset/ab36137540d5/

Log:	merge numpy-indexing-by-arrays-bool, this adds some basic indexing
	by bool arrays (only when shapes match) as well as bitwise_and/or
	and a few helper functions. More work on indexing by arrays to be
	done in the future

diff --git a/pypy/module/micronumpy/__init__.py b/pypy/module/micronumpy/__init__.py
--- a/pypy/module/micronumpy/__init__.py
+++ b/pypy/module/micronumpy/__init__.py
@@ -86,6 +86,8 @@
         ("subtract", "subtract"),
         ('sqrt', 'sqrt'),
         ("tan", "tan"),
+        ('bitwise_and', 'bitwise_and'),
+        ('bitwise_or', 'bitwise_or'),
     ]:
         interpleveldefs[exposed] = "interp_ufuncs.get(space).%s" % impl
 
diff --git a/pypy/module/micronumpy/interp_dtype.py b/pypy/module/micronumpy/interp_dtype.py
--- a/pypy/module/micronumpy/interp_dtype.py
+++ b/pypy/module/micronumpy/interp_dtype.py
@@ -47,6 +47,10 @@
     def getitem(self, storage, i):
         return self.itemtype.read(storage, self.itemtype.get_element_size(), i, 0)
 
+    def getitem_bool(self, storage, i):
+        isize = self.itemtype.get_element_size()
+        return self.itemtype.read_bool(storage, isize, i, 0)
+
     def setitem(self, storage, i, box):
         self.itemtype.store(storage, self.itemtype.get_element_size(), i, 0, box)
 
@@ -85,6 +89,12 @@
     def descr_get_shape(self, space):
         return space.newtuple([])
 
+    def is_int_type(self):
+        return self.kind == SIGNEDLTR or self.kind == UNSIGNEDLTR
+
+    def is_bool_type(self):
+        return self.kind == BOOLLTR
+
 W_Dtype.typedef = TypeDef("dtype",
     __module__ = "numpypy",
     __new__ = interp2app(W_Dtype.descr__new__.im_func),
diff --git a/pypy/module/micronumpy/interp_iter.py b/pypy/module/micronumpy/interp_iter.py
--- a/pypy/module/micronumpy/interp_iter.py
+++ b/pypy/module/micronumpy/interp_iter.py
@@ -4,6 +4,19 @@
 from pypy.module.micronumpy.strides import calculate_broadcast_strides,\
      calculate_slice_strides
 
+# structures to describe slicing
+
+class Chunk(object):
+    def __init__(self, start, stop, step, lgt):
+        self.start = start
+        self.stop = stop
+        self.step = step
+        self.lgt = lgt
+
+    def extend_shape(self, shape):
+        if self.step != 0:
+            shape.append(self.lgt)
+
 class BaseTransform(object):
     pass
 
@@ -38,11 +51,18 @@
         self.size = size
 
     def next(self, shapelen):
+        return self._next(1)
+
+    def _next(self, ofs):
         arr = instantiate(ArrayIterator)
         arr.size = self.size
-        arr.offset = self.offset + 1
+        arr.offset = self.offset + ofs
         return arr
 
+    def next_no_increase(self, shapelen):
+        # a hack to make JIT believe this is always virtual
+        return self._next(0)
+
     def done(self):
         return self.offset >= self.size
 
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
@@ -2,14 +2,15 @@
 from pypy.interpreter.error import OperationError, operationerrfmt
 from pypy.interpreter.gateway import interp2app, NoneNotWrapped
 from pypy.interpreter.typedef import TypeDef, GetSetProperty
-from pypy.module.micronumpy import interp_ufuncs, interp_dtype, signature
+from pypy.module.micronumpy import interp_ufuncs, interp_dtype, signature,\
+     interp_boxes
 from pypy.module.micronumpy.strides import calculate_slice_strides
 from pypy.rlib import jit
 from pypy.rpython.lltypesystem import lltype, rffi
 from pypy.tool.sourcetools import func_with_new_name
 from pypy.rlib.rstring import StringBuilder
 from pypy.module.micronumpy.interp_iter import ArrayIterator, OneDimIterator,\
-     SkipLastAxisIterator
+     SkipLastAxisIterator, Chunk, ViewIterator
 
 numpy_driver = jit.JitDriver(
     greens=['shapelen', 'sig'],
@@ -39,7 +40,24 @@
     get_printable_location=signature.new_printable_location('slice'),
     name='numpy_slice',
 )
-
+count_driver = jit.JitDriver(
+    greens=['shapelen'],
+    virtualizables=['frame'],
+    reds=['s', 'frame', 'iter', 'arr'],
+    name='numpy_count'
+)
+filter_driver = jit.JitDriver(
+    greens=['shapelen', 'sig'],
+    virtualizables=['frame'],
+    reds=['concr', 'argi', 'ri', 'frame', 'v', 'res', 'self'],
+    name='numpy_filter',
+)
+filter_set_driver = jit.JitDriver(
+    greens=['shapelen', 'sig'],
+    virtualizables=['frame'],
+    reds=['idx', 'idxi', 'frame', 'arr'],
+    name='numpy_filterset',
+)
 
 def _find_shape_and_elems(space, w_iterable):
     shape = [space.len_w(w_iterable)]
@@ -270,6 +288,9 @@
     descr_gt = _binop_impl("greater")
     descr_ge = _binop_impl("greater_equal")
 
+    descr_and = _binop_impl("bitwise_and")
+    descr_or = _binop_impl("bitwise_or")
+
     def _binop_right_impl(ufunc_name):
         def impl(self, space, w_other):
             w_other = scalar_w(space,
@@ -479,11 +500,69 @@
     def _prepare_slice_args(self, space, w_idx):
         if (space.isinstance_w(w_idx, space.w_int) or
             space.isinstance_w(w_idx, space.w_slice)):
-            return [space.decode_index4(w_idx, self.shape[0])]
-        return [space.decode_index4(w_item, self.shape[i]) for i, w_item in
+            return [Chunk(*space.decode_index4(w_idx, self.shape[0]))]
+        return [Chunk(*space.decode_index4(w_item, self.shape[i])) for i, w_item in
                 enumerate(space.fixedview(w_idx))]
 
+    def count_all_true(self, arr):
+        sig = arr.find_sig()
+        frame = sig.create_frame(self)
+        shapelen = len(arr.shape)
+        s = 0
+        iter = None
+        while not frame.done():
+            count_driver.jit_merge_point(arr=arr, frame=frame, iter=iter, s=s,
+                                         shapelen=shapelen)
+            iter = frame.get_final_iter()
+            s += arr.dtype.getitem_bool(arr.storage, iter.offset)
+            frame.next(shapelen)
+        return s
+
+    def getitem_filter(self, space, arr):
+        concr = arr.get_concrete()
+        size = self.count_all_true(concr)
+        res = W_NDimArray(size, [size], self.find_dtype())
+        ri = ArrayIterator(size)
+        shapelen = len(self.shape)
+        argi = concr.create_iter()
+        sig = self.find_sig()
+        frame = sig.create_frame(self)
+        v = None
+        while not frame.done():
+            filter_driver.jit_merge_point(concr=concr, argi=argi, ri=ri,
+                                          frame=frame, v=v, res=res, sig=sig,
+                                          shapelen=shapelen, self=self)
+            if concr.dtype.getitem_bool(concr.storage, argi.offset):
+                v = sig.eval(frame, self)
+                res.setitem(ri.offset, v)
+                ri = ri.next(1)
+            else:
+                ri = ri.next_no_increase(1)
+            argi = argi.next(shapelen)
+            frame.next(shapelen)
+        return res
+
+    def setitem_filter(self, space, idx, val):
+        size = self.count_all_true(idx)
+        arr = SliceArray([size], self.dtype, self, val)
+        sig = arr.find_sig()
+        shapelen = len(self.shape)
+        frame = sig.create_frame(arr)
+        idxi = idx.create_iter()
+        while not frame.done():
+            filter_set_driver.jit_merge_point(idx=idx, idxi=idxi, sig=sig,
+                                              frame=frame, arr=arr,
+                                              shapelen=shapelen)
+            if idx.dtype.getitem_bool(idx.storage, idxi.offset):
+                sig.eval(frame, arr)
+                frame.next_from_second(1)
+            frame.next_first(shapelen)
+            idxi = idxi.next(shapelen)
+
     def descr_getitem(self, space, w_idx):
+        if (isinstance(w_idx, BaseArray) and w_idx.shape == self.shape and
+            w_idx.find_dtype().is_bool_type()):
+            return self.getitem_filter(space, w_idx)
         if self._single_item_result(space, w_idx):
             concrete = self.get_concrete()
             item = concrete._index_of_single_item(space, w_idx)
@@ -493,6 +572,11 @@
 
     def descr_setitem(self, space, w_idx, w_value):
         self.invalidated()
+        if (isinstance(w_idx, BaseArray) and w_idx.shape == self.shape and
+            w_idx.find_dtype().is_bool_type()):
+            return self.get_concrete().setitem_filter(space,
+                                                      w_idx.get_concrete(),
+                                             convert_to_array(space, w_value))
         if self._single_item_result(space, w_idx):
             concrete = self.get_concrete()
             item = concrete._index_of_single_item(space, w_idx)
@@ -509,9 +593,8 @@
     def create_slice(self, chunks):
         shape = []
         i = -1
-        for i, (start_, stop, step, lgt) in enumerate(chunks):
-            if step != 0:
-                shape.append(lgt)
+        for i, chunk in enumerate(chunks):
+            chunk.extend_shape(shape)
         s = i + 1
         assert s >= 0
         shape += self.shape[s:]
@@ -724,8 +807,7 @@
                                          frame=frame,
                                          ri=ri,
                                          self=self, result=result)
-            result.dtype.setitem(result.storage, ri.offset,
-                                 sig.eval(frame, self))
+            result.setitem(ri.offset, sig.eval(frame, self))
             frame.next(shapelen)
             ri = ri.next(shapelen)
         return result
@@ -945,7 +1027,7 @@
                             builder.append('\n' + indent)
                         else:
                             builder.append(indent)
-                    view = self.create_slice([(i, 0, 0, 1)]).get_concrete()
+                    view = self.create_slice([Chunk(i, 0, 0, 1)]).get_concrete()
                     view.to_str(space, comma, builder, indent=indent + ' ',
                                                     use_ellipsis=use_ellipsis)
                 if i < self.shape[0] - 1:
@@ -962,7 +1044,7 @@
                         builder.append(indent)
                 # create_slice requires len(chunks) > 1 in order to reduce
                 # shape
-                view = self.create_slice([(i, 0, 0, 1)]).get_concrete()
+                view = self.create_slice([Chunk(i, 0, 0, 1)]).get_concrete()
                 view.to_str(space, comma, builder, indent=indent + ' ',
                                                     use_ellipsis=use_ellipsis)
                 i += 1
@@ -1091,6 +1173,10 @@
                                parent)
         self.start = start
 
+    def create_iter(self):
+        return ViewIterator(self.start, self.strides, self.backstrides,
+                            self.shape)
+
     def setshape(self, space, new_shape):
         if len(self.shape) < 1:
             return
@@ -1137,6 +1223,9 @@
         self.shape = new_shape
         self.calc_strides(new_shape)
 
+    def create_iter(self):
+        return ArrayIterator(self.size)
+
     def create_sig(self):
         return signature.ArraySignature(self.dtype)
 
@@ -1191,6 +1280,7 @@
     arr = W_NDimArray(size, shape[:], dtype=dtype, order=order)
     shapelen = len(shape)
     arr_iter = ArrayIterator(arr.size)
+    # XXX we might want to have a jitdriver here
     for i in range(len(elems_w)):
         w_elem = elems_w[i]
         dtype.setitem(arr.storage, arr_iter.offset,
@@ -1257,6 +1347,9 @@
     __gt__ = interp2app(BaseArray.descr_gt),
     __ge__ = interp2app(BaseArray.descr_ge),
 
+    __and__ = interp2app(BaseArray.descr_and),
+    __or__ = interp2app(BaseArray.descr_or),
+
     __repr__ = interp2app(BaseArray.descr_repr),
     __str__ = interp2app(BaseArray.descr_str),
     __array_interface__ = GetSetProperty(BaseArray.descr_array_iface),
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
@@ -249,15 +249,16 @@
 
 
 class W_Ufunc2(W_Ufunc):
-    _immutable_fields_ = ["comparison_func", "func", "name"]
+    _immutable_fields_ = ["comparison_func", "func", "name", "int_only"]
     argcount = 2
 
     def __init__(self, func, name, promote_to_float=False, promote_bools=False,
-        identity=None, comparison_func=False):
+        identity=None, comparison_func=False, int_only=False):
 
         W_Ufunc.__init__(self, name, promote_to_float, promote_bools, identity)
         self.func = func
         self.comparison_func = comparison_func
+        self.int_only = int_only
 
     def call(self, space, args_w):
         from pypy.module.micronumpy.interp_numarray import (Call2,
@@ -268,6 +269,7 @@
         w_rhs = convert_to_array(space, w_rhs)
         calc_dtype = find_binop_result_dtype(space,
             w_lhs.find_dtype(), w_rhs.find_dtype(),
+            int_only=self.int_only,
             promote_to_float=self.promote_to_float,
             promote_bools=self.promote_bools,
         )
@@ -304,10 +306,12 @@
 
 
 def find_binop_result_dtype(space, dt1, dt2, promote_to_float=False,
-    promote_bools=False):
+    promote_bools=False, int_only=False):
     # dt1.num should be <= dt2.num
     if dt1.num > dt2.num:
         dt1, dt2 = dt2, dt1
+    if int_only and (not dt1.is_int_type() or not dt2.is_int_type()):
+        raise OperationError(space.w_TypeError, space.wrap("Unsupported types"))
     # Some operations promote op(bool, bool) to return int8, rather than bool
     if promote_bools and (dt1.kind == dt2.kind == interp_dtype.BOOLLTR):
         return interp_dtype.get_dtype_cache(space).w_int8dtype
@@ -425,6 +429,10 @@
             ("add", "add", 2, {"identity": 0}),
             ("subtract", "sub", 2),
             ("multiply", "mul", 2, {"identity": 1}),
+            ("bitwise_and", "bitwise_and", 2, {"identity": 1,
+                                               'int_only': True}),
+            ("bitwise_or", "bitwise_or", 2, {"identity": 0,
+                                             'int_only': True}),
             ("divide", "div", 2, {"promote_bools": True}),
             ("mod", "mod", 2, {"promote_bools": True}),
             ("power", "pow", 2, {"promote_bools": True}),
@@ -476,7 +484,7 @@
         extra_kwargs["identity"] = identity
 
         func = ufunc_dtype_caller(space, ufunc_name, op_name, argcount,
-            comparison_func=extra_kwargs.get("comparison_func", False)
+            comparison_func=extra_kwargs.get("comparison_func", False),
         )
         if argcount == 1:
             ufunc = W_Ufunc1(func, ufunc_name, **extra_kwargs)
diff --git a/pypy/module/micronumpy/signature.py b/pypy/module/micronumpy/signature.py
--- a/pypy/module/micronumpy/signature.py
+++ b/pypy/module/micronumpy/signature.py
@@ -82,6 +82,16 @@
         for i in range(len(self.iterators)):
             self.iterators[i] = self.iterators[i].next(shapelen)
 
+    @unroll_safe
+    def next_from_second(self, shapelen):
+        """ Don't increase the first iterator
+        """
+        for i in range(1, len(self.iterators)):
+            self.iterators[i] = self.iterators[i].next(shapelen)
+
+    def next_first(self, shapelen):
+        self.iterators[0] = self.iterators[0].next(shapelen)
+
     def get_final_iter(self):
         final_iter = promote(self.final_iter)
         if final_iter < 0:
diff --git a/pypy/module/micronumpy/strides.py b/pypy/module/micronumpy/strides.py
--- a/pypy/module/micronumpy/strides.py
+++ b/pypy/module/micronumpy/strides.py
@@ -10,12 +10,12 @@
     rstart = start
     rshape = []
     i = -1
-    for i, (start_, stop, step, lgt) in enumerate(chunks):
-        if step != 0:
-            rstrides.append(strides[i] * step)
-            rbackstrides.append(strides[i] * (lgt - 1) * step)
-            rshape.append(lgt)
-        rstart += strides[i] * start_
+    for i, chunk in enumerate(chunks):
+        if chunk.step != 0:
+            rstrides.append(strides[i] * chunk.step)
+            rbackstrides.append(strides[i] * (chunk.lgt - 1) * chunk.step)
+            rshape.append(chunk.lgt)
+        rstart += strides[i] * chunk.start
     # add a reminder
     s = i + 1
     assert s >= 0
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
@@ -2,6 +2,7 @@
 import py
 from pypy.module.micronumpy.test.test_base import BaseNumpyAppTest
 from pypy.module.micronumpy.interp_numarray import W_NDimArray, shape_agreement
+from pypy.module.micronumpy.interp_iter import Chunk
 from pypy.module.micronumpy import signature
 from pypy.interpreter.error import OperationError
 from pypy.conftest import gettestobjspace
@@ -37,53 +38,54 @@
 
     def test_create_slice_f(self):
         a = W_NDimArray(10 * 5 * 3, [10, 5, 3], MockDtype(), 'F')
-        s = a.create_slice([(3, 0, 0, 1)])
+        s = a.create_slice([Chunk(3, 0, 0, 1)])
         assert s.start == 3
         assert s.strides == [10, 50]
         assert s.backstrides == [40, 100]
-        s = a.create_slice([(1, 9, 2, 4)])
+        s = a.create_slice([Chunk(1, 9, 2, 4)])
         assert s.start == 1
         assert s.strides == [2, 10, 50]
         assert s.backstrides == [6, 40, 100]
-        s = a.create_slice([(1, 5, 3, 2), (1, 2, 1, 1), (1, 0, 0, 1)])
+        s = a.create_slice([Chunk(1, 5, 3, 2), Chunk(1, 2, 1, 1), Chunk(1, 0, 0, 1)])
         assert s.shape == [2, 1]
         assert s.strides == [3, 10]
         assert s.backstrides == [3, 0]
-        s = a.create_slice([(0, 10, 1, 10), (2, 0, 0, 1)])
+        s = a.create_slice([Chunk(0, 10, 1, 10), Chunk(2, 0, 0, 1)])
         assert s.start == 20
         assert s.shape == [10, 3]
 
     def test_create_slice_c(self):
         a = W_NDimArray(10 * 5 * 3, [10, 5, 3], MockDtype(), 'C')
-        s = a.create_slice([(3, 0, 0, 1)])
+        s = a.create_slice([Chunk(3, 0, 0, 1)])
         assert s.start == 45
         assert s.strides == [3, 1]
         assert s.backstrides == [12, 2]
-        s = a.create_slice([(1, 9, 2, 4)])
+        s = a.create_slice([Chunk(1, 9, 2, 4)])
         assert s.start == 15
         assert s.strides == [30, 3, 1]
         assert s.backstrides == [90, 12, 2]
-        s = a.create_slice([(1, 5, 3, 2), (1, 2, 1, 1), (1, 0, 0, 1)])
+        s = a.create_slice([Chunk(1, 5, 3, 2), Chunk(1, 2, 1, 1),
+                            Chunk(1, 0, 0, 1)])
         assert s.start == 19
         assert s.shape == [2, 1]
         assert s.strides == [45, 3]
         assert s.backstrides == [45, 0]
-        s = a.create_slice([(0, 10, 1, 10), (2, 0, 0, 1)])
+        s = a.create_slice([Chunk(0, 10, 1, 10), Chunk(2, 0, 0, 1)])
         assert s.start == 6
         assert s.shape == [10, 3]
 
     def test_slice_of_slice_f(self):
         a = W_NDimArray(10 * 5 * 3, [10, 5, 3], MockDtype(), 'F')
-        s = a.create_slice([(5, 0, 0, 1)])
+        s = a.create_slice([Chunk(5, 0, 0, 1)])
         assert s.start == 5
-        s2 = s.create_slice([(3, 0, 0, 1)])
+        s2 = s.create_slice([Chunk(3, 0, 0, 1)])
         assert s2.shape == [3]
         assert s2.strides == [50]
         assert s2.parent is a
         assert s2.backstrides == [100]
         assert s2.start == 35
-        s = a.create_slice([(1, 5, 3, 2)])
-        s2 = s.create_slice([(0, 2, 1, 2), (2, 0, 0, 1)])
+        s = a.create_slice([Chunk(1, 5, 3, 2)])
+        s2 = s.create_slice([Chunk(0, 2, 1, 2), Chunk(2, 0, 0, 1)])
         assert s2.shape == [2, 3]
         assert s2.strides == [3, 50]
         assert s2.backstrides == [3, 100]
@@ -91,16 +93,16 @@
 
     def test_slice_of_slice_c(self):
         a = W_NDimArray(10 * 5 * 3, [10, 5, 3], MockDtype(), order='C')
-        s = a.create_slice([(5, 0, 0, 1)])
+        s = a.create_slice([Chunk(5, 0, 0, 1)])
         assert s.start == 15 * 5
-        s2 = s.create_slice([(3, 0, 0, 1)])
+        s2 = s.create_slice([Chunk(3, 0, 0, 1)])
         assert s2.shape == [3]
         assert s2.strides == [1]
         assert s2.parent is a
         assert s2.backstrides == [2]
         assert s2.start == 5 * 15 + 3 * 3
-        s = a.create_slice([(1, 5, 3, 2)])
-        s2 = s.create_slice([(0, 2, 1, 2), (2, 0, 0, 1)])
+        s = a.create_slice([Chunk(1, 5, 3, 2)])
+        s2 = s.create_slice([Chunk(0, 2, 1, 2), Chunk(2, 0, 0, 1)])
         assert s2.shape == [2, 3]
         assert s2.strides == [45, 1]
         assert s2.backstrides == [45, 2]
@@ -108,14 +110,14 @@
 
     def test_negative_step_f(self):
         a = W_NDimArray(10 * 5 * 3, [10, 5, 3], MockDtype(), 'F')
-        s = a.create_slice([(9, -1, -2, 5)])
+        s = a.create_slice([Chunk(9, -1, -2, 5)])
         assert s.start == 9
         assert s.strides == [-2, 10, 50]
         assert s.backstrides == [-8, 40, 100]
 
     def test_negative_step_c(self):
         a = W_NDimArray(10 * 5 * 3, [10, 5, 3], MockDtype(), order='C')
-        s = a.create_slice([(9, -1, -2, 5)])
+        s = a.create_slice([Chunk(9, -1, -2, 5)])
         assert s.start == 135
         assert s.strides == [-30, 3, 1]
         assert s.backstrides == [-120, 12, 2]
@@ -124,7 +126,7 @@
         a = W_NDimArray(10 * 5 * 3, [10, 5, 3], MockDtype(), 'F')
         r = a._index_of_single_item(self.space, self.newtuple(1, 2, 2))
         assert r == 1 + 2 * 10 + 2 * 50
-        s = a.create_slice([(0, 10, 1, 10), (2, 0, 0, 1)])
+        s = a.create_slice([Chunk(0, 10, 1, 10), Chunk(2, 0, 0, 1)])
         r = s._index_of_single_item(self.space, self.newtuple(1, 0))
         assert r == a._index_of_single_item(self.space, self.newtuple(1, 2, 0))
         r = s._index_of_single_item(self.space, self.newtuple(1, 1))
@@ -134,7 +136,7 @@
         a = W_NDimArray(10 * 5 * 3, [10, 5, 3], MockDtype(), 'C')
         r = a._index_of_single_item(self.space, self.newtuple(1, 2, 2))
         assert r == 1 * 3 * 5 + 2 * 3 + 2
-        s = a.create_slice([(0, 10, 1, 10), (2, 0, 0, 1)])
+        s = a.create_slice([Chunk(0, 10, 1, 10), Chunk(2, 0, 0, 1)])
         r = s._index_of_single_item(self.space, self.newtuple(1, 0))
         assert r == a._index_of_single_item(self.space, self.newtuple(1, 2, 0))
         r = s._index_of_single_item(self.space, self.newtuple(1, 1))
@@ -1302,15 +1304,25 @@
         assert isinstance(i['data'][0], int)
         raises(TypeError, getattr, array(3), '__array_interface__')
 
+    def test_array_indexing_one_elem(self):
+        skip("not yet")
+        from _numpypy import array, arange
+        raises(IndexError, 'arange(3)[array([3.5])]')
+        a = arange(3)[array([1])]
+        assert a == 1
+        assert a[0] == 1
+        raises(IndexError,'arange(3)[array([15])]')
+        assert arange(3)[array([-3])] == 0
+        raises(IndexError,'arange(3)[array([-15])]')
+        assert arange(3)[array(1)] == 1
+
     def test_fill(self):
         from _numpypy import array
-
         a = array([1, 2, 3])
         a.fill(10)
         assert (a == [10, 10, 10]).all()
         a.fill(False)
         assert (a == [0, 0, 0]).all()
-
         b = a[:1]
         b.fill(4)
         assert (b == [4]).all()
@@ -1324,6 +1336,24 @@
         d.fill(100)
         assert d == 100
 
+    def test_array_indexing_bool(self):
+        from _numpypy import arange
+        a = arange(10)
+        assert (a[a > 3] == [4, 5, 6, 7, 8, 9]).all()
+        a = arange(10).reshape(5, 2)
+        assert (a[a > 3] == [4, 5, 6, 7, 8, 9]).all()
+        assert (a[a & 1 == 1] == [1, 3, 5, 7, 9]).all()
+
+    def test_array_indexing_bool_setitem(self):
+        from _numpypy import arange, array
+        a = arange(6)
+        a[a > 3] = 15
+        assert (a == [0, 1, 2, 3, 15, 15]).all()
+        a = arange(6).reshape(3, 2)
+        a[a & 1 == 1] = array([8, 9, 10])
+        assert (a == [[0, 8], [2, 9], [4, 10]]).all()
+
+
 
 class AppTestSupport(BaseNumpyAppTest):
     def setup_class(cls):
diff --git a/pypy/module/micronumpy/test/test_ufuncs.py b/pypy/module/micronumpy/test/test_ufuncs.py
--- a/pypy/module/micronumpy/test/test_ufuncs.py
+++ b/pypy/module/micronumpy/test/test_ufuncs.py
@@ -355,11 +355,20 @@
         raises(ValueError, maximum.reduce, [])
 
     def test_reduceND(self):
-        from numpypy import add, arange
+        from _numpypy import add, arange
         a = arange(12).reshape(3, 4)
         assert (add.reduce(a, 0) == [12, 15, 18, 21]).all()
         assert (add.reduce(a, 1) == [6.0, 22.0, 38.0]).all()
 
+    def test_bitwise(self):
+        from _numpypy import bitwise_and, bitwise_or, arange, array
+        a = arange(6).reshape(2, 3)
+        assert (a & 1 == [[0, 1, 0], [1, 0, 1]]).all()
+        assert (a & 1 == bitwise_and(a, 1)).all()
+        assert (a | 1 == [[1, 1, 3], [3, 5, 5]]).all()
+        assert (a | 1 == bitwise_or(a, 1)).all()
+        raises(TypeError, 'array([1.0]) & 1')
+
     def test_comparisons(self):
         import operator
         from _numpypy import equal, not_equal, less, less_equal, greater, greater_equal
diff --git a/pypy/module/micronumpy/test/test_zjit.py b/pypy/module/micronumpy/test/test_zjit.py
--- a/pypy/module/micronumpy/test/test_zjit.py
+++ b/pypy/module/micronumpy/test/test_zjit.py
@@ -217,6 +217,7 @@
         # This is the sum of the ops for both loops, however if you remove the
         # optimization then you end up with 2 float_adds, so we can still be
         # sure it was optimized correctly.
+        py.test.skip("too fragile")
         self.check_resops({'setinteriorfield_raw': 4, 'getfield_gc': 22,
                            'getarrayitem_gc': 4, 'getarrayitem_gc_pure': 2,
                            'getfield_gc_pure': 8,
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
@@ -94,6 +94,9 @@
             width, storage, i, offset
         ))
 
+    def read_bool(self, storage, width, i, offset):
+        raise NotImplementedError
+
     def store(self, storage, width, i, offset, box):
         value = self.unbox(box)
         libffi.array_setitem(clibffi.cast_type_to_ffitype(self.T),
@@ -168,6 +171,7 @@
     @simple_binary_op
     def min(self, v1, v2):
         return min(v1, v2)
+    
 
 class Bool(BaseType, Primitive):
     T = lltype.Bool
@@ -185,6 +189,11 @@
         else:
             return self.False
 
+
+    def read_bool(self, storage, width, i, offset):
+        return libffi.array_getitem(clibffi.cast_type_to_ffitype(self.T),
+                                    width, storage, i, offset)
+
     def coerce_subtype(self, space, w_subtype, w_item):
         # Doesn't return subclasses so it can return the constants.
         return self._coerce(space, w_item)
@@ -253,6 +262,14 @@
             assert v == 0
             return 0
 
+    @simple_binary_op
+    def bitwise_and(self, v1, v2):
+        return v1 & v2
+
+    @simple_binary_op
+    def bitwise_or(self, v1, v2):
+        return v1 | v2
+
 class Int8(BaseType, Integer):
     T = rffi.SIGNEDCHAR
     BoxType = interp_boxes.W_Int8Box


More information about the pypy-commit mailing list