How do I declare return type based on argument type?

user_1015595 0
Viewed 859

This question is similar to Python Typing: declare return value type based on function argument but different enough that it doesn't fit into a comment.

I have the following function:

T = TypeVar('T', dict, str)

def fun(t: T) -> T:
    if t == dict:
        return {"foo": "bar"}
    else:
        return "foo"

I'd like to be able to call it like this:

a_string = fun(str)
a_dict = fun(dict)

Pylance throws this error on the second line:

Expression of type "dict[str, str]" cannot be assigned to return type "T@fun"

And this error on the last line:

Expression of type "Literal['foo']" cannot be assigned to return type "T@fun"

According to this answer, I should be able to do something like this:

T = TypeVar('T', dict, str)

def fun(t: Type[T]) -> T:
    if t == dict:
        return t({"foo": "bar"})
    else:
        return t("foo")

This removes the error on the second line but causes a different error on the last line:

No overloads for "__init__" match the provided arguments
  Argument types: (Literal['foo'])

I studied this answer for a long time before I was finally able to make it work:

T = TypeVar('T', dict, str)

def fun(t: Callable[..., T]) -> T:
    if t == dict:
        return t({"foo": "bar"})
    else:
        return t("foo")

The problem with this is I don't understand why it works. And I don't understand why the others don't.

python-typing python
2 Answers

The problem with this is I don't understand why it works

The last one works because a type is a callable. So the typing here says that fun takes something that, given something, will return a T, and you supply that in the form of a type.

And I don't understand why the others don't.

The first version doesn't work because there is no binding between the left and right side of ->. So you can pass a typevar, but you cannot specify that the specific on the left must be the specific on the right. To put it differently, if the argument is a dict, then according to the signature, the return type doesn't have to be a dict, but rather still the types specified by the typevar (it is a bit confusing, because the same symbol - T - appears in both. However, T means here "one of ...").

The second version doesn't work because the return is indicating that you're returning a type as well, as far as the typechecker is concerned.

user_35107360

You could also use typing.overload here.

from typing import TypeVar, Type, overload, Callable


T = TypeVar('T')


@overload
def fun(t: Type[dict]) -> dict:
    ...


@overload
def fun(t: Type[str]) -> str:
    ...


def fun(t: Callable[..., T]) -> T
    if t == dict:
        return t({"foo": "bar"})
    else:
        return t("foo")

This should allow fun(dict) and fun(str), but no calls with a different type, while also ensuring that fun does, indeed, return a value of the type passed as an argument.

user_11268410
Related