Python - 创建一个类作为二维 numpy 数组的索引

Question

我有一些索引类代码TwoDimIndex。我想用它来索引像arr[idx]这样的numpy二维数组。下面是该类的代码。

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

类TwoDimIndex用作索引，例如arr[TwoDimIndex()]。

Jupyter Notebook 和 Python 3.8 的所有代码。但是执行此代码时出现错误。

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

有什么方法可以让 TwoDimIndex 类的实例成为一个 numpy 二维数组索引？

Answer 1

我找到了从元组继承的简短解决方案。

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

不确定它是否正确，但它有效。