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

Answer 2

To solve this, you need a second generic type argument, to represent the return type of foo.

SelfStr = TypeVar("SelfStr", bound="Container[str, Any]", covariant=True)

The Any is okay. We'll see that later.

So far so good. Let's define the Container:

class Container(Generic[T, SelfStr]):
    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 __repr__(self):
        return type(self).__name__ + "(" + repr(self._contents) + ")"

Note that we had to ignore the types in foo. This is because mypy has inferred the type of type(self) incorrectly. It thinks that type(self) returns Container[...] (or a subclass), but in fact it returns Container (or a subclass). You'll see that when we get to running this code.

Next, we need some way of creating a container. We want the type to look like Container[T, Container[str, Container[str, ...]]].

In the first line of the class declaration, we made the second type-parameter of the class be a SelfStr, which is itself Container[str, Any]. This means that definition of SelfStr should become bounded to Container[str, SelfStr], so we should get an upper bound of Container[str, Container[str, ...]] as we wanted. This works: it will only allow our recursive type (or subclasses) or Any. Unfortunately, mypy won't do inference on recursive generic types, reporting test.Container[builtins.int, <nothing>], so we have to do the heavy lifting. Time for some ✨ magic ✨.

_ContainerStr: TypeAlias = Container[str, "_ContainerStr"]
ContainerComplete: TypeAlias = Container[T, _ContainerStr]

The _ContainerStr alias will give us the recursive part of the signature. We then expose ContainerComplete, which we can use as a constructor, for example:

ContainerComplete[int]([1,2,3])

Awesome! But what about subclasses? We just have to do the same thing again, for our subclass:

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

All done! Now let's demonstrate:

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

sc2 = sc.foo()
reveal_type(sc2)

print(sc2.time_travel())

Putting everything together, we get:

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

if not TYPE_CHECKING:
    reveal_type = print

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

class Container(Generic[T, SelfStr]):
    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 __repr__(self):
        return type(self).__name__ + "(" + repr(self._contents) + ")"
_ContainerStr: TypeAlias = Container[str, "_ContainerStr"]
ContainerComplete: TypeAlias = Container[T, _ContainerStr]

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

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

sc2 = sc.foo()
reveal_type(sc2)

print(sc2.time_travel())

And the output looks like this (you need a recent version of mypy):

$ mypy test.py
test.py:17: note: Revealed type is "Type[test.Container[T`1, SelfStr`2]]"
test.py:33: note: Revealed type is "test.SuperContainer[builtins.int, test.SuperContainer[builtins.str, ...]]"
test.py:36: note: Revealed type is "test.SuperContainer[builtins.str, test.SuperContainer[builtins.str, ...]]"
Success: no issues found in 1 source file
$ python test.py
<__main__.SuperContainer object at 0x7f30165582d0>
<class '__main__.SuperContainer'>
<__main__.SuperContainer object at 0x7f3016558390>
magic
$

You can remove a lot of the boilerplate using metaclasses. This has the added advantage that it's inherited. If you override __call__, you can even get isinstance working properly (it doesn't work with the generic type aliases *Complete, it still works fine for the classes themselves).

Note that this only partially works in PyCharm: