What is the type annotation for return value of __getitem__?

user_8471995 0
Viewed 214

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?
pytorch type-hinting python
2 Answers

For this you can use Unions as the actual annotation of the function and overload to allow you type checker to know that a slice is only returned when a slice is given and an int is given when only an int is given.

from typing import overload, Union

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

    @overload
    def __getitem__(self, index: int) -> int: ...

    @overload
    def __getitem__(self, index: slice) -> list[int]: ...

    def __getitem__(self, index: Union[slice, int]) -> Union[list[int], int]: # Or Slice[int] if that's a thing
        return self.some_list[index]

If your example is not only for list of ints but for general lists, you can change to Generic types.

user_140787580

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.

user_177698150
Related