[ty] Ensure annotation/type expressions in stub files are always deferred (#21401)

Author: AlexWaygood

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

@@ -396,3 +396,34 @@ B = NewType("B", list[Any])
 T = TypeVar("T")
 C = NewType("C", list[T])  # TODO: should be "error: [invalid-newtype]"
 ```
+
+## Forward references in stub files
+
+Stubs natively support forward references, so patterns that would raise `NameError` at runtime are
+allowed in stub files:
+
+`stub.pyi`:
+
+```pyi
+from typing import NewType
+
+N = NewType("N", A)
+
+class A: ...
+```
+
+`main.py`:
+
+```py
+from stub import N, A
+
+n = N(A())  # fine
+
+def f(x: A): ...
+
+f(n)  # fine
+
+class Invalid: ...
+
+bad = N(Invalid())  # error: [invalid-argument-type]
+```

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

@@ -266,7 +266,48 @@ from typing import TypeVar
 
 # error: [invalid-legacy-type-variable]
 T = TypeVar("T", invalid_keyword=True)
+```
+
+### Forward references in stubs
+
+Stubs natively support forward references, so patterns that would raise `NameError` at runtime are
+allowed in stub files:
+
+`stub.pyi`:
+
+```pyi
+from typing import TypeVar
+
+T = TypeVar("T", bound=A, default=B)
+U = TypeVar("U", C, D)
+
+class A: ...
+class B(A): ...
+class C: ...
+class D: ...
+
+def f(x: T) -> T: ...
+def g(x: U) -> U: ...
+```
+
+`main.py`:
+
+```py
+from stub import f, g, A, B, C, D
+
+reveal_type(f(A()))  # revealed: A
+reveal_type(f(B()))  # revealed: B
+reveal_type(g(C()))  # revealed: C
+reveal_type(g(D()))  # revealed: D
+
+# TODO: one diagnostic would probably be sufficient here...?
+#
+# error: [invalid-argument-type] "Argument type `C` does not satisfy upper bound `A` of type variable `T`"
+# error: [invalid-argument-type] "Argument to function `f` is incorrect: Expected `B`, found `C`"
+reveal_type(f(C()))  # revealed: B
 
+# error: [invalid-argument-type]
+reveal_type(g(A()))  # revealed: Unknown
 ```
 
 ### Constructor signature versioning

crates/ty_python_semantic/resources/mdtest/narrow/isinstance.md

@@ -82,7 +82,7 @@ def _(x: int | str | bytes | memoryview | range):
     if isinstance(x, int | str):
         reveal_type(x)  # revealed: int | str
     elif isinstance(x, bytes | memoryview):
-        reveal_type(x)  # revealed: bytes | memoryview[Unknown]
+        reveal_type(x)  # revealed: bytes | memoryview[int]
     else:
         reveal_type(x)  # revealed: range
 ```
@@ -242,11 +242,11 @@ def _(flag: bool):
 def _(flag: bool):
     x = 1 if flag else "a"
 
-    # error: [invalid-argument-type] "Argument to function `isinstance` is incorrect: Expected `type | UnionType | tuple[Unknown, ...]`, found `Literal["a"]"
+    # error: [invalid-argument-type] "Argument to function `isinstance` is incorrect: Expected `type | UnionType | tuple[Divergent, ...]`, found `Literal["a"]"
     if isinstance(x, "a"):
         reveal_type(x)  # revealed: Literal[1, "a"]
 
-    # error: [invalid-argument-type] "Argument to function `isinstance` is incorrect: Expected `type | UnionType | tuple[Unknown, ...]`, found `Literal["int"]"
+    # error: [invalid-argument-type] "Argument to function `isinstance` is incorrect: Expected `type | UnionType | tuple[Divergent, ...]`, found `Literal["int"]"
     if isinstance(x, "int"):
         reveal_type(x)  # revealed: Literal[1, "a"]
 ```

crates/ty_python_semantic/resources/mdtest/narrow/issubclass.md

@@ -283,7 +283,7 @@ def flag() -> bool:
 
 t = int if flag() else str
 
-# error: [invalid-argument-type] "Argument to function `issubclass` is incorrect: Expected `type | UnionType | tuple[Unknown, ...]`, found `Literal["str"]"
+# error: [invalid-argument-type] "Argument to function `issubclass` is incorrect: Expected `type | UnionType | tuple[Divergent, ...]`, found `Literal["str"]"
 if issubclass(t, "str"):
     reveal_type(t)  # revealed: <class 'int'> | <class 'str'>
 

crates/ty_python_semantic/resources/mdtest/paramspec.md

@@ -102,6 +102,20 @@ Other values are invalid.
 P4 = ParamSpec("P4", default=int)
 ```
 
+### Forward references in stub files
+
+Stubs natively support forward references, so patterns that would raise `NameError` at runtime are
+allowed in stub files:
+
+```pyi
+from typing_extensions import ParamSpec
+
+P = ParamSpec("P", default=[A, B])
+
+class A: ...
+class B: ...
+```
+
 ### PEP 695
 
 ```toml

crates/ty_python_semantic/src/types/infer/builder.rs

@@ -6921,6 +6921,7 @@ impl<'db, 'ast> TypeInferenceBuilder<'db, 'ast> {
         ty
     }
 
+    #[track_caller]
     fn store_expression_type(&mut self, expression: &ast::Expr, ty: Type<'db>) {
         if self.deferred_state.in_string_annotation() {
             // Avoid storing the type of expressions that are part of a string annotation because

crates/ty_python_semantic/src/types/infer/builder/annotation_expression.rs

@@ -18,7 +18,18 @@ impl<'db> TypeInferenceBuilder<'db, '_> {
         annotation: &ast::Expr,
         deferred_state: DeferredExpressionState,
     ) -> TypeAndQualifiers<'db> {
-        let previous_deferred_state = std::mem::replace(&mut self.deferred_state, deferred_state);
+        // `DeferredExpressionState::InStringAnnotation` takes precedence over other deferred states.
+        // However, if it's not a stringified annotation, we must still ensure that annotation expressions
+        // are always deferred in stub files.
+        let state = if deferred_state.in_string_annotation() {
+            deferred_state
+        } else if self.in_stub() {
+            DeferredExpressionState::Deferred
+        } else {
+            deferred_state
+        };
+
+        let previous_deferred_state = std::mem::replace(&mut self.deferred_state, state);
         let annotation_ty = self.infer_annotation_expression_impl(annotation);
         self.deferred_state = previous_deferred_state;
         annotation_ty

crates/ty_python_semantic/src/types/infer/builder/type_expression.rs

@@ -20,6 +20,18 @@ 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> {
+        // `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 {
+            DeferredExpressionState::None => {
+                if self.in_stub() {
+                    self.deferred_state = DeferredExpressionState::Deferred;
+                }
+            }
+            DeferredExpressionState::InStringAnnotation(_) | DeferredExpressionState::Deferred => {}
+        }
+
         let mut ty = self.infer_type_expression_no_store(expression);
         let divergent = Type::divergent(Some(self.scope()));
         if ty.has_divergent_type(self.db(), divergent) {