[ty] Improve robustness of various type-qualifier-related checks

crates/ty_python_semantic/resources/mdtest/type_qualifiers/classvar.md

@@ -333,18 +333,18 @@ python-version = "3.12"
 ```
 
 ```py
-from typing import ClassVar
+from typing import ClassVar, TypedDict
 from ty_extensions import reveal_mro
 
-# error: [invalid-type-form] "`ClassVar` annotations are only allowed in class-body scopes"
+# error: [invalid-type-form] "`ClassVar` is only allowed in class bodies"
 x: ClassVar[int] = 1
 
 class C:
     def __init__(self) -> None:
         # error: [invalid-type-form] "`ClassVar` annotations are not allowed for non-name targets"
         self.x: ClassVar[int] = 1
 
-        # error: [invalid-type-form] "`ClassVar` annotations are only allowed in class-body scopes"
+        # error: [invalid-type-form] "`ClassVar` is only allowed in class bodies"
         y: ClassVar[int] = 1
 
 # error: [invalid-type-form] "Type qualifier `typing.ClassVar` is not allowed in parameter annotations"
@@ -369,6 +369,12 @@ class Foo(ClassVar[tuple[int]]): ...
 # TODO: Show `Unknown` instead of `@Todo` type in the MRO; or ignore `ClassVar` and show the MRO as if `ClassVar` was not there
 # revealed: (<class 'Foo'>, @Todo(Inference of subscript on special form), <class 'object'>)
 reveal_mro(Foo)
+
+class Foo(TypedDict):
+    # error: [invalid-type-form] "`ClassVar` is not allowed in TypedDict fields"
+    x: ClassVar[int]
+    # error: [invalid-type-form] "`ClassVar` is not allowed in TypedDict fields"
+    y: ClassVar
 ```
 
 [`typing.classvar`]: https://docs.python.org/3/library/typing.html#typing.ClassVar

crates/ty_python_semantic/resources/mdtest/type_qualifiers/final.md

@@ -662,7 +662,7 @@ python-version = "3.12"
 ```
 
 ```py
-from typing import Final, ClassVar, Annotated
+from typing import Final, ClassVar, Annotated, TypedDict
 from ty_extensions import reveal_mro
 
 LEGAL_A: Final[int] = 1
@@ -703,6 +703,18 @@ class Foo(Final[tuple[int]]): ...
 # TODO: Show `Unknown` instead of `@Todo` type in the MRO; or ignore `Final` and show the MRO as if `Final` was not there
 # revealed: (<class 'Foo'>, @Todo(Inference of subscript on special form), <class 'object'>)
 reveal_mro(Foo)
+
+class Foo(TypedDict):
+    # error: [invalid-type-form] "`Final` is not allowed in TypedDict fields"
+    # error: [invalid-typed-dict-statement] "TypedDict item cannot have a value"
+    a: Final[int] = 42
+    # error: [invalid-type-form] "`Final` is not allowed in TypedDict fields"
+    # error: [invalid-typed-dict-statement] "TypedDict item cannot have a value"
+    b: Final = 56
+    # error: [invalid-type-form] "`Final` is not allowed in TypedDict fields"
+    c: Final[int]
+    # error: [invalid-type-form] "`Final` is not allowed in TypedDict fields"
+    d: Final
 ```
 
 ### Attribute assignment outside `__init__`

crates/ty_python_semantic/resources/mdtest/type_qualifiers/initvar.md

@@ -144,9 +144,10 @@ class AlsoWrong:
 `InitVar` annotations are not allowed outside of dataclass attribute annotations:
 
 ```py
+from typing import TypedDict
 from dataclasses import InitVar, dataclass
 
-# error: [invalid-type-form] "`InitVar` annotations are only allowed in class-body scopes"
+# error: [invalid-type-form] "`InitVar` is only allowed in dataclass fields"
 x: InitVar[int] = 1
 
 # error: [invalid-type-form] "Type qualifier `dataclasses.InitVar` is not allowed in parameter annotations"
@@ -158,7 +159,11 @@ def g() -> InitVar[int]:
     return 1
 
 class C:
-    # TODO: this would ideally be an error
+    # error: [invalid-type-form] "`InitVar` is only allowed in dataclass fields"
+    x: InitVar[int]
+
+class D(TypedDict):
+    # error: [invalid-type-form] "`InitVar` is not allowed in TypedDict fields"
     x: InitVar[int]
 
 @dataclass

crates/ty_python_semantic/resources/mdtest/typed_dict.md

@@ -2772,9 +2772,9 @@ from typing import TypedDict
 x: TypedDict = {"name": "Alice"}
 ```
 
-### `ReadOnly`, `Required` and `NotRequired` not allowed in parameter annotations
+### `ReadOnly`, `Required` and `NotRequired` not allowed in parameter annotations or return annotations
 
-```py
+```pyi
 from typing_extensions import Required, NotRequired, ReadOnly
 
 def bad(
@@ -2785,29 +2785,62 @@ def bad(
     # error: [invalid-type-form] "Type qualifier `typing.ReadOnly` is not allowed in parameter annotations"
     c: ReadOnly[int],
 ): ...
+
+# error: [invalid-type-form] "Type qualifier `typing.Required` is not allowed in return type annotations"
+def bad2() -> Required[int]: ...
+
+# error: [invalid-type-form] "Type qualifier `typing.NotRequired` is not allowed in return type annotations"
+def bad2() -> NotRequired[int]: ...
+
+# error: [invalid-type-form] "Type qualifier `typing.ReadOnly` is not allowed in return type annotations"
+def bad2() -> ReadOnly[int]: ...
+```
+
+### `Required`, `NotRequired` and `ReadOnly` require exactly one argument
+
+```py
+from typing_extensions import TypedDict, ReadOnly, Required, NotRequired
+
+class Foo(TypedDict):
+    a: Required  # error: [invalid-type-form] "`Required` may not be used without a type argument"
+    b: Required[()]  # error: [invalid-type-form] "Type qualifier `typing.Required` expected exactly 1 argument, got 0"
+    c: Required[int, str]  # error: [invalid-type-form] "Type qualifier `typing.Required` expected exactly 1 argument, got 2"
+    d: NotRequired  # error: [invalid-type-form] "`NotRequired` may not be used without a type argument"
+    e: NotRequired[()]  # error: [invalid-type-form] "Type qualifier `typing.NotRequired` expected exactly 1 argument, got 0"
+    # error: [invalid-type-form] "Type qualifier `typing.NotRequired` expected exactly 1 argument, got 2"
+    f: NotRequired[int, str]
+    g: ReadOnly  # error: [invalid-type-form] "`ReadOnly` may not be used without a type argument"
+    h: ReadOnly[()]  # error: [invalid-type-form] "Type qualifier `typing.ReadOnly` expected exactly 1 argument, got 0"
+    i: ReadOnly[int, str]  # error: [invalid-type-form] "Type qualifier `typing.ReadOnly` expected exactly 1 argument, got 2"
 ```
 
-### `Required` and `NotRequired` not allowed outside `TypedDict`
+### `Required`, `NotRequired` and `ReadOnly` are not allowed outside `TypedDict`
 
 ```py
-from typing_extensions import Required, NotRequired, TypedDict
+from typing_extensions import Required, NotRequired, TypedDict, ReadOnly
 
 # error: [invalid-type-form] "`Required` is only allowed in TypedDict fields"
 x: Required[int]
 # error: [invalid-type-form] "`NotRequired` is only allowed in TypedDict fields"
 y: NotRequired[str]
+# error: [invalid-type-form] "`ReadOnly` is only allowed in TypedDict fields"
+z: ReadOnly[str]
 
 class MyClass:
     # error: [invalid-type-form] "`Required` is only allowed in TypedDict fields"
     x: Required[int]
     # error: [invalid-type-form] "`NotRequired` is only allowed in TypedDict fields"
     y: NotRequired[str]
+    # error: [invalid-type-form] "`ReadOnly` is only allowed in TypedDict fields"
+    z: ReadOnly[str]
 
 def f():
     # error: [invalid-type-form] "`Required` is only allowed in TypedDict fields"
     x: Required[int] = 1
     # error: [invalid-type-form] "`NotRequired` is only allowed in TypedDict fields"
     y: NotRequired[str] = ""
+    # error: [invalid-type-form] "`ReadOnly` is only allowed in TypedDict fields"
+    z: ReadOnly[str]
 
 # fine
 MyFunctionalTypedDict = TypedDict("MyFunctionalTypedDict", {"not-an-identifier": Required[int]})

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

@@ -18,6 +18,7 @@ use ruff_python_stdlib::typing::as_pep_585_generic;
 use ruff_text_size::{Ranged, TextRange};
 use rustc_hash::{FxHashMap, FxHashSet};
 use smallvec::SmallVec;
+use strum::IntoEnumIterator;
 use ty_module_resolver::{KnownModule, ModuleName, resolve_module};
 
 use super::deferred;
@@ -100,6 +101,7 @@ use crate::types::mro::DynamicMroErrorKind;
 use crate::types::newtype::NewType;
 use crate::types::set_theoretic::RecursivelyDefined;
 use crate::types::signatures::CallableSignature;
+use crate::types::special_form::TypeQualifier;
 use crate::types::subclass_of::SubclassOfInner;
 use crate::types::tuple::{Tuple, TupleLength, TupleSpecBuilder, TupleType};
 use crate::types::type_alias::{ManualPEP695TypeAliasType, PEP695TypeAliasType};
@@ -3864,8 +3866,11 @@ impl<'db, 'ast> TypeInferenceBuilder<'db, 'ast> {
             );
 
             if !annotated.qualifiers.is_empty() {
-                for qualifier in [TypeQualifiers::CLASS_VAR, TypeQualifiers::INIT_VAR] {
-                    if annotated.qualifiers.contains(qualifier)
+                for qualifier in TypeQualifier::iter() {
+                    if !qualifier.is_valid_for_non_name_targets()
+                        && annotated
+                            .qualifiers
+                            .contains(TypeQualifiers::from(qualifier))
                         && let Some(builder) = self
                             .context
                             .report_lint(&INVALID_TYPE_FORM, annotation.as_ref())
@@ -4141,45 +4146,117 @@ impl<'db, 'ast> TypeInferenceBuilder<'db, 'ast> {
         }
 
         if !declared.qualifiers.is_empty() {
-            let current_scope_id = self.scope().file_scope_id(self.db());
-            let current_scope = self.index.scope(current_scope_id);
-            if current_scope.kind() != ScopeKind::Class {
-                for qualifier in [TypeQualifiers::CLASS_VAR, TypeQualifiers::INIT_VAR] {
-                    if declared.qualifiers.contains(qualifier)
-                        && let Some(builder) =
-                            self.context.report_lint(&INVALID_TYPE_FORM, annotation)
-                    {
-                        builder.into_diagnostic(format_args!(
-                            "`{name}` annotations are only allowed in class-body scopes",
-                            name = qualifier.name()
-                        ));
+            for qualifier in TypeQualifier::iter() {
+                if !declared
+                    .qualifiers
+                    .contains(TypeQualifiers::from(qualifier))
+                {
+                    continue;
+                }
+                let current_scope_id = self.scope().file_scope_id(self.db());
+
+                if self.index.scope(current_scope_id).kind() != ScopeKind::Class {
+                    match qualifier {
+                        TypeQualifier::Final => {}
+                        TypeQualifier::ClassVar => {
+                            if let Some(builder) =
+                                self.context.report_lint(&INVALID_TYPE_FORM, annotation)
+                            {
+                                builder
+                                    .into_diagnostic("`ClassVar` is only allowed in class bodies");
+                            }
+                        }
+                        TypeQualifier::InitVar => {
+                            if let Some(builder) =
+                                self.context.report_lint(&INVALID_TYPE_FORM, annotation)
+                            {
+                                builder.into_diagnostic(
+                                    "`InitVar` is only allowed in dataclass fields",
+                                );
+                            }
+                        }
+                        TypeQualifier::NotRequired
+                        | TypeQualifier::ReadOnly
+                        | TypeQualifier::Required => {
+                            if let Some(builder) =
+                                self.context.report_lint(&INVALID_TYPE_FORM, annotation)
+                            {
+                                builder.into_diagnostic(format_args!(
+                                    "`{name}` is only allowed in TypedDict fields",
+                                    name = qualifier.name()
+                                ));
+                            }
+                        }
                     }
+
+                    continue;
                 }
-            }
 
-            // `Required`, `NotRequired`, and `ReadOnly` are only valid inside TypedDict classes.
-            if declared.qualifiers.intersects(
-                TypeQualifiers::REQUIRED | TypeQualifiers::NOT_REQUIRED | TypeQualifiers::READ_ONLY,
-            ) {
-                let in_typed_dict = current_scope.kind() == ScopeKind::Class
-                    && nearest_enclosing_class(self.db(), self.index, self.scope())
-                        .is_some_and(|class| class.is_typed_dict(self.db()));
-                if !in_typed_dict {
-                    for qualifier in [
-                        TypeQualifiers::REQUIRED,
-                        TypeQualifiers::NOT_REQUIRED,
-                        TypeQualifiers::READ_ONLY,
-                    ] {
-                        if declared.qualifiers.contains(qualifier)
-                            && let Some(builder) =
+                let nearest_enclosing_class =
+                    nearest_enclosing_class(self.db(), self.index, self.scope());
+                let class_kind = nearest_enclosing_class.and_then(|class| {
+                    CodeGeneratorKind::from_class(self.db(), ClassLiteral::Static(class), None)
+                });
+
+                match class_kind {
+                    Some(CodeGeneratorKind::TypedDict) => match qualifier {
+                        TypeQualifier::ClassVar | TypeQualifier::Final | TypeQualifier::InitVar => {
+                            let Some(builder) =
                                 self.context.report_lint(&INVALID_TYPE_FORM, annotation)
-                        {
+                            else {
+                                continue;
+                            };
+                            builder.into_diagnostic(format_args!(
+                                "`{name}` is not allowed in TypedDict fields",
+                                name = qualifier.name()
+                            ));
+                        }
+                        TypeQualifier::NotRequired
+                        | TypeQualifier::ReadOnly
+                        | TypeQualifier::Required => {}
+                    },
+                    Some(CodeGeneratorKind::DataclassLike(_)) => match qualifier {
+                        TypeQualifier::NotRequired
+                        | TypeQualifier::ReadOnly
+                        | TypeQualifier::Required => {
+                            let Some(builder) =
+                                self.context.report_lint(&INVALID_TYPE_FORM, annotation)
+                            else {
+                                continue;
+                            };
+                            builder.into_diagnostic(format_args!(
+                                "`{name}` is not allowed in dataclass fields",
+                                name = qualifier.name()
+                            ));
+                        }
+                        TypeQualifier::ClassVar | TypeQualifier::Final | TypeQualifier::InitVar => {
+                        }
+                    },
+                    Some(CodeGeneratorKind::NamedTuple) | None => match qualifier {
+                        TypeQualifier::NotRequired
+                        | TypeQualifier::Required
+                        | TypeQualifier::ReadOnly => {
+                            let Some(builder) =
+                                self.context.report_lint(&INVALID_TYPE_FORM, annotation)
+                            else {
+                                continue;
+                            };
                             builder.into_diagnostic(format_args!(
                                 "`{name}` is only allowed in TypedDict fields",
                                 name = qualifier.name()
                             ));
                         }
-                    }
+                        TypeQualifier::InitVar => {
+                            let Some(builder) =
+                                self.context.report_lint(&INVALID_TYPE_FORM, annotation)
+                            else {
+                                continue;
+                            };
+                            builder
+                                .into_diagnostic("`InitVar` is only allowed in dataclass fields");
+                        }
+                        TypeQualifier::ClassVar | TypeQualifier::Final => {}
+                    },
                 }
             }
         }

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

@@ -85,25 +85,30 @@ impl<'db> TypeInferenceBuilder<'db, '_> {
         ) -> TypeAndQualifiers<'db> {
             let special_case = match ty {
                 Type::SpecialForm(special_form) => match special_form {
-                    SpecialFormType::TypeQualifier(TypeQualifier::InitVar) => {
-                        if let Some(builder) =
-                            builder.context.report_lint(&INVALID_TYPE_FORM, annotation)
-                        {
-                            builder.into_diagnostic(
-                                "`InitVar` may not be used without a type argument",
-                            );
+                    SpecialFormType::TypeQualifier(qualifier) => {
+                        match qualifier {
+                            TypeQualifier::InitVar
+                            | TypeQualifier::ReadOnly
+                            | TypeQualifier::NotRequired
+                            | TypeQualifier::Required => {
+                                if let Some(builder) =
+                                    builder.context.report_lint(&INVALID_TYPE_FORM, annotation)
+                                {
+                                    builder.into_diagnostic(format_args!(
+                                        "`{}` may not be used without a type argument",
+                                        qualifier.name(),
+                                    ));
+                                }
+                            }
+                            TypeQualifier::ClassVar | TypeQualifier::Final => {}
                         }
+
                         Some(TypeAndQualifiers::new(
                             Type::unknown(),
                             TypeOrigin::Declared,
-                            TypeQualifiers::INIT_VAR,
+                            TypeQualifiers::from(qualifier),
                         ))
                     }
-                    SpecialFormType::TypeQualifier(qualifier) => Some(TypeAndQualifiers::new(
-                        Type::unknown(),
-                        TypeOrigin::Declared,
-                        TypeQualifiers::from(qualifier),
-                    )),
                     SpecialFormType::TypeAlias if pep_613_policy == PEP613Policy::Allowed => {
                         Some(TypeAndQualifiers::declared(ty))
                     }