Search code examples
pythonpytorchpython-typing

What is the type annotation for return value of __getitem__?

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?
Solution

  • 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.