CODEWITHSUNDEEP
pythontype-hintingmypy
Mario Ishac
Mario Ishac

Reputation: 5877

How to type hint function accepting both union and any type of that union?

$ cat foo.py

from typing import overload, Union, TypeVar

T = Union[int, str]
SubT = TypeVar("SubT", int, str)

@overload
def a(t: T) -> T:
    ...

@overload
def a(t: SubT) -> SubT:
    ...

def a(t: T) -> T:
    return t

$ mypy foo.py

foo.py:7: error: Overloaded function signatures 1 and 2 overlap with incompatible return types

Why are the return types incompatible, and how can I make this typecheck? I want these example return types:

v_1: Union[int, str] = 1
v_2: int = 2
v_3: str = "3"

a(v_1)  # Want Union[int, str].
a(v_2)  # Want int.
a(v_3)  # Want str.

I want to avoid having to explicitly overload a for int and str, as SubT has many more than 2 constraints in reality.

If I remove the first overload, a(v_1) wouldn't type check. If I remove the second overload, a(v_2) and a(v_3) would have their return value type hinted as Union[int, str] instead of int and str respectively.

In my actual problem, I have an Iterable of either homogeneous types under SubT (the TypeVar) or heterogeneous types under T (the Union). I want to write an a that will operate on the elements of both without losing type granulariy if it present in the homogeneous case.

Upvotes: 1

Views: 900

Answers (1)

Ry-
Ry-

Reputation: 224862

TypeVar will let you offer both the types of the union and the union itself as options:

T = TypeVar("T", int, str, Union[int, str])

def a(t: T) -> T:
    return t

Upvotes: 1

Related Questions