How to correctly asses the typing of an index when writing type hints?

Question

I am writing a custom Dataset class for a particular deep learning problem and the question of how to annotate the index of __getitem__ using type hints arose.

Given

    def __getitem__(self, i: int) -> Tuple[Tensor, Tensor]: 
        return self.X[i], self.y[i]

what would be the best way to annotate i given the fact that it can be a slice (eg. X[:i])?

Answer 1

So, as I said, i can be and integer or a slice object.

So you can try:

def __getitem__(self, i: Union[int, slice]) -> Tuple[Tensor, Tensor]:
    ...