Commits

Brian Kearns committed b23a46d

implement array.itemset

Comments (0)

Files changed (2)

pypy/module/micronumpy/interp_numarray.py

         raise OperationError(space.w_NotImplementedError, space.wrap(
             "non-int arg not supported"))
 
+    def descr_itemset(self, space, args_w):
+        if len(args_w) == 0:
+            raise OperationError(space.w_ValueError, space.wrap(
+                "itemset must have at least one argument"))
+        if len(args_w) != len(self.get_shape()) + 1:
+            raise OperationError(space.w_ValueError, space.wrap(
+                "incorrect number of indices for array"))
+        self.descr_setitem(space, space.newtuple(args_w[:-1]), args_w[-1])
+
     def descr___array__(self, space, w_dtype=None):
         if not space.is_none(w_dtype):
             raise OperationError(space.w_NotImplementedError, space.wrap(
         raise OperationError(space.w_NotImplementedError, space.wrap(
             "getfield not implemented yet"))
 
-    def descr_itemset(self, space, w_arg):
-        raise OperationError(space.w_NotImplementedError, space.wrap(
-            "itemset not implemented yet"))
-
     @unwrap_spec(new_order=str)
     def descr_newbyteorder(self, space, new_order=NPY_SWAP):
         return self.descr_view(space,
     flat = GetSetProperty(W_NDimArray.descr_get_flatiter,
                           W_NDimArray.descr_set_flatiter),
     item = interp2app(W_NDimArray.descr_item),
+    itemset = interp2app(W_NDimArray.descr_itemset),
     real = GetSetProperty(W_NDimArray.descr_get_real,
                           W_NDimArray.descr_set_real),
     imag = GetSetProperty(W_NDimArray.descr_get_imag,

pypy/module/micronumpy/test/test_numarray.py

         assert b[0] == 1
         assert b[1] == 'ab'
 
+    def test_itemset(self):
+        import numpy as np
+        a = np.array(range(5))
+        exc = raises(ValueError, a.itemset)
+        assert exc.value[0] == 'itemset must have at least one argument'
+        exc = raises(ValueError, a.itemset, 1, 2, 3)
+        assert exc.value[0] == 'incorrect number of indices for array'
+        a.itemset(1, 5)
+        assert a[1] == 5
+        a = np.array(range(6)).reshape(2, 3)
+        a.itemset(1, 2, 100)
+        assert a[1, 2] == 100
+
     def test_index(self):
         import numpy as np
         a = np.array([1], np.uint16)