Author: mattip Branch: numppy-flatitter Changeset: r51471:60b724406de5 Date: 2012-01-19 01:46 +0200 http://bitbucket.org/pypy/pypy/changeset/60b724406de5/
Log: add setitem, getitem to flatitter 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 @@ -1299,13 +1299,18 @@ size = 1 for sh in arr.shape: size *= sh - self.strides = [arr.strides[-1]] - self.backstrides = [arr.backstrides[-1]] - ViewArray.__init__(self, size, [size], arr.dtype, arr.order, - arr) + if arr.strides[-1] < arr.strides[0]: + self.strides = [arr.strides[-1]] + self.backstrides = [arr.backstrides[-1]] + else: + self.strides = [arr.strides[0]] + self.backstrides = [arr.backstrides[0]] + ViewArray.__init__(self, size, [size], arr.dtype, order=arr.order, + parent=arr) self.shapelen = len(arr.shape) self.iter = OneDimIterator(arr.start, self.strides[0], self.shape[0]) + self.base = arr def descr_next(self, space): if self.iter.done(): @@ -1317,9 +1322,42 @@ def descr_iter(self): return self + def descr_getitem(self, space, w_idx): + if not space.isinstance_w(w_idx, space.w_int): + raise OperationError(space.w_ValueError, space.wrap( + "non-integer indexing not supported yet")) + _i = space.int_w(w_idx) + if _i<0: + i = self.size + _i + else: + i = _i + if i >= self.size or i < 0: + raise operationerrfmt(space.w_IndexError, + "index (%d) out of range (%d<=index<%d", + _i, -self.size, self.size) + result = self.getitem(self.base.start + i * self.strides[0]) + return result + + def descr_setitem(self, space, w_idx, w_value): + if not space.isinstance_w(w_idx, space.w_int): + raise OperationError(space.w_ValueError, space.wrap( + "non-integer indexing not supported yet")) + _i = space.int_w(w_idx) + if _i<0: + i = self.size + _i + else: + i = _i + if i >= self.size or i < 0: + raise operationerrfmt(space.w_IndexError, + "index (%d) out of range (%d<=index<%d", + _i, -self.size, self.size) + self.setitem(self.base.start + i * self.strides[0], w_value) + W_FlatIterator.typedef = TypeDef( 'flatiter', next = interp2app(W_FlatIterator.descr_next), __iter__ = interp2app(W_FlatIterator.descr_iter), + __getitem__ = interp2app(W_FlatIterator.descr_getitem), + __setitem__ = interp2app(W_FlatIterator.descr_setitem), ) W_FlatIterator.acceptable_as_base_class = False 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 @@ -276,6 +276,12 @@ for i in xrange(5): assert a[i] == b[i] + def test_getitem_nd(self): + from _numpypy import arange + a = arange(15).reshape(3, 5) + assert a[1, 3] == 8 + assert a.T[1, 2] == 11 + def test_setitem(self): from _numpypy import array a = array(range(5)) @@ -1286,6 +1292,29 @@ a = ones((2, 2)) assert list(((a + a).flat)) == [2, 2, 2, 2] + def test_flatiter_getitem(self): + from _numpypy import arange + a = arange(10) + assert a.flat[3] == 3 + assert a[2:].flat[3] == 5 + assert (a + a).flat[3] == 6 + assert a[::2].flat[3] == 6 + assert a.reshape(2,5).flat[3] == 3 + b = a.flat + b.next() + b.next() + b.next() + assert b[3] == 3 + assert b[-2] == 8 + raises(IndexError, "b[11]") + raises(IndexError, "b[-11]") + + def test_flatiter_transpose(self): + from _numpypy import arange + a = arange(10) + skip('out-of-order transformations do not work yet') + assert a.reshape(2,5).T.flat[3] == 6 + def test_slice_copy(self): from _numpypy import zeros a = zeros((10, 10)) _______________________________________________ pypy-commit mailing list pypy-commit@python.org http://mail.python.org/mailman/listinfo/pypy-commit