[FIXED] Python - create a class as index for 2-dimension numpy array

Issue

I have some code of index class TwoDimIndex. I wanna use it to indexes of numpy 2-dimension array like arr[idx]. Below is code of that class.

import numpy as np

class TwoDimIndex:
    def __init__(self, i1, i2):
        self.i1 = i1
        self.i2 = i2
        pass

    def __index__(self):
        return self.i1, self.i2 
        # return np.array([self.i1, self.i2]) - Tried this too
        
    def __eq__(self, other):
        return self.i1 == other
        
    def __int__(self):
        return self.i1
        
    def __hash__(self):
        return hash(self.__int__())

# Don't edit code of this function please
def useful_func():
    idx = TwoDimIndex(1, 1) # Can edit create instance
    arr_two_dim = np.array([[0, 1], [2, 3]])
    print(idx.__index__() == (1, 1))
    print(arr_two_dim[1, 1], arr_two_dim[idx.__index__()]) # Success
    print(arr_two_dim[idx]) # Crash here
    # But I want this code work!!!

useful_func()

Class TwoDimIndex is use to be index, for example arr[TwoDimIndex()].

All code for Jupyter Notebook and Python 3.8. But I get error when execute this code.

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

Is there any ways to make an instance of a class TwoDimIndex an numpy 2-d array index?

Solution

I find short solution with inherit from tuple.

class TwoDimIndex(tuple):
    def __init__(self, tup):
        self.tup = tuple(tup)
        pass

    def __index__(self):
        return self.tup

def useful_func():
    idx = TwoDimIndex([1, 1]) # Edit at create class
    arr_two_dim = np.array([[0, 1], [2, 3]])
    print(idx.__index__() == (1, 1))
    print(arr_two_dim[1, 1], arr_two_dim[idx.__index__()]) # Success
    print(arr_two_dim[idx]) # Success now

useful_func()

Not sure it is right way but it work.

Answered By - Константин Ушаков