Python typing, mypy infer return type based on class method return type

Question

Consider the case when I have different classes implementing the same method while returning different types.

class A:
    def method(self) -> float:
        return 3.14

class B:
    def method(self) -> str:
        return 'a string'

def do_method(x):
    return x.method()

r = do_method(A())
reveal_type(r)  # Revealed type is 'Any'

Mypy is not being able to infer the exact return type of function do_method() which depends on its argument x . How can I help Mypy achieve this?

Note: Please also consider that number of such classes that I want to use with the function do_method() is too many, so one doesn't want to change them all.

Answer 1

You could use generic protocol to do what you need. But it should be minded that mypy requires the covariance of the return type of the function of protocol when it is a TypeVar , so we must explicitly state this by covariant=True , otherwise the variable is considered as an invariant by default.

A covariant return type of a method is one that can be replaced by a "narrower" type when the method is overridden in a subclass.

from typing import TypeVar, Protocol

T = TypeVar('T', covariant=True)

class Proto(Protocol[T]):
    def method(self) -> T: ...
    
class A:
    def method(self) -> float:
        return 3.14
        
class B:
    def method(self) -> str:
        return 'a string'
        
def do_method(x: Proto[T]) -> T:
    return x.method()
    
r1 = do_method(A())
reveal_type(r1)  # Revealed type is 'builtins.float*'
r2 = do_method(B())
reveal_type(r2)  # Revealed type is 'builtins.str*'