[pypy-commit] pypy ndarray-ptp: refactored put and added more tests.

[andrewsmedina]

Author: Andrews Medina <andrewsmedina at gmail.com>
Branch: ndarray-ptp
Changeset: r65139:935fc2578913
Date: 2013-06-28 00:54 -0300
http://bitbucket.org/pypy/pypy/changeset/935fc2578913/

Log:	refactored put and added more tests.

diff --git a/pypy/module/micronumpy/interp_arrayops.py b/pypy/module/micronumpy/interp_arrayops.py
--- a/pypy/module/micronumpy/interp_arrayops.py
+++ b/pypy/module/micronumpy/interp_arrayops.py
@@ -197,23 +197,35 @@
 def put(space, w_arr, w_indices, w_values, mode='raise'):
     from pypy.module.micronumpy import constants
     from pypy.module.micronumpy.support import int_w
+
     arr = convert_to_array(space, w_arr)
+
     if mode not in constants.MODES:
         raise OperationError(space.w_ValueError,
                              space.wrap("mode %s not known" % (mode,)))
-    indices = convert_to_array(space, w_indices)
-    values = convert_to_array(space, w_values)
-    if not indices:
+    if not w_indices:
         raise OperationError(space.w_ValueError,
                              space.wrap("indice list cannot be empty"))
-    if not values:
+    if not w_values:
         raise OperationError(space.w_ValueError,
                              space.wrap("value list cannot be empty"))
+
     dtype = arr.get_dtype()
-    val_iter = values.create_iter()
-    ind_iter = indices.create_iter()
-    while not ind_iter.done():
-        index = int_w(space, ind_iter.getitem())
+
+    if space.isinstance_w(w_indices, space.w_list):
+        indices = space.listview(w_indices)
+    else:
+        indices = [w_indices]
+
+    if space.isinstance_w(w_values, space.w_list):
+        values = space.listview(w_values)
+    else:
+        values = [w_values]
+
+    v_idx = 0
+    for idx in indices:
+        index = int_w(space, idx)
+
         if index < 0 or index >= arr.get_size():
             if constants.MODES[mode] == constants.MODE_RAISE:
                 raise OperationError(space.w_ValueError, space.wrap(
@@ -226,9 +238,13 @@
                     index = 0
                 else:
                     index = arr.get_size() - 1
-        arr.setitem(space, [index], val_iter.getitem().convert_to(dtype))
-        ind_iter.next()
-        val_iter.next()
+
+        value = values[v_idx]
+
+        if v_idx + 1 < len(values):
+            v_idx += 1
+
+        arr.setitem(space, [index], dtype.coerce(space, value))
 
 
 def diagonal(space, arr, offset, axis1, axis2):
diff --git a/pypy/module/micronumpy/test/test_arrayops.py b/pypy/module/micronumpy/test/test_arrayops.py
--- a/pypy/module/micronumpy/test/test_arrayops.py
+++ b/pypy/module/micronumpy/test/test_arrayops.py
@@ -136,8 +136,14 @@
     def test_put_basic(self):
         from numpypy import arange, array
         a = arange(5)
-        a.put([0,2], [-44, -55])
+        a.put([0, 2], [-44, -55])
         assert (a == array([-44, 1, -55, 3, 4])).all()
+        a = arange(5)
+        a.put([3, 4], 9)
+        assert (a == array([0, 1, 2, 9, 9])).all()
+        a = arange(5)
+        a.put(1, [7, 8])
+        assert (a == array([0, 7, 2, 3, 4])).all()
 
     def test_put_modes(self):
         from numpypy import array, arange