Premature variable pinning for union of types in a Sequence

[DouweM]

The following works in pyright, but not in mypy, ty, or (which is why I'm here!) pyrefly:

from __future__ import annotations

from dataclasses import dataclass

from typing_extensions import (
    Generic,
    Sequence,
    TypeVar,
    assert_type,
)

T = TypeVar("T")


@dataclass
class Agent(Generic[T]):
    output_type: Sequence[type[T]]


class Foo:
    pass


class Bar:
    pass


# pyright - works
# mypy - error: Expression is of type "Agent[object]", not "Agent[Foo | Bar]"  [assert-type]
# pyrefly - assert_type(Agent[Foo], Agent[Bar | Foo]) failed + Argument `list[type[Bar] | type[Foo]]` is not assignable to parameter `output_type` with type `Sequence[type[Foo]]` in function `Agent.__init__`
# ty - `Agent[Foo | Bar]` and `Agent[Unknown]` are not equivalent types
assert_type(Agent([Foo, Bar]), Agent[Foo | Bar])

# pyright - works
# mypy - error: Expression is of type "Agent[Never]", not "Agent[int | str]"  [assert-type]
# pyrefly - assert_type(Agent[int], Agent[str | int]) failed + Argument `list[type[str] | type[int]]` is not assignable to parameter `output_type` with type `Sequence[type[int]]` in function `Agent.__init__`
# ty - `Agent[int | str]` and `Agent[Unknown]` are not equivalent types
assert_type(Agent([int, str]), Agent[int | str])

# works
assert_type(Agent[Foo | Bar]([Foo, Bar]), Agent[Foo | Bar])

# works
assert_type(Agent[int | str]([int, str]), Agent[int | str])

It would be great to see this work in pyrefly, but if there's a good reason the other 2 out of 3 typecheckers don't support this I'd love to understand why!

• This is related to a new PydanticAI feature, if you're curious check out Support functions as output_type, as well as lists of functions and other types pydantic/pydantic-ai#1785 (comment)
• Issues for other type checkers:
  • Sequence of types python/mypy#19142
  • Handle type[T] for type variables astral-sh/ty#501

[yangdanny97]

Weird, we're somehow pinning the type of the type var to just Foo when we call the constructor.

The contextual typing works if it's a regular assignment, like x: Sequence[type[Foo | Bar]] = [Foo, Bar]

[yangdanny97]

Oh, I have a theory, we might be breaking up the union and doing is_subset_eq on the union members one by one, pinning the var to the first union member it encounters.

cc @ndmitchell dealing with the side effects here seems tricky

[ndmitchell]

I think we do break up the union and pin it that way. I think for "x <: A | B" that's the wrong thing to do and we should pin Var before descending into the union.

[samwgoldman]

Here's a reduced example:

from typing import Sequence

def f[T](ts: Sequence[type[T]]) -> T: ...

class A: ...
class B: ...

f([A, B])

And here's what's happening in the solver:

list[type[A] | type[B]] <: Sequence[type[T]]
    type[A] | type[B] <: type[T]
        type[A] <: type[T]
            A <: T  ==>  T = A
        type[B] <: type[T]
            B <: T
                B <: A -- error

So this is arguably a manifestation of #105, but it's worth noting that we can side-step that issue for the specific case in this one, because we are allowed to make a simplification from type[A] | type[B] to type[A | B], which would change the derivation to:

list[type[A | B]] <: Sequence[type[T]]
    type[A | B] <: type[T]
        A | B <: T  ==>  T = A | B