[pypy-commit] pypy default: copy logic from getitem to setitem for int idexing of record ndarrays

[mattip]

Author: mattip <matti.picus at gmail.com>
Branch: 
Changeset: r80364:1f1f0b0f5f07
Date: 2015-10-21 08:49 +1100
http://bitbucket.org/pypy/pypy/changeset/1f1f0b0f5f07/

Log:	copy logic from getitem to setitem for int idexing of record
	ndarrays

diff --git a/pypy/module/micronumpy/boxes.py b/pypy/module/micronumpy/boxes.py
--- a/pypy/module/micronumpy/boxes.py
+++ b/pypy/module/micronumpy/boxes.py
@@ -591,6 +591,14 @@
     def descr_setitem(self, space, w_item, w_value):
         if space.isinstance_w(w_item, space.w_basestring):
             item = space.str_w(w_item)
+        elif space.isinstance_w(w_item, space.w_int):
+            indx = space.int_w(w_item)
+            try:
+                item = self.dtype.names[indx][0]
+            except IndexError:
+                if indx < 0:
+                    indx += len(self.dtype.names)
+                raise oefmt(space.w_IndexError, "invalid index (%d)", indx)
         else:
             raise oefmt(space.w_IndexError, "invalid index")
         try:
diff --git a/pypy/module/micronumpy/test/test_ndarray.py b/pypy/module/micronumpy/test/test_ndarray.py
--- a/pypy/module/micronumpy/test/test_ndarray.py
+++ b/pypy/module/micronumpy/test/test_ndarray.py
@@ -645,6 +645,12 @@
         for i in xrange(5):
             assert a[i] == i
 
+    def test_setitem_record(self):
+        from numpy import zeros
+        trie =  zeros(200, dtype= [ ("A","uint32"),("C","uint32"), ])
+        trie[0][0] = 1
+        assert trie[0]['A'] == 1
+
     def test_setitem_array(self):
         import numpy as np
         a = np.array((-1., 0, 1))/0.