[pypy-commit] pypy numpy-refactor: single item setitem

[fijal]

Author: Maciej Fijalkowski <fijall at gmail.com>
Branch: numpy-refactor
Changeset: r56934:782a9c9df7e2
Date: 2012-08-30 15:24 +0200
http://bitbucket.org/pypy/pypy/changeset/782a9c9df7e2/

Log:	single item setitem

diff --git a/pypy/module/micronumpy/arrayimpl/concrete.py b/pypy/module/micronumpy/arrayimpl/concrete.py
--- a/pypy/module/micronumpy/arrayimpl/concrete.py
+++ b/pypy/module/micronumpy/arrayimpl/concrete.py
@@ -13,7 +13,7 @@
         self.size = array.size
 
     def setitem(self, elem):
-        self.dtype.setitem(self.array, self.offset, elem)
+        self.array.setitem(self.offset, elem)
 
     def getitem(self):
         return self.array.getitem(self.offset)
@@ -65,6 +65,9 @@
     def getitem(self, index):
         return self.dtype.getitem(self, index)
 
+    def setitem(self, index, value):
+        self.dtype.setitem(self, index, value)
+
     # -------------------- applevel get/setitem -----------------------
 
     @jit.unroll_safe
@@ -115,3 +118,13 @@
             chunks = self._prepare_slice_args(space, w_index)
             return chunks.apply(self)
 
+    def descr_setitem(self, space, w_index, w_value):
+        try:
+            item = self._single_item_index(space, w_index)
+            self.setitem(item, self.dtype.coerce(space, w_value))
+        except IndexError:
+            w_value = support.convert_to_array(space, w_value)
+            chunks = self._prepare_slice_args(space, w_index)
+            view = chunks.apply(self)
+            view.setslice(space, w_value)
+
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
@@ -3,7 +3,7 @@
 from pypy.interpreter.error import operationerrfmt
 from pypy.interpreter.typedef import TypeDef, GetSetProperty
 from pypy.interpreter.gateway import interp2app, unwrap_spec
-from pypy.module.micronumpy import interp_dtype, interp_ufuncs
+from pypy.module.micronumpy import interp_dtype, interp_ufuncs, support
 from pypy.module.micronumpy.arrayimpl import create_implementation
 from pypy.module.micronumpy.strides import find_shape_and_elems
 from pypy.tool.sourcetools import func_with_new_name
@@ -56,6 +56,13 @@
             return self.getitem_filter(space, w_idx)
         return self.implementation.descr_getitem(space, w_idx)
 
+    def descr_setitem(self, space, w_idx, w_value):
+        if (isinstance(w_idx, W_NDimArray) and w_idx.shape == self.shape and
+            w_idx.find_dtype().is_bool_type()):
+            return self.setitem_filter(space, w_idx,
+                                       support.convert_to_array(space, w_value))
+        self.implementation.descr_setitem(space, w_idx, w_value)
+
     def create_iter(self):
         return self.implementation.create_iter()
 
@@ -88,6 +95,7 @@
     __add__ = interp2app(W_NDimArray.descr_add),
 
     __getitem__ = interp2app(W_NDimArray.descr_getitem),
+    __setitem__ = interp2app(W_NDimArray.descr_setitem),
 
     dtype = GetSetProperty(W_NDimArray.descr_get_dtype),
     shape = GetSetProperty(W_NDimArray.descr_get_shape,