Why pyright can't correctly infer the return type of overload here

[maple3142]

In the following code, when RET0 = int and RET1 = str, pyright can correctly infer that g(0) is int and g(1) is str.

But if RET0 = Optional[int] and RET1 = Optional[str], then the return type of both g(0) and g(1) becomes int | str | None, so the optional types somehow got flattened.

from typing import Literal, Union, overload, TypeVar, Optional, Any


RET0 = int
RET1 = str

# RET0 = Optional[int]
# RET1 = Optional[str]


@overload
def f(arg: Literal[0]) -> RET0: ...
@overload
def f(arg: Literal[1]) -> RET1: ...


def f(arg: Union[Literal[0], Literal[1]]) -> Union[RET0, RET1]:
    if arg == 0:
        return 0
    return "1"


T = TypeVar("T", Literal[0], Literal[1])


def g(arg: T):
    return f(arg)


reveal_type(f)
reveal_type(f(0))
reveal_type(f(1))

reveal_type(g)
reveal_type(g(0))  # should be RET0
reveal_type(g(1))  # should be RET1

[erictraut]

You're attempting to use a value-constrained type variable here. And you're constraining it to literal values, no less. I recommend against using value-constrained type variables in your code. Their behaviors are not well defined in the typing spec, and you will see different assumptions and behaviors across type checkers. Value-constrained type variables are unique to the Python type system. You won't find them in any other type system for good reason.

I recommend switching away from value-constrained type variables. An overload is the appropriate and recommended approach in this case.

[maple3142]

Thanks for the reply! So in Python type hints, the best way two write a wrapper function (g) for a overloaded function (f), whose return type depends on the argument literals, is to define all the overloads manually?