[ty] Infer type variables within generic unions

crates/ty_python_semantic/resources/mdtest/external/sqlalchemy.md

@@ -107,44 +107,34 @@ We can also specify particular columns to select:
 
 ```py
 stmt = select(User.id, User.name)
-# TODO: should be `Select[tuple[int, str]]`
-reveal_type(stmt)  # revealed: Select[tuple[Unknown, Unknown]]
+reveal_type(stmt)  # revealed: Select[tuple[int, str]]
 
 ids_and_names = session.execute(stmt).all()
-# TODO: should be `Sequence[Row[tuple[int, str]]]`
-reveal_type(ids_and_names)  # revealed: Sequence[Row[tuple[Unknown, Unknown]]]
+reveal_type(ids_and_names)  # revealed: Sequence[Row[tuple[int, str]]]
 
 for row in session.execute(stmt):
-    # TODO: should be `Row[tuple[int, str]]`
-    reveal_type(row)  # revealed: Row[tuple[Unknown, Unknown]]
+    reveal_type(row)  # revealed: Row[tuple[int, str]]
 
 for user_id, name in session.execute(stmt).tuples():
-    # TODO: should be `int`
-    reveal_type(user_id)  # revealed: Unknown
-    # TODO: should be `str`
-    reveal_type(name)  # revealed: Unknown
+    reveal_type(user_id)  # revealed: int
+    reveal_type(name)  # revealed: str
 
 result = session.execute(stmt)
 row = result.one_or_none()
 assert row is not None
 (user_id, name) = row._tuple()
-# TODO: should be `int`
-reveal_type(user_id)  # revealed: Unknown
-# TODO: should be `str`
-reveal_type(name)  # revealed: Unknown
+reveal_type(user_id)  # revealed: int
+reveal_type(name)  # revealed: str
 
 stmt = select(User.id).where(User.name == "Alice")
 
-# TODO: should be `Select[tuple[int]]`
-reveal_type(stmt)  # revealed: Select[tuple[Unknown]]
+reveal_type(stmt)  # revealed: Select[tuple[int]]
 
 alice_id = session.scalars(stmt).first()
-# TODO: should be `int | None`
-reveal_type(alice_id)  # revealed: Unknown | None
+reveal_type(alice_id)  # revealed: int | None
 
 alice_id = session.scalar(stmt)
-# TODO: should be `int | None`
-reveal_type(alice_id)  # revealed: Unknown | None
+reveal_type(alice_id)  # revealed: int | None
 ```
 
 Using the legacy `query` API also works:
@@ -166,15 +156,12 @@ And similarly when specifying particular columns:
 
 ```py
 query = session.query(User.id, User.name)
-# TODO: should be `RowReturningQuery[tuple[int, str]]`
-reveal_type(query)  # revealed: RowReturningQuery[tuple[Unknown, Unknown]]
+reveal_type(query)  # revealed: RowReturningQuery[tuple[int, str]]
 
-# TODO: should be `list[Row[tuple[int, str]]]`
-reveal_type(query.all())  # revealed: list[Row[tuple[Unknown, Unknown]]]
+reveal_type(query.all())  # revealed: list[Row[tuple[int, str]]]
 
 for row in query:
-    # TODO: should be `Row[tuple[int, str]]`
-    reveal_type(row)  # revealed: Row[tuple[Unknown, Unknown]]
+    reveal_type(row)  # revealed: Row[tuple[int, str]]
 ```
 
 ## Async API
@@ -203,8 +190,6 @@ async def test_async(session: AsyncSession):
     stmt = select(User.id, User.name)
     result = await session.execute(stmt)
     for user_id, name in result.tuples():
-        # TODO: should be `int`
-        reveal_type(user_id)  # revealed: Unknown
-        # TODO: should be `str`
-        reveal_type(name)  # revealed: Unknown
+        reveal_type(user_id)  # revealed: int
+        reveal_type(name)  # revealed: str
 ```

crates/ty_python_semantic/resources/mdtest/generics/legacy/functions.md

@@ -347,6 +347,138 @@ reveal_type(tuple_param("a", ("a", 1)))  # revealed: tuple[Literal["a"], Literal
 reveal_type(tuple_param(1, ("a", 1)))  # revealed: tuple[Literal["a"], Literal[1]]
 ```
 
+When a union parameter contains generic classes like `P[T] | Q[T]`, we can infer the typevar from
+the actual argument even for non-final classes.
+
+```py
+from typing import TypeVar, Generic
+
+T = TypeVar("T")
+
+class P(Generic[T]):
+    x: T
+
+class Q(Generic[T]):
+    x: T
+
+def extract_t(x: P[T] | Q[T]) -> T:
+    raise NotImplementedError
+
+reveal_type(extract_t(P[int]()))  # revealed: int
+reveal_type(extract_t(Q[str]()))  # revealed: str
+```
+
+Passing anything else results in an error:
+
+```py
+# error: [invalid-argument-type]
+reveal_type(extract_t([1, 2]))  # revealed: Unknown
+```
+
+This also works when different union elements have different typevars:
+
+```py
+S = TypeVar("S")
+
+def extract_both(x: P[T] | Q[S]) -> tuple[T, S]:
+    raise NotImplementedError
+
+reveal_type(extract_both(P[int]()))  # revealed: tuple[int, Unknown]
+reveal_type(extract_both(Q[str]()))  # revealed: tuple[Unknown, str]
+```
+
+Inference also works when passing subclasses of the generic classes in the union.
+
+```py
+class SubP(P[T]):
+    pass
+
+class SubQ(Q[T]):
+    pass
+
+reveal_type(extract_t(SubP[int]()))  # revealed: int
+reveal_type(extract_t(SubQ[str]()))  # revealed: str
+
+reveal_type(extract_both(SubP[int]()))  # revealed: tuple[int, Unknown]
+reveal_type(extract_both(SubQ[str]()))  # revealed: tuple[Unknown, str]
+```
+
+When a type is a subclass of both `P` and `Q` with different specializations, we cannot infer a
+single type for `T` in `extract_t`, because `P` and `Q` are invariant. However, we can still infer
+both types in a call to `extract_both`:
+
+```py
+class PandQ(P[int], Q[str]):
+    pass
+
+# TODO: Ideally, we would return `Unknown` here.
+# error: [invalid-argument-type]
+reveal_type(extract_t(PandQ()))  # revealed: int | str
+
+reveal_type(extract_both(PandQ()))  # revealed: tuple[int, str]
+```
+
+When non-generic types are part of the union, we can still infer typevars for the remaining generic
+types:
+
+```py
+def extract_optional_t(x: None | P[T]) -> T:
+    raise NotImplementedError
+
+reveal_type(extract_optional_t(None))  # revealed: Unknown
+reveal_type(extract_optional_t(P[int]()))  # revealed: int
+```
+
+Passing anything else results in an error:
+
+```py
+# error: [invalid-argument-type]
+reveal_type(extract_optional_t(Q[str]()))  # revealed: Unknown
+```
+
+If the union contains contains parent and child of a generic class, we ideally pick the union
+element that is more precise:
+
+```py
+class Base(Generic[T]):
+    x: T
+
+class Sub(Base[T]): ...
+
+def f(t: Base[T] | Sub[T | None]) -> T:
+    raise NotImplementedError
+
+reveal_type(f(Base[int]()))  # revealed: int
+# TODO: Should ideally be `str`
+reveal_type(f(Sub[str | None]()))  # revealed: str | None
+```
+
+If we have a case like the following, where only one of the union elements matches due to the
+typevar bound, we do not emit a specialization error:
+
+```py
+from typing import TypeVar
+
+I_int = TypeVar("I_int", bound=int)
+S_str = TypeVar("S_str", bound=str)
+
+class P(Generic[T]):
+    value: T
+
+def f(t: P[I_int] | P[S_str]) -> tuple[I_int, S_str]:
+    raise NotImplementedError
+
+reveal_type(f(P[int]()))  # revealed: tuple[int, Unknown]
+reveal_type(f(P[str]()))  # revealed: tuple[Unknown, str]
+```
+
+However, if we pass something that does not match _any_ union element, we do emit an error:
+
+```py
+# error: [invalid-argument-type]
+reveal_type(f(P[bytes]()))  # revealed: tuple[Unknown, Unknown]
+```
+
 ## Inferring nested generic function calls
 
 We can infer type assignments in nested calls to multiple generic functions. If they use the same

crates/ty_python_semantic/resources/mdtest/generics/pep695/functions.md

@@ -310,6 +310,127 @@ reveal_type(tuple_param("a", ("a", 1)))  # revealed: tuple[Literal["a"], Literal
 reveal_type(tuple_param(1, ("a", 1)))  # revealed: tuple[Literal["a"], Literal[1]]
 ```
 
+When a union parameter contains generic classes like `P[T] | Q[T]`, we can infer the typevar from
+the actual argument even for non-final classes.
+
+```py
+class P[T]:
+    x: T  # invariant
+
+class Q[T]:
+    x: T  # invariant
+
+def extract_t[T](x: P[T] | Q[T]) -> T:
+    raise NotImplementedError
+
+reveal_type(extract_t(P[int]()))  # revealed: int
+reveal_type(extract_t(Q[str]()))  # revealed: str
+```
+
+Passing anything else results in an error:
+
+```py
+# error: [invalid-argument-type]
+reveal_type(extract_t([1, 2]))  # revealed: Unknown
+```
+
+This also works when different union elements have different typevars:
+
+```py
+def extract_both[T, S](x: P[T] | Q[S]) -> tuple[T, S]:
+    raise NotImplementedError
+
+reveal_type(extract_both(P[int]()))  # revealed: tuple[int, Unknown]
+reveal_type(extract_both(Q[str]()))  # revealed: tuple[Unknown, str]
+```
+
+Inference also works when passing subclasses of the generic classes in the union.
+
+```py
+class SubP[T](P[T]):
+    pass
+
+class SubQ[T](Q[T]):
+    pass
+
+reveal_type(extract_t(SubP[int]()))  # revealed: int
+reveal_type(extract_t(SubQ[str]()))  # revealed: str
+
+reveal_type(extract_both(SubP[int]()))  # revealed: tuple[int, Unknown]
+reveal_type(extract_both(SubQ[str]()))  # revealed: tuple[Unknown, str]
+```
+
+When a type is a subclass of both `P` and `Q` with different specializations, we cannot infer a
+single type for `T` in `extract_t`, because `P` and `Q` are invariant. However, we can still infer
+both types in a call to `extract_both`:
+
+```py
+class PandQ(P[int], Q[str]):
+    pass
+
+# TODO: Ideally, we would return `Unknown` here.
+# error: [invalid-argument-type]
+reveal_type(extract_t(PandQ()))  # revealed: int | str
+
+reveal_type(extract_both(PandQ()))  # revealed: tuple[int, str]
+```
+
+When non-generic types are part of the union, we can still infer typevars for the remaining generic
+types:
+
+```py
+def extract_optional_t[T](x: None | P[T]) -> T:
+    raise NotImplementedError
+
+reveal_type(extract_optional_t(None))  # revealed: Unknown
+reveal_type(extract_optional_t(P[int]()))  # revealed: int
+```
+
+Passing anything else results in an error:
+
+```py
+# error: [invalid-argument-type]
+reveal_type(extract_optional_t(Q[str]()))  # revealed: Unknown
+```
+
+If the union contains contains parent and child of a generic class, we ideally pick the union
+element that is more precise:
+
+```py
+class Base[T]:
+    x: T
+
+class Sub[T](Base[T]): ...
+
+def f[T](t: Base[T] | Sub[T | None]) -> T:
+    raise NotImplementedError
+
+reveal_type(f(Base[int]()))  # revealed: int
+# TODO: Should ideally be `str`
+reveal_type(f(Sub[str | None]()))  # revealed: str | None
+```
+
+If we have a case like the following, where only one of the union elements matches due to the
+typevar bound, we do not emit a specialization error:
+
+```py
+class P[T]:
+    value: T
+
+def f[I: int, S: str](t: P[I] | P[S]) -> tuple[I, S]:
+    raise NotImplementedError
+
+reveal_type(f(P[int]()))  # revealed: tuple[int, Unknown]
+reveal_type(f(P[str]()))  # revealed: tuple[Unknown, str]
+```
+
+However, if we pass something that does not match _any_ union element, we do emit an error:
+
+```py
+# error: [invalid-argument-type]
+reveal_type(f(P[bytes]()))  # revealed: tuple[Unknown, Unknown]
+```
+
 ## Inferring nested generic function calls
 
 We can infer type assignments in nested calls to multiple generic functions. If they use the same