Distribute return type over union

[randolf-scholz]

Inspired by microsoft/pyright#10673, I looked at this simple example: Code sample in pyright playground

from collections.abc import Sequence

def listify[T](x: Sequence[T]) -> list[T]:
    return list(x)

def func(x: Sequence[str] | Sequence[int]):
    y = listify(x)
    reveal_type(y)  # list[str | int]
    print(y)

It feels like the revealed type is not actually what we want here. The desirable inference is imo list[str] | list[int].

Is there simple way to annotate listify to make the return type distribute over unions, or is this a general type-checker limitation?

What seems to happen here is that we assign Sequence[str] | Sequence[int] <: Sequence[str | int] = Sequence[T], so we pick T = int | str. However, in principle wouldn't it be more precise if the type checker, when seeing a UnionType as an argument, passes every member separately and yields the UnionType of the individual return types?

I guess with multiple arguments this runs into combinatoric explosion...

[ryux1]

Your intuition is supported by the current overload-call algorithm, but the distribution is specific to overload resolution. An ordinary generic call solves T against the argument as a whole, so here it can solve T = str | int and the result is consequently list[str | int].

For this finite set of cases, overloads express the correlation:

from collections.abc import Sequence
from typing import overload

@overload
def listify(x: Sequence[str], /) -> list[str]: ...
@overload
def listify(x: Sequence[int], /) -> list[int]: ...

def listify[T](x: Sequence[T], /) -> list[T]:
    return list(x)

With x: Sequence[str] | Sequence[int], neither overload initially accepts the union as a whole. The typing spec's overload algorithm then performs argument type expansion, checks each union member separately, and unions the successful return types. The result should therefore be list[str] | list[int].

There isn't currently a general annotation that says “map this type constructor distributively over every member of an arbitrary union”. If listify also needs to be publicly generic for types beyond these finite cases, adding an unrestricted generic overload as a fallback defeats this particular trick: that overload can accept the original union directly and infer T = str | int. In that case you need either explicit overloads for the domain where the distinction matters, a separate unrestricted generic function, or narrowing at the call site.

And yes, the potential combinatorial growth is handled explicitly in the overload algorithm: expandable arguments are considered one at a time, and their successful return types are combined.