Type hint TypeVars inside @overload decorators with *args

[JackAshwell11]

Hi, so this is a simplified version of my current problem. The code:

from typing import TypeVar, overload


class Base:
    base_var: int


class SubOne(Base):
    component_one_var: float


class SubTwo(Base):
    component_two_var: str


class SubThree(Base):
    component_three_var: bool


T = TypeVar("T", bound=Base)
T1 = TypeVar("T1", bound=Base)
T2 = TypeVar("T2", bound=Base)


@overload
def foo(__a_one: type[T]) -> tuple[int, tuple[T]]:
    ...


@overload
def foo(__a_one: type[T], __a_two: type[T1]) -> tuple[int, tuple[T, T1]]:
    ...


@overload
def foo(
    __a_one: type[T], __a_two: type[T1], __a_three: type[T2]
) -> tuple[int, tuple[T, T1, T2]]:
    ...


def foo(*args: type[Base]) -> tuple[int, tuple[Base, ...]]:
    return tuple(1, tuple(arg() for arg in args))


for res_id, res_one in foo(SubOne):
    print(res_id, res_one)  # Type hinting works here


for res_id, res_two in foo(SubOne, SubTwo):
    print(res_id, res_two)  # Type hinting doesn't works here


for res_id, res_three in foo(SubOne, SubTwo, SubThree):
    print(res_id, res_three)  # Type hinting doesn't works here

The foo() function is meant to take in an arbitrary number of parameters. The function then returns a collection of instances which match those parameter types so that each return value in the tuple is correctly typed to the specific class. However, this works for having one parameter, but doesn't work for having multiple parameters. Is there any way to fix this?

[erictraut]

Is foo supposed to return an Iterable of tuples? Based on how you're using it in the for loops, it looks like that's your intent, but the return type annotation for foo indicates that it returns a single tuple, not an Iterable of tuples.

This type checks:

@overload
def foo(__a_one: type[T]) -> Iterable[tuple[int, tuple[T]]]:
    ...

@overload
def foo(__a_one: type[T], __a_two: type[T1]) -> Iterable[tuple[int, tuple[T, T1]]]:
    ...

@overload
def foo(
    __a_one: type[T], __a_two: type[T1], __a_three: type[T2]
) -> Iterable[tuple[int, tuple[T, T1, T2]]]:
    ...

def foo(*args: type[Base]) -> Iterable[tuple[int, tuple[Base, ...]]]:
    return [(1, tuple(arg() for arg in args))]

[JackAshwell11]

Hi, so I did post a slightly incorrect snippet of code, but your advice seems to have put me on the correct path. The only problem now is that the sub_one variables in the second and third for loop seem to have the type of tuple[SubOne] when they should have the type of SubOne:

from typing import Iterable, TypeVar, overload


class Base:
    base_var: int


class SubOne(Base):
    component_one_var: float


class SubTwo(Base):
    component_two_var: str


class SubThree(Base):
    component_three_var: bool


T = TypeVar("T", bound=Base)
T1 = TypeVar("T1", bound=Base)
T2 = TypeVar("T2", bound=Base)


registry: dict[int, dict[type[Base], Base]] = {
    1: {SubOne: SubOne(), SubTwo: SubTwo(), SubThree: SubThree()},
    2: {SubOne: SubOne(), SubTwo: SubTwo()},
    3: {SubOne: SubOne()},
    4: {},
    5: {SubOne: SubOne(), SubThree: SubThree()},
    6: {SubTwo: SubTwo(), SubThree: SubThree()},
}


@overload
def foo(__a_one: type[T]) -> Iterable[tuple[int, tuple[T]]]:
    ...


@overload
def foo(__a_one: type[T], __a_two: type[T1]) -> Iterable[tuple[int, tuple[T, T1]]]:
    ...


@overload
def foo(
    __a_one: type[T],
    __a_two: type[T1],
    __a_three: type[T2],
) -> Iterable[tuple[int, tuple[T, T1, T2]]]:
    ...


def foo(*args: type[Base]) -> Iterable[tuple[int, tuple[Base, ...]]]:
    ids = [
        res_id
        for res_id, components in registry.items()
        if all(arg in components for arg in args)
    ]

    return [
        (
            object_id,
            tuple(registry[object_id][arg] for arg in args),
        )
        for object_id in ids
    ]


for res_id, sub_one in foo(SubOne):
    print(res_id, sub_one)


for res_id, (sub_one, sub_two) in foo(SubOne, SubTwo):
    print(res_id, sub_one, sub_two)


for res_id, (sub_one, sub_two, sub_three) in foo(SubOne, SubTwo, SubThree):
    print(res_id, sub_one, sub_two, sub_three)

[erictraut]

Your first overload signature has a return type of Iterable[tuple[int, tuple[T]]], but the implementation returns Iterable[tuple[int, T]] in this case. I'm not sure whether the implementation or the overload is incorrect, but they don't match.

I presume that you're using mypy. Mypy treats the first assignment of a variable as an implicit type declaration and then enforces that declaration for all other assignments. Since the first assignment of sub_one is assigning a value of type tuple[type[SubOne]], it assumes that all subsequent assignments must also conform to this type. If you use pyright as your type checker, it doesn't assume that the first assignment is an implicit type declaration, so your code type checks without errors.

[JackAshwell11]

Is there a way I can fix that error with mypy?

[erictraut]

The answer depends on how you decide to fix the bug in your code. The most straightforward way is to change your implementation of foo so it returns a type that's consistent with the overload definition above. Then adjust the first loop so it unpacks the one-entry tuple.

for res_id, (sub_one, ) in foo(SubOne):
    print(res_id, sub_one)

If you want to keep your implementation of foo as it is, then you'll need to change the first overload so it matches the implementation (i.e. change its return type to Iterable[tuple[int, T]]. Then you'll need to explicitly declare that sub_one can accept either a tuple[SubOne] or a SubOne:

sub_one: tuple[SubOne] | SubOne

[JackAshwell11]

Seems like just changing the loop to for res_id, (sub_one, ) in foo(SubOne): without chaning foo() declaration fixes all the mypy and pyright errors. It now successfully works so thanks a ton!