How to type result based on constructor parameters?

[vnmabus]

I am trying to type code like this:

class A:

    def __init__(self, return_tuple=False):
        self.return_tuple = return_tuple
    
    def __call__(self):
        if self.return_tuple:
            return (1, 2)
        return 1

I have tried to make A generic and use overloads on the self parameter in several ways but none seems to work:

class A(Generic[T]):

    def __init__(self, return_tuple: T = False) -> None:
        self.return_tuple = return_tuple

    @overload
    def __call__(self: A[Literal[True]]) -> Tuple[int, int]:
        pass

    @overload
    def __call__(self: A[Literal[False]]) -> int
        pass
    
    def __call__(self) -> Union[int, Tuple[int, int]]:
        if self.return_tuple:
            return (1, 2)
        return 1

where T is either

T = TypeVar("T", bool)

or

T = TypeVar("T", Literal[True], Literal[False])

The first case mark errors in the definition of the TypeVar, the Generic and the optional parameter in __init__. The second one marks an error only in the optional parameter (even if it is essentially the same type). reveal_type does not have the expected result in both cases.

Am I doing something wrong or is this case just not currently taken into consideration?

[freundTech]

This isn't something you can do with mypy.

Mypy is a static type checker. The type you want to describe is inherently dynamic. It depends on the value of a attribute/parameter at runtime. Mypy can only check types with are known at "compile time".

[JelleZijlstra]

I do believe @vnmabus's example should work; Literal types are part of the type system now. Maybe we don't handle them properly in overload resolution though. Here's a self-contained example:

from typing import Generic, TypeVar, Literal, overload, Tuple, Union

T = TypeVar("T", Literal[True], Literal[False])

class A(Generic[T]):
    @overload
    def __call__(self: A[Literal[True]]) -> Tuple[int, int]:
        pass

    @overload
    def __call__(self: A[Literal[False]]) -> int:
        pass
    
    def __call__(self) -> Union[int, Tuple[int, int]]:
        pass

a: A[Literal[False]]
reveal_type(a)  # 'main.A[Literal[False]]'
reveal_type(a())  # 'Tuple[builtins.int, builtins.int]'

The call to a() should match only the second overload, but instead mypy picks the first.

[erictraut]

Here's a solution that works in pyright and almost works in mypy. I'm not sure why, but mypy appears to choose the wrong __call__ overload on the last line.

from typing import Generic, Literal, Tuple, TypeVar, Union, overload

T = TypeVar("T", bound=bool)

class A(Generic[T]):
    @overload
    def __init__(self: A[Literal[False]], return_tuple: Literal[False] = ...) -> None:
        ...

    @overload
    def __init__(self: A[Literal[True]], return_tuple: Literal[True]) -> None:
        ...

    def __init__(self, return_tuple: bool = False) -> None:
        self.return_tuple = return_tuple

    @overload
    def __call__(self: A[Literal[True]]) -> Tuple[int, int]:
        ...

    @overload
    def __call__(self: A[Literal[False]]) -> int:
        ...

    def __call__(self) -> Union[int, Tuple[int, int]]:
        if self.return_tuple:
            return (1, 2)
        return 1

a1 = A(return_tuple=True)
reveal_type(a1)
reveal_type(a1())  # Tuple[int, int]

a2 = A()
reveal_type(a2)
reveal_type(a2())  # int