[pypy-commit] pypy numpy-setslice: Fixed setslice and added extra tests for slice of slice. Everything appears to be working for setting slices.

[justinpeel]

Author: Justin Peel <notmuchtotell at gmail.com>
Branch: numpy-setslice
Changeset: r45747:818b17e123de
Date: 2011-07-18 18:15 -0600
http://bitbucket.org/pypy/pypy/changeset/818b17e123de/

Log:	Fixed setslice and added extra tests for slice of slice. Everything
	appears to be working for setting slices.

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
@@ -402,7 +402,6 @@
         BaseArray.__init__(self)
         self.signature = signature
         self.parent = parent
-        self.storage = parent.storage
         self.invalidates = parent.invalidates
 
     def get_concrete(self):
@@ -436,17 +435,20 @@
         if isinstance(parent, SingleDimSlice):
             self.start = parent.calc_index(start)
             self.stop = parent.calc_index(stop)
-            self.step = parent.step * self.step
+            self.step = parent.step * step
+            self.parent = parent.parent
         else:
             self.start = start
             self.stop = stop
             self.step = step
+            self.parent = parent
         self.size = slice_length
 
     def find_size(self):
         return self.size
 
     def _sliceloop1(self, start, stop, step, arr):
+        storage = self.parent.storage
         signature = Signature()
         new_sig = self.signature.transition(signature)
         i = start
@@ -454,11 +456,12 @@
         while i < stop:
             slice_driver1.jit_merge_point(signature=signature, self=self,
                     step=step, stop=stop, i=i, j=j, arr=arr)
-            self.storage[i] = arr.eval(j)
+            storage[i] = arr.eval(j)
             j += 1
             i += step
 
     def _sliceloop2(self, start, stop, step, arr):
+        storage = self.parent.storage
         signature = Signature()
         new_sig = self.signature.transition(signature)
         i = start
@@ -466,7 +469,7 @@
         while i > stop:
             slice_driver2.jit_merge_point(signature=signature, self=self,
                     step=step, stop=stop, i=i, j=j, arr=arr)
-            self.storage[i] = arr.eval(j)
+            storage[i] = arr.eval(j)
             j += 1
             i += step
 
@@ -553,12 +556,6 @@
 
     def setslice(self, space, start, stop, step, slice_length, arr):
         i = start
-        #if stop < 0:
-        #    stop += self.find_size()
-        #if step > 0:
-        #    stop = min(stop, self.find_size())
-        #else:
-        #    stop = max(stop, 0)
         if not isinstance(arr, BaseArray):
             arr = convert_to_array(space, arr)
         if step > 0:
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
@@ -67,16 +67,25 @@
         a[1:4:2] = b
         assert a[1] == 0.
         assert a[3] == 1.
+
+    def test_setslice_of_slice_array(self):
+        from numpy import array, zeros
+        a = zeros(5)
         a[::2] = array([9., 10., 11.])
         assert a[0] == 9.
         assert a[2] == 10.
         assert a[4] == 11.
-        a[1:4:2][::-1] = b
+        a[1:4:2][::-1] = array([1., 2.])
         assert a[0] == 9.
-        assert a[1] == 1.
+        assert a[1] == 2.
         assert a[2] == 10.
-        assert a[3] == 0.
+        assert a[3] == 1.
         assert a[4] == 11.
+        a = zeros(10)
+        a[::2][::-1][::2] = array(range(1,4))
+        a[8] = 1.
+        a[4] = 2.
+        a[0] = 3.
 
     def test_setslice_list(self):
         from numpy import array