[pypy-commit] pypy default: provide ndarray.__index__()

[bdkearns]

Author: Brian Kearns <bdkearns at gmail.com>
Branch: 
Changeset: r68458:42d72b7298bc
Date: 2013-12-17 23:51 -0500
http://bitbucket.org/pypy/pypy/changeset/42d72b7298bc/

Log:	provide ndarray.__index__()

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
@@ -992,7 +992,8 @@
         elif shape == [1]:
             value = self.descr_getitem(space, space.wrap(0))
         else:
-            raise OperationError(space.w_TypeError, space.wrap("only length-1 arrays can be converted to Python scalars"))
+            raise OperationError(space.w_TypeError, space.wrap(
+                "only length-1 arrays can be converted to Python scalars"))
         if self.get_dtype().is_str_or_unicode():
             raise OperationError(space.w_TypeError, space.wrap(
                 "don't know how to convert scalar number to int"))
@@ -1006,7 +1007,8 @@
         elif shape == [1]:
             value = self.descr_getitem(space, space.wrap(0))
         else:
-            raise OperationError(space.w_TypeError, space.wrap("only length-1 arrays can be converted to Python scalars"))
+            raise OperationError(space.w_TypeError, space.wrap(
+                "only length-1 arrays can be converted to Python scalars"))
         if self.get_dtype().is_str_or_unicode():
             raise OperationError(space.w_TypeError, space.wrap(
                 "don't know how to convert scalar number to long"))
@@ -1020,12 +1022,30 @@
         elif shape == [1]:
             value = self.descr_getitem(space, space.wrap(0))
         else:
-            raise OperationError(space.w_TypeError, space.wrap("only length-1 arrays can be converted to Python scalars"))
+            raise OperationError(space.w_TypeError, space.wrap(
+                "only length-1 arrays can be converted to Python scalars"))
         if self.get_dtype().is_str_or_unicode():
             raise OperationError(space.w_TypeError, space.wrap(
                 "don't know how to convert scalar number to float"))
         return space.float(value)
 
+    def descr_index(self, space):
+        shape = self.get_shape()
+        if len(shape) == 0:
+            assert isinstance(self.implementation, scalar.Scalar)
+            value = space.wrap(self.implementation.get_scalar_value())
+        elif shape == [1]:
+            value = self.descr_getitem(space, space.wrap(0))
+        else:
+            raise OperationError(space.w_TypeError, space.wrap(
+                "only integer arrays with one element "
+                "can be converted to an index"))
+        if not self.get_dtype().is_int_type() or self.get_dtype().is_bool_type():
+            raise OperationError(space.w_TypeError, space.wrap(
+                "only integer arrays with one element "
+                "can be converted to an index"))
+        return value.item(space)
+
     def descr_reduce(self, space):
         from rpython.rlib.rstring import StringBuilder
         from pypy.interpreter.mixedmodule import MixedModule
@@ -1204,6 +1224,7 @@
     __long__ = interp2app(W_NDimArray.descr_long),
     __float__ = interp2app(W_NDimArray.descr_float),
     __buffer__ = interp2app(W_NDimArray.descr_get_data),
+    __index__ = interp2app(W_NDimArray.descr_index),
 
     __pos__ = interp2app(W_NDimArray.descr_pos),
     __neg__ = interp2app(W_NDimArray.descr_neg),
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
@@ -2718,6 +2718,17 @@
         assert b[0] == 1
         assert b[1] == 'ab'
 
+    def test_index(self):
+        import numpy as np
+        a = np.array([1], np.uint16)
+        i = a.__index__()
+        assert type(i) is int
+        assert i == 1
+        for a in [np.array('abc'), np.array([1,2]), np.array([True])]:
+            exc = raises(TypeError, a.__index__)
+            assert exc.value.message == 'only integer arrays with one element ' \
+                                        'can be converted to an index'
+
     def test_int_array_index(self):
         from numpypy import array
         assert (array([])[[]] == []).all()