Changing type argument of Python with subclassing

Question

Python's typing system allows for generics in classes:

class A(Generic[T]):
    def get_next(self) -> T

which is very handy. However, even in 3.11 with the Self type, I cannot find a way to change the type argument (the T) without specifying the class name. Here's the recommended usage from PEP 673: Self Type: https://peps.python.org/pep-0673/a

class Container(Generic[T]):
    def foo(
        self: Container[T],
    ) -> Container[str]:
        # maybe implementing something like:
        return self.__class__([str(x) for x in self])

The problem is if I want to subclass container:

class SuperContainer(Container[T]):
    def time_travel(self): ...

And then if I have an instance of SuperContainer and call foo on it, the typing will be wrong, and think that it's a Container not SuperContainer.

sc = SuperContainer([1, 2, 3])
sc2 = sc.foo()
reveal_type(sc2)  # mypy: Container[str]
sc2.time_travel()  # typing error: only SuperContainers can time-travel
isinstance(sc2, SuperContainer)  # True

Is there an accepted way to allow a program to change the type argument in the superclass that preserves the typing of the subclass?

Answer 1

I have no clue how this works, but I made it work with Type[T] too. I genuinely cannot explain this code so I'm just gonna copy and paste, better people than me can tell you why.

from typing import TypeVar, Generic, Any, TypeAlias, TYPE_CHECKING, Type

if not TYPE_CHECKING:
    reveal_type = print

T = TypeVar('T')
SelfStr = TypeVar("SelfStr", bound="Container[str, Any, Any]", covariant=True)
SelfTypeT = TypeVar("SelfTypeT", bound="Container[Type[Any], Any, Any]", covariant=True)

class Container(Generic[T, SelfStr, SelfTypeT]):
    def __init__(self, contents: list[T]):
        self._contents = contents

    def __iter__(self):
        return iter(self._contents)

    def foo(self) -> SelfStr:
        reveal_type(type(self))
        # Mypy is wrong here: it thinks that type(self) is already annotated, but in fact the type parameters are erased.
        return type(self)([str(x) for x in self])  # type: ignore

    def get_types(self) -> SelfTypeT:
        return type(self)([type(x) for x in self])  # type: ignore

    def __repr__(self):
        return type(self).__name__ + "(" + repr(self._contents) + ")"
_ContainerStr: TypeAlias = Container[str, "_ContainerStr", "ContainerComplete[Type[str]]"]
_ContainerTypeT: TypeAlias = Container[Type[T], "_ContainerStr", "_ContainerTypeT[Type[type]]"]
ContainerComplete: TypeAlias = Container[T, _ContainerStr, _ContainerTypeT[T]]

class SuperContainer(Container[T, SelfStr, SelfTypeT]):
    def time_travel(self):
        return "magic"
_SuperContainerStr: TypeAlias = SuperContainer[str, "_SuperContainerStr", "SuperContainerComplete[Type[str]]"]
_SuperContainerTypeT: TypeAlias = SuperContainer[Type[T], "_SuperContainerStr", "_SuperContainerTypeT[Type[type]]"]
SuperContainerComplete: TypeAlias = SuperContainer[T, _SuperContainerStr, _SuperContainerTypeT[T]]

sc = SuperContainerComplete[int]([3, 4, 5])
reveal_type(sc)

sc2 = sc.foo()
reveal_type(sc2)

sc3 = sc.get_types()
reveal_type(sc3)

class Base:
    pass

class Impl1(Base):
    pass

class Impl2(Base):
    pass

sc4 = SuperContainerComplete[Base]([Impl1(), Impl2()])

sc5 = sc4.foo()
reveal_type(sc5)

sc6 = sc4.get_types()
reveal_type(sc6)

print(sc2.time_travel())

Mypy and Python are both happy with this code so I guess It Works (TM).