Search code examples
pythonarraysnumpyindexingmagic-methods

numpy ndarray indexing with __index__ method

I don't understand how indexing of a numpy ndarray works, when using a custom class instance as the index.

I have the following code:

import numpy as np

class MyClass:
    def __index__(self):
        return 1,2

foo = np.array([[1,2,3],[4,5,6]])
bar = MyClass()

print(foo[1,2])
print(foo[bar])

I expect to get the same result (6) from both print functions. But from the second one, where the class instance is used a the index, I receive an error:

IndexError: only integers, slices (`:`), ellipsis (`...`), numpy.newaxis (`None`) and integer or boolean arrays are valid indices

If I call the __index__ method explicitly with

print(foo[bar.__index__()])

it works. But this defeats the purpose of the magic method.

If I call the array with just one index, everything works fine:

import numpy as np

class MyClass:
    def __index__(self):
        return 1

foo = np.array([[1,2,3],[4,5,6]])
bar = MyClass()

print(foo[1])
print(foo[bar])

>>> [4 5 6]
>>> [4 5 6]

So what I don't get:

  • The ndarray can use the output of the __index__ method for multiple dimensions. Seen when calling it explicitly.
  • The ndarray does coll the __index__ method. Seen in the second example.
  • But for some reason, this does not work at the same time. Why?

Did I miss something, or does the ndarray not support this kind of indexing?

I just want to add, that it apparently doesn't matter, how the __index__ method outputs its result. I tried:

return a, b
return (a, b)
return tuple((a, b))

None of them worked for me.

Solution

  • As mentioned here, __index__ method Must return an integer.

    That's why your attempt didn't work, while the "one index" example worked.