Return type of slice for a user-made container in python

I am creating a custom container that returns an instance of itself when sliced:

from typing import Union, List

class CustomContainer:
    def __init__(self, values: List[int]):
        self.values = values

    def __getitem__(self, item: Union[int, slice]) -> Union[int, CustomContainer]:
        if isinstance(item, slice):
            return CustomContainer(self.values[item])
        return self.values[item]

This works but comes with the following problem:

a = CustomContainer([1, 2])
b = a[0]  # is always int, but recognized as both int and CustomContainer
c = a[:]  # is always CustomContainer, but recognized as both int and CustomContainer

# Non-scalable solution: Forced type hint
d: int = a[0]
e: CustomContainer = a[:]

If I change the return type of __getitem__ to only int (my original approach), then a[0] correctly shows type int, but a[:] is considered a list instead of a CustomContainer. As far as I understand, there used to be a function in python2 to define how slices are created, but it was removed in python3.

Is there a way to give the proper type hint without having to force the type hint every time I use my container?

Answer

You want to use typing.overload, which allows you to register multiple different signatures of a function with a type checker. Functions decorated with @overload are ignored at runtime, so you’ll typically just fill the body with a literal ellipsis ..., pass, or a docstring. This also means that you have to keep at least one version of the function that isn’t decorated with @overload, which will be the actual function used at runtime.

If you take a look at typeshed, the repository of stub files used by most major type-checkers for checking the standard library, you’ll see this is the technique they use for annotating __getitem__ methods in custom containers such as collections.UserList. In your case, you’d annotate your method like this:

from typing import overload, Union, List

class CustomContainer:
    def __init__(self, values: List[int]):
        self.values = values
        
    @overload
    def __getitem__(self, item: int) -> int:
        """Signature when the function is passed an int"""
        
    @overload
    def __getitem__(self, item: slice) -> CustomContainer:
        """Signature when the function is passed a slice"""

    def __getitem__(self, item: Union[int, slice]) -> Union[int, CustomContainer]:
        """Actual runtime implementation"""

        if isinstance(item, slice):
            return CustomContainer(self.values[item])
        return self.values[item]

a = CustomContainer([1, 2])
b = a[0]
c = a[:]

reveal_type(b)
reveal_type(c)

Run it through MyPy, and it tells us:

main.py:24: note: Revealed type is "builtins.int"
main.py:25: note: Revealed type is "__main__.CustomContainer"

Further reading

The mypy docs for @overload can be found here.