[pypy-commit] pypy numpy-multidim: basic slice support

[fijal]

Author: Maciej Fijalkowski <fijall at gmail.com>
Branch: numpy-multidim
Changeset: r48583:252e03277a09
Date: 2011-10-28 19:46 +0200
http://bitbucket.org/pypy/pypy/changeset/252e03277a09/

Log:	basic slice support

diff --git a/pypy/module/micronumpy/compile.py b/pypy/module/micronumpy/compile.py
--- a/pypy/module/micronumpy/compile.py
+++ b/pypy/module/micronumpy/compile.py
@@ -53,7 +53,10 @@
         return False
 
     def decode_index4(self, w_idx, size):
-        return (self.int_w(w_idx), 0, 0, 1)
+        if isinstance(w_idx, IntObject):
+            return (self.int_w(w_idx), 0, 0, 1)
+        else:
+            return (0, size, 1, size)
 
     @specialize.argtype(1)
     def wrap(self, obj):
@@ -63,7 +66,7 @@
             return BoolObject(obj)
         elif isinstance(obj, int):
             return IntObject(obj)
-        raise Exception
+        return obj
 
     def newlist(self, items):
         return ListObject(items)
@@ -138,6 +141,9 @@
     def __init__(self, items):
         self.items = items
 
+class SliceObject(W_Root):
+    tp = FakeSpace.w_slice
+
 class InterpreterState(object):
     def __init__(self, code):
         self.code = code
@@ -210,11 +216,11 @@
 
     def execute(self, interp):
         w_lhs = self.lhs.execute(interp)
+        if isinstance(self.rhs, SliceConstant):
+            w_rhs = self.rhs.wrap(interp.space)
+        else:
+            w_rhs = self.rhs.execute(interp)
         assert isinstance(w_lhs, BaseArray)
-        if isinstance(self.rhs, SliceConstant):
-            # XXX interface has changed on multidim branch
-            raise NotImplementedError
-        w_rhs = self.rhs.execute(interp)
         if self.name == '+':
             w_res = w_lhs.descr_add(interp.space, w_rhs)
         elif self.name == '*':
@@ -223,12 +229,10 @@
             w_res = w_lhs.descr_sub(interp.space, w_rhs)            
         elif self.name == '->':
             if isinstance(w_rhs, Scalar):
-                index = int(interp.space.float_w(
-                    w_rhs.value.wrap(interp.space)))
-                dtype = interp.space.fromcache(W_Float64Dtype)
-                return Scalar(dtype, w_lhs.get_concrete().eval(index))
-            else:
-                raise NotImplementedError
+                w_rhs = w_rhs.eval(0).wrap(interp.space)
+                assert isinstance(w_rhs, FloatObject)
+                w_rhs = IntObject(int(w_rhs.floatval))
+            w_res = w_lhs.descr_getitem(interp.space, w_rhs)
         else:
             raise NotImplementedError
         if not isinstance(w_res, BaseArray):
@@ -293,6 +297,9 @@
     def __init__(self):
         pass
 
+    def wrap(self, space):
+        return SliceObject()
+
     def __repr__(self):
         return 'slice()'
 
diff --git a/pypy/module/micronumpy/test/test_compile.py b/pypy/module/micronumpy/test/test_compile.py
--- a/pypy/module/micronumpy/test/test_compile.py
+++ b/pypy/module/micronumpy/test/test_compile.py
@@ -161,10 +161,9 @@
         assert interp.results[0].value.val == 256
 
     def test_slice(self):
-        py.test.skip("in progress")
         interp = self.run("""
         a = [1,2,3,4]
         b = a -> :
         b -> 3
         """)
-        assert interp.results[0].value.val == 3
+        assert interp.results[0].value.val == 4