Commits

Brian Kearns committed 42d72b7

provide ndarray.__index__()

  • Participants
  • Parent commits e8735af

Comments (0)

Files changed (2)

File pypy/module/micronumpy/interp_numarray.py

         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"))
         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"))
         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
     __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),

File pypy/module/micronumpy/test/test_numarray.py

         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()