[ty] support imported PEP 613 type aliases

Author: carljm

crates/ruff_benchmark/benches/ty_walltime.rs

@@ -143,7 +143,7 @@ static FREQTRADE: Benchmark = Benchmark::new(
         max_dep_date: "2025-06-17",
         python_version: PythonVersion::PY312,
     },
-    525,
+    600,
 );
 
 static PANDAS: Benchmark = Benchmark::new(
@@ -163,7 +163,7 @@ static PANDAS: Benchmark = Benchmark::new(
         max_dep_date: "2025-06-17",
         python_version: PythonVersion::PY312,
     },
-    3000,
+    4000,
 );
 
 static PYDANTIC: Benchmark = Benchmark::new(
@@ -181,7 +181,7 @@ static PYDANTIC: Benchmark = Benchmark::new(
         max_dep_date: "2025-06-17",
         python_version: PythonVersion::PY39,
     },
-    5000,
+    7000,
 );
 
 static SYMPY: Benchmark = Benchmark::new(

crates/ty_ide/src/completion.rs

@@ -2136,10 +2136,10 @@ C.<CURSOR>
         );
 
         assert_snapshot!(
-            builder.skip_keywords().skip_builtins().type_signatures().build().snapshot(), @r"
+            builder.skip_keywords().skip_builtins().type_signatures().build().snapshot(), @r###"
         meta_attr :: int
         mro :: bound method <class 'C'>.mro() -> list[type]
-        __annotate__ :: @Todo | None
+        __annotate__ :: (() -> dict[str, Any]) | None
         __annotations__ :: dict[str, Any]
         __base__ :: type | None
         __bases__ :: tuple[type, ...]
@@ -2182,7 +2182,7 @@ C.<CURSOR>
         __text_signature__ :: str | None
         __type_params__ :: tuple[TypeVar | ParamSpec | TypeVarTuple, ...]
         __weakrefoffset__ :: int
-        ");
+        "###);
     }
 
     #[test]
@@ -2331,14 +2331,14 @@ Quux.<CURSOR>
         );
 
         assert_snapshot!(
-            builder.skip_keywords().skip_builtins().type_signatures().build().snapshot(), @r"
+            builder.skip_keywords().skip_builtins().type_signatures().build().snapshot(), @r###"
         mro :: bound method <class 'Quux'>.mro() -> list[type]
         some_attribute :: int
         some_class_method :: bound method <class 'Quux'>.some_class_method() -> int
         some_method :: def some_method(self) -> int
         some_property :: property
         some_static_method :: def some_static_method(self) -> int
-        __annotate__ :: @Todo | None
+        __annotate__ :: (() -> dict[str, Any]) | None
         __annotations__ :: dict[str, Any]
         __base__ :: type | None
         __bases__ :: tuple[type, ...]
@@ -2381,7 +2381,7 @@ Quux.<CURSOR>
         __text_signature__ :: str | None
         __type_params__ :: tuple[TypeVar | ParamSpec | TypeVarTuple, ...]
         __weakrefoffset__ :: int
-        ");
+        "###);
     }
 
     #[test]

crates/ty_python_semantic/resources/mdtest/annotations/unsupported_special_forms.md

@@ -12,11 +12,8 @@ P = ParamSpec("P")
 Ts = TypeVarTuple("Ts")
 R_co = TypeVar("R_co", covariant=True)
 
-Alias: TypeAlias = int
-
 def f(*args: Unpack[Ts]) -> tuple[Unpack[Ts]]:
     reveal_type(args)  # revealed: tuple[@Todo(`Unpack[]` special form), ...]
-    reveal_type(Alias)  # revealed: @Todo(Support for `typing.TypeAlias`)
     return args
 
 def g() -> TypeGuard[int]: ...

crates/ty_python_semantic/resources/mdtest/attributes.md

@@ -2208,9 +2208,9 @@ reveal_type(False.real)  # revealed: Literal[0]
 All attribute access on literal `bytes` types is currently delegated to `builtins.bytes`:
 
 ```py
-# revealed: bound method Literal[b"foo"].join(iterable_of_bytes: Iterable[@Todo(Support for `typing.TypeAlias`)], /) -> bytes
+# revealed: bound method Literal[b"foo"].join(iterable_of_bytes: Iterable[Buffer], /) -> bytes
 reveal_type(b"foo".join)
-# revealed: bound method Literal[b"foo"].endswith(suffix: @Todo(Support for `typing.TypeAlias`) | tuple[@Todo(Support for `typing.TypeAlias`), ...], start: SupportsIndex | None = None, end: SupportsIndex | None = None, /) -> bool
+# revealed: bound method Literal[b"foo"].endswith(suffix: Buffer | tuple[Buffer, ...], start: SupportsIndex | None = None, end: SupportsIndex | None = None, /) -> bool
 reveal_type(b"foo".endswith)
 ```
 

crates/ty_python_semantic/resources/mdtest/binary/instances.md

@@ -313,8 +313,7 @@ reveal_type(A() + "foo")  # revealed: A
 reveal_type("foo" + A())  # revealed: A
 
 reveal_type(A() + b"foo")  # revealed: A
-# TODO should be `A` since `bytes.__add__` doesn't support `A` instances
-reveal_type(b"foo" + A())  # revealed: bytes
+reveal_type(b"foo" + A())  # revealed: A
 
 reveal_type(A() + ())  # revealed: A
 reveal_type(() + A())  # revealed: A

crates/ty_python_semantic/resources/mdtest/binary/integers.md

@@ -54,10 +54,8 @@ reveal_type(2**largest_u32)  # revealed: int
 
 def variable(x: int):
     reveal_type(x**2)  # revealed: int
-    # TODO: should be `Any` (overload 5 on `__pow__`), requires correct overload matching
-    reveal_type(2**x)  # revealed: int
-    # TODO: should be `Any` (overload 5 on `__pow__`), requires correct overload matching
-    reveal_type(x**x)  # revealed: int
+    reveal_type(2**x)  # revealed: Any
+    reveal_type(x**x)  # revealed: Any
 ```
 
 If the second argument is \<0, a `float` is returned at runtime. If the first argument is \<0 but

crates/ty_python_semantic/resources/mdtest/call/methods.md

@@ -598,9 +598,9 @@ from typing_extensions import Self
 
 reveal_type(object.__new__)  # revealed: def __new__(cls) -> Self@__new__
 reveal_type(object().__new__)  # revealed: def __new__(cls) -> Self@__new__
-# revealed: Overload[(cls, x: @Todo(Support for `typing.TypeAlias`) = Literal[0], /) -> Self@__new__, (cls, x: str | bytes | bytearray, /, base: SupportsIndex) -> Self@__new__]
+# revealed: Overload[(cls, x: str | Buffer | SupportsInt | SupportsIndex | SupportsTrunc = Literal[0], /) -> Self@__new__, (cls, x: str | bytes | bytearray, /, base: SupportsIndex) -> Self@__new__]
 reveal_type(int.__new__)
-# revealed: Overload[(cls, x: @Todo(Support for `typing.TypeAlias`) = Literal[0], /) -> Self@__new__, (cls, x: str | bytes | bytearray, /, base: SupportsIndex) -> Self@__new__]
+# revealed: Overload[(cls, x: str | Buffer | SupportsInt | SupportsIndex | SupportsTrunc = Literal[0], /) -> Self@__new__, (cls, x: str | bytes | bytearray, /, base: SupportsIndex) -> Self@__new__]
 reveal_type((42).__new__)
 
 class X:

crates/ty_python_semantic/resources/mdtest/call/open.md

@@ -10,13 +10,13 @@ import pickle
 
 reveal_type(open(""))  # revealed: TextIOWrapper[_WrappedBuffer]
 reveal_type(open("", "r"))  # revealed: TextIOWrapper[_WrappedBuffer]
-reveal_type(open("", "rb"))  # revealed: @Todo(`builtins.open` return type)
+reveal_type(open("", "rb"))  # revealed: BufferedReader[_BufferedReaderStream]
 
 with open("foo.pickle", "rb") as f:
     x = pickle.load(f)  # fine
 
 def _(mode: str):
-    reveal_type(open("", mode))  # revealed: @Todo(`builtins.open` return type)
+    reveal_type(open("", mode))  # revealed: IO[Any]
 ```
 
 ## `os.fdopen`
@@ -29,7 +29,7 @@ import os
 
 reveal_type(os.fdopen(0))  # revealed: TextIOWrapper[_WrappedBuffer]
 reveal_type(os.fdopen(0, "r"))  # revealed: TextIOWrapper[_WrappedBuffer]
-reveal_type(os.fdopen(0, "rb"))  # revealed: @Todo(`os.fdopen` return type)
+reveal_type(os.fdopen(0, "rb"))  # revealed: BufferedReader[_BufferedReaderStream]
 
 with os.fdopen(0, "rb") as f:
     x = pickle.load(f)  # fine
@@ -43,9 +43,9 @@ And similarly for `Path.open()`:
 from pathlib import Path
 import pickle
 
-reveal_type(Path("").open())  # revealed: @Todo(`Path.open` return type)
-reveal_type(Path("").open("r"))  # revealed: @Todo(`Path.open` return type)
-reveal_type(Path("").open("rb"))  # revealed: @Todo(`Path.open` return type)
+reveal_type(Path("").open())  # revealed: TextIOWrapper[_WrappedBuffer]
+reveal_type(Path("").open("r"))  # revealed: TextIOWrapper[_WrappedBuffer]
+reveal_type(Path("").open("rb"))  # revealed: BufferedReader[_BufferedReaderStream]
 
 with Path("foo.pickle").open("rb") as f:
     x = pickle.load(f)  # fine
@@ -61,7 +61,7 @@ import pickle
 
 reveal_type(NamedTemporaryFile())  # revealed: _TemporaryFileWrapper[bytes]
 reveal_type(NamedTemporaryFile("r"))  # revealed: _TemporaryFileWrapper[str]
-reveal_type(NamedTemporaryFile("rb"))  # revealed: @Todo(`tempfile.NamedTemporaryFile` return type)
+reveal_type(NamedTemporaryFile("rb"))  # revealed: _TemporaryFileWrapper[bytes]
 
 with NamedTemporaryFile("rb") as f:
     x = pickle.load(f)  # fine

crates/ty_python_semantic/resources/mdtest/expression/lambda.md

@@ -127,7 +127,7 @@ x = lambda y: y
 reveal_type(x.__code__)  # revealed: CodeType
 reveal_type(x.__name__)  # revealed: str
 reveal_type(x.__defaults__)  # revealed: tuple[Any, ...] | None
-reveal_type(x.__annotations__)  # revealed: dict[str, @Todo(Support for `typing.TypeAlias`)]
+reveal_type(x.__annotations__)  # revealed: dict[str, Any]
 reveal_type(x.__dict__)  # revealed: dict[str, Any]
 reveal_type(x.__doc__)  # revealed: str | None
 reveal_type(x.__kwdefaults__)  # revealed: dict[str, Any] | None

crates/ty_python_semantic/resources/mdtest/pep613_type_aliases.md

@@ -1,8 +1,124 @@
 # PEP 613 type aliases
 
-## No panics
+PEP 613 type aliases are simple assignment statements, annotated with `typing.TypeAlias` to mark
+them as a type alias. At runtime, they behave the same as implicit type aliases. Our support for
+them is currently the same as for implicit type aliases, but we don't reproduce the full
+implicit-type-alias test suite here, just some particularly interesting cases.
 
-We do not fully support PEP 613 type aliases yet. For now, just make sure that we don't panic:
+## Basic
+
+### as `TypeAlias`
+
+```py
+from typing import TypeAlias
+
+IntOrStr: TypeAlias = int | str
+
+def _(x: IntOrStr):
+    reveal_type(x)  # revealed: int | str
+```
+
+### as `typing.TypeAlias`
+
+```py
+import typing
+
+IntOrStr: typing.TypeAlias = int | str
+
+def _(x: IntOrStr):
+    reveal_type(x)  # revealed: int | str
+```
+
+## Can be used as value
+
+Because PEP 613 type aliases are just annotated assignments, they can be used as values, like a
+legacy type expression (and unlike a PEP 695 type alias). We might prefer this wasn't allowed, but
+people do use it.
+
+```py
+from typing import TypeAlias
+
+MyExc: TypeAlias = Exception
+
+try:
+    raise MyExc("error")
+except MyExc as e:
+    reveal_type(e)  # revealed: Exception
+```
+
+## Unknown type in PEP 604 union
+
+If we run into an unexpected type in a PEP 604 union in the RHS of a PEP 613 type alias, we still
+understand it as a union type, just with an unknown element.
+
+```py
+from typing import TypeAlias
+from nonexistent import unknown_type  # error: [unresolved-import]
+
+MyAlias: TypeAlias = int | unknown_type | str
+
+def _(x: MyAlias):
+    reveal_type(x)  # revealed: int | Unknown | str
+```
+
+## Callable type in union
+
+```py
+from typing import TypeAlias, Callable
+
+MyAlias: TypeAlias = int | Callable[[str], int]
+
+def _(x: MyAlias):
+    # TODO: int | (str) -> int
+    reveal_type(x)  # revealed: int | ((str, /) -> int)
+```
+
+## Subscripted generic alias in union
+
+```py
+from typing import TypeAlias, TypeVar
+
+T = TypeVar("T")
+
+Alias1: TypeAlias = list[T] | set[T]
+MyAlias: TypeAlias = int | Alias1[str]
+
+def _(x: MyAlias):
+    # TODO: int | list[str] | set[str]
+    reveal_type(x)  # revealed: int | @Todo(Specialization of union type alias)
+```
+
+## Imported
+
+`alias.py`:
+
+```py
+from typing import TypeAlias
+
+MyAlias: TypeAlias = int | str
+```
+
+`main.py`:
+
+```py
+from alias import MyAlias
+
+def _(x: MyAlias):
+    reveal_type(x)  # revealed: int | str
+```
+
+## String literal in RHS
+
+```py
+from typing import TypeAlias
+
+IntOrStr: TypeAlias = "int | str"
+
+def _(x: IntOrStr):
+    reveal_type(x)  # revealed: int | str
+```
+
+## Cyclic
 
 ```py
 from typing import TypeAlias
@@ -18,6 +134,26 @@ def _(rec: RecursiveHomogeneousTuple):
     reveal_type(rec)  # revealed: tuple[Divergent, ...]
 ```
 
+## Conditionally imported on Python < 3.10
+
+```toml
+[environment]
+python-version = "3.9"
+```
+
+```py
+try:
+    # error: [unresolved-import]
+    from typing import TypeAlias
+except ImportError:
+    from typing_extensions import TypeAlias
+
+MyAlias: TypeAlias = int
+
+def _(x: MyAlias):
+    reveal_type(x)  # revealed: int
+```
+
 ## PEP-613 aliases in stubs are deferred
 
 Although the right-hand side of a PEP-613 alias is a value expression, inference of this value is
@@ -46,7 +182,31 @@ f(stub.B())
 
 class Unrelated: ...
 
-# TODO: we should emit `[invalid-argument-type]` here
-# (the alias is a `@Todo` because it's imported from another file)
+# error: [invalid-argument-type]
 f(Unrelated())
 ```
+
+## Invalid position
+
+`typing.TypeAlias` must be used as the sole annotation in an annotated assignment. Use in any other
+context is an error.
+
+```py
+from typing import TypeAlias
+
+# error: [invalid-type-form]
+def _(x: TypeAlias):
+    reveal_type(x)  # revealed: Unknown
+
+# error: [invalid-type-form]
+y: list[TypeAlias] = []
+```
+
+## RHS is required
+
+```py
+from typing import TypeAlias
+
+# error: [invalid-type-form]
+Empty: TypeAlias
+```