What is the type annotation for return value of __getitem__?

Question

I was making a dataset class by inheriting torch.utils.data.Dataset and came across the following issue.

Unlike previous functions which returns a fixed type of values, __getitem__ doesn't. For example,

class MyExample:
    def __init__(self, some_list: list[int]):
        self.some_list = some_list

    def __getitem__(self, index):
        return self.some_list[index]

MyExample[<index>] will return int, while MyExample[<slice>] will return slice of int. VScode intellisense automatically write T_co for its type annotation, but I don't understand what this means.

• How should I annotate such a function?
• What do I need to learn in order to understand such an annotation?

Answer 1

You can use Union to describe the possibility of multiple different return types.

from typing import Union

class MyExample():
    def __init__(self, some_list: list[int]):
        self.some_list = some_list

    def __getitem__(self, index: Union[int, slice]) -> Union[int, list[int]]:
        return self.some_list[index]

a = MyExample([1, 2, 3, 4, 5])
print(a[2])
print(a[1:3])

mypy reports no issues with the above code snippet.