From 3e7e91724c116019d00a6cfe15ce74683de60da1 Mon Sep 17 00:00:00 2001
From: Alex Waygood <Alex.Waygood@Gmail.com>
Date: Fri, 14 Nov 2025 21:07:02 +0000
Subject: [PATCH] [ty] Further improve details around which expressions should
 be deferred in stub files (#21456)

## Summary

- Always restore the previous `deferred_state` after parsing a type
expression: we don't want that state leaking out into other contexts
where we shouldn't be deferring expression inference
- Always defer the right-hand-side of a PEP-613 type alias in a stub
file, allowing for forward references on the right-hand side of `T:
TypeAlias = X | Y` in a stub file

Addresses @carljm's review in
https://github.com/astral-sh/ruff/pull/21401#discussion_r2524260153

## Test Plan

I added a regression test for a regression that the first version of
this PR introduced (we need to make sure the r.h.s. of a PEP-613
`TypeAlias`es is always deferred in a stub file)
---
 .../resources/mdtest/pep613_type_aliases.md   | 37 ++++++++++++++++++-
 .../src/types/infer/builder.rs                | 21 +++++++++--
 .../types/infer/builder/type_expression.rs    |  7 +++-
 3 files changed, 58 insertions(+), 7 deletions(-)

diff --git a/crates/ty_python_semantic/resources/mdtest/pep613_type_aliases.md b/crates/ty_python_semantic/resources/mdtest/pep613_type_aliases.md
index de3851ddab..5272be683c 100644
--- a/crates/ty_python_semantic/resources/mdtest/pep613_type_aliases.md
+++ b/crates/ty_python_semantic/resources/mdtest/pep613_type_aliases.md
@@ -1,6 +1,8 @@
 # PEP 613 type aliases
 
-We do not support PEP 613 type aliases yet. For now, just make sure that we don't panic:
+## No panics
+
+We do not fully support PEP 613 type aliases yet. For now, just make sure that we don't panic:
 
 ```py
 from typing import TypeAlias
@@ -15,3 +17,36 @@ RecursiveHomogeneousTuple: TypeAlias = tuple[int | "RecursiveHomogeneousTuple",
 def _(rec: RecursiveHomogeneousTuple):
     reveal_type(rec)  # revealed: tuple[Divergent, ...]
 ```
+
+## 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
+deferred in a stub file, allowing for forward references:
+
+`stub.pyi`:
+
+```pyi
+from typing import TypeAlias
+
+MyAlias: TypeAlias = A | B
+
+class A: ...
+class B: ...
+```
+
+`module.py`:
+
+```py
+import stub
+
+def f(x: stub.MyAlias): ...
+
+f(stub.A())
+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)
+f(Unrelated())
+```
diff --git a/crates/ty_python_semantic/src/types/infer/builder.rs b/crates/ty_python_semantic/src/types/infer/builder.rs
index aace19bb45..d3d4673b26 100644
--- a/crates/ty_python_semantic/src/types/infer/builder.rs
+++ b/crates/ty_python_semantic/src/types/infer/builder.rs
@@ -5486,10 +5486,23 @@ impl<'db, 'ast> TypeInferenceBuilder<'db, 'ast> {
                 self.dataclass_field_specifiers = specifiers;
             }
 
-            let inferred_ty = self.infer_maybe_standalone_expression(
-                value,
-                TypeContext::new(Some(declared.inner_type())),
-            );
+            // We defer the r.h.s. of PEP-613 `TypeAlias` assignments in stub files.
+            let declared_type = declared.inner_type();
+            let previous_deferred_state = self.deferred_state;
+
+            if matches!(
+                declared_type,
+                Type::SpecialForm(SpecialFormType::TypeAlias)
+                    | Type::Dynamic(DynamicType::TodoTypeAlias)
+            ) && self.in_stub()
+            {
+                self.deferred_state = DeferredExpressionState::Deferred;
+            }
+
+            let inferred_ty = self
+                .infer_maybe_standalone_expression(value, TypeContext::new(Some(declared_type)));
+
+            self.deferred_state = previous_deferred_state;
 
             self.dataclass_field_specifiers.clear();
 
diff --git a/crates/ty_python_semantic/src/types/infer/builder/type_expression.rs b/crates/ty_python_semantic/src/types/infer/builder/type_expression.rs
index 513361ff4e..c33c058892 100644
--- a/crates/ty_python_semantic/src/types/infer/builder/type_expression.rs
+++ b/crates/ty_python_semantic/src/types/infer/builder/type_expression.rs
@@ -20,10 +20,12 @@ use crate::types::{
 impl<'db> TypeInferenceBuilder<'db, '_> {
     /// Infer the type of a type expression.
     pub(super) fn infer_type_expression(&mut self, expression: &ast::Expr) -> Type<'db> {
+        let previous_deferred_state = self.deferred_state;
+
         // `DeferredExpressionState::InStringAnnotation` takes precedence over other states.
         // However, if it's not a stringified annotation, we must still ensure that annotation expressions
         // are always deferred in stub files.
-        match self.deferred_state {
+        match previous_deferred_state {
             DeferredExpressionState::None => {
                 if self.in_stub() {
                     self.deferred_state = DeferredExpressionState::Deferred;
@@ -31,8 +33,9 @@ impl<'db> TypeInferenceBuilder<'db, '_> {
             }
             DeferredExpressionState::InStringAnnotation(_) | DeferredExpressionState::Deferred => {}
         }
-
         let mut ty = self.infer_type_expression_no_store(expression);
+        self.deferred_state = previous_deferred_state;
+
         let divergent = Type::divergent(Some(self.scope()));
         if ty.has_divergent_type(self.db(), divergent) {
             ty = divergent;