Adding overloads in an overridden subclass method

I have been trying to use overload so that type checkers can infer the return type of a method that is overloaded and is also overridden in a subclass but it doesn’t seem possible without type: ignore at least with mypy and pyright. This is a minimal example:

from __future__ import annotations
from typing import overload

class A:
    @overload
    def method(self, arg1: A) -> A: ...
    @overload
    def method(self, arg1: list[A]) -> list[A]: ...

    def method(self, arg1: A | list[A]) -> A | list[A]:
        assert False

class B(A):
    @overload # type: ignore
    def method(self, arg1: B) -> B: ...
    @overload
    def method(self, arg1: A) -> A: ...
    @overload
    def method(self, arg1: list[A]) -> list[A]: ...

    def method(self, arg1: A | list[A]) -> A | list[A]: # pyright: ignore
        assert False

reveal_type(B().method(B())) # B

My primary goal here is to get reveal_type to infer the type correctly as shown at the bottom which seems to work with both mypy and pyright. The second two overloads listed in class B are identical to the ones listed in class A as is the signature of the actual methods. The only difference is that in class B there is an additional more specific overload saying that if B.method is called with an instance of B rather than A then there will be a more specific return type B rather than A. This more specific overload is what is needed for reveal_type to work for the example shown.

Neither mypy nor pyright accepts the code as is and so two type: ignore are added, one for each checker, at different lines. Without those I get:

$ mypy p.py 
p.py:14: error: Signature of "method" incompatible with supertype "A"  [override]
p.py:14: note:      Superclass:
p.py:14: note:          @overload
p.py:14: note:          def method(self, arg1: A) -> A
p.py:14: note:          @overload
p.py:14: note:          def method(self, arg1: list[A]) -> list[A]
p.py:14: note:      Subclass:
p.py:14: note:          @overload
p.py:14: note:          def method(self, arg1: B) -> B
p.py:14: note:          @overload
p.py:14: note:          def method(self, arg1: A) -> A
p.py:14: note:          @overload
p.py:14: note:          def method(self, arg1: list[A]) -> list[A]
p.py:24: note: Revealed type is "p.B"
Found 1 error in 1 file (checked 1 source file)

$ pyright p.py
WARNING: there is a new pyright version available (v1.1.388 -> v1.1.390).
Please install the new version or set PYRIGHT_PYTHON_FORCE_VERSION to `latest`

p.py
  p.py:21:9 - error: Method "method" overrides class "A" in an incompatible manner
    Override does not handle all overloads of base method (reportIncompatibleMethodOverride)
  p.py:24:13 - information: Type of "B().method(B())" is "B"
1 error, 0 warnings, 1 information 

As far as I can tell there is no way to have different overloads in a subclass with either mypy or pyright without using type: ignore. If the parent class A has overloads for method then any subclass that overrides method must have exactly the same list of overloads (which must be explicitly repeated) or both mypy and pyright will reject it. The message from pyright Override does not handle all overloads of base method and mypy’s Signature of "method" incompatible with supertype "A" both seem incorrect to me.

Maybe I misunderstand something but I don’t think that there is anything incorrect about the additional overload: it is just a more specific version of one of the existing overloads. Both checkers seem to be able to understand it as far as reveal_type is concerned but they report errors anyway.

Is there some better way to write this or do I just need type: ignore?

When overriding a method, the type of the override must be assignable to the type of the original method. For detailed assignability rules for callables, refer to this section of the typing spec. In particular, this subsection discusses assignability of overloaded callables.

One way to make this type check without error is to use a method-scoped type variable in your overload definition.

class B(A):
    @overload
    def method[T: A](self, arg1: T) -> T: ...
    ...
2 Likes

Thanks. That works for the two classes I showed but it doesn’t work as I want it to when there are more classes. The problem with this approach and also with using Self is that it applies recursively to all subclasses. If I have class C(B): pass then C().method(C()) should be of type B rather than C:

from __future__ import annotations
from typing import overload, TypeVar

T = TypeVar('T', bound='A')

class A:
    @overload
    def method(self, arg1: T) -> T: ...
    @overload
    def method(self, arg1: list[A]) -> list[A]: ...

    def method(self, arg1: T | list[A]) -> T | list[A]:
        assert False

class B(A):
    @overload
    def method(self, arg1: T) -> T: ...
    @overload
    def method(self, arg1: list[A]) -> list[A]: ...

    def method(self, arg1: T | list[A]) -> T | list[A]:
        assert False

class C(B):
    pass

reveal_type(B().method(B())) # B
reveal_type(C().method(C())) # C (but should be B)

Is the override in class B not assignable to the one in class A? The spec says:

If a callable B is overloaded with two or more signatures, it is assignable to callable A if at least one of the overloaded signatures in B is assignable to A

And:

If a callable A is overloaded with two or more signatures, callable B is assignable to A if B is assignable to all of the signatures in A :

The overloads in B.method are an exact superset of the overloads in A.method so B.method should be assignable to A.method.