[ty] Validate writes to TypedDict keys

Author: sharkdp

crates/ty/docs/rules.md

@@ -28,6 +28,15 @@ mdtest path: crates/ty_python_semantic/resources/mdtest/typed_dict.md
 14 | 
 15 | def access_with_str_key(person: Person, str_key: str):
 16 |     person[str_key]  # error: [invalid-key]
+17 | 
+18 | def write_to_key_with_wrong_type(person: Person):
+19 |     person["age"] = "42"  # error: [invalid-assignment]
+20 | 
+21 | def write_to_non_existing_key(person: Person):
+22 |     person["naem"] = "Alice"  # error: [invalid-key]
+23 | 
+24 | def write_to_non_literal_string_key(person: Person, str_key: str):
+25 |     person[str_key] = "Alice"  # error: [invalid-assignment]
 ```
 
 # Diagnostics
@@ -71,7 +80,54 @@ error[invalid-key]: TypedDict `Person` cannot be indexed with a key of type `str
 15 | def access_with_str_key(person: Person, str_key: str):
 16 |     person[str_key]  # error: [invalid-key]
    |            ^^^^^^^
+17 |
+18 | def write_to_key_with_wrong_type(person: Person):
    |
 info: rule `invalid-key` is enabled by default
 
 ```
+
+```
+error[invalid-assignment]: Invalid assignment to key "age" with declared type `int | None` on TypedDict `Person`
+  --> src/mdtest_snippet.py:19:5
+   |
+18 | def write_to_key_with_wrong_type(person: Person):
+19 |     person["age"] = "42"  # error: [invalid-assignment]
+   |     ------ -----    ^^^^ value of type `Literal["42"]`
+   |     |      |
+   |     |      key has declared type `int | None`
+   |     TypedDict `Person`
+20 |
+21 | def write_to_non_existing_key(person: Person):
+   |
+info: rule `invalid-assignment` is enabled by default
+
+```
+
+```
+error[invalid-key]: Invalid key access on TypedDict `Person`
+  --> src/mdtest_snippet.py:22:5
+   |
+21 | def write_to_non_existing_key(person: Person):
+22 |     person["naem"] = "Alice"  # error: [invalid-key]
+   |     ------ ^^^^^^ Unknown key "naem" - did you mean "name"?
+   |     |
+   |     TypedDict `Person`
+23 |
+24 | def write_to_non_literal_string_key(person: Person, str_key: str):
+   |
+info: rule `invalid-key` is enabled by default
+
+```
+
+```
+error[invalid-assignment]: Can not assign value of type `Literal["Alice"]` to key of type `str` on TypedDict `Person`
+  --> src/mdtest_snippet.py:25:12
+   |
+24 | def write_to_non_literal_string_key(person: Person, str_key: str):
+25 |     person[str_key] = "Alice"  # error: [invalid-assignment]
+   |            ^^^^^^^
+   |
+info: rule `invalid-assignment` is enabled by default
+
+```

crates/ty_python_semantic/resources/mdtest/snapshots/typed_dict.md_-_`TypedDict`_-_Diagnostics_(e5289abf5c570c29).snap

@@ -69,7 +69,8 @@ Assignments to keys are also validated:
 ```py
 # TODO: this should be an error
 alice["name"] = None
-# TODO: this should be an error
+
+# error: [invalid-assignment] "Invalid assignment to key "name" with declared type `str` on TypedDict `Person`: value of type `None`"
 bob["name"] = None
 ```
 
@@ -78,7 +79,8 @@ Assignments to non-existing keys are disallowed:
 ```py
 # TODO: this should be an error
 alice["extra"] = True
-# TODO: this should be an error
+
+# error: [invalid-key] "Invalid key access on TypedDict `Person`: Unknown key "extra""
 bob["extra"] = True
 ```
 
@@ -138,6 +140,8 @@ reveal_type(alice["name"])  # revealed: Unknown
 
 ## Key-based access
 
+### Reading
+
 ```py
 from typing import TypedDict, Final, Literal, Any
 
@@ -169,6 +173,54 @@ def _(person: Person, literal_key: Literal["age"], union_of_keys: Literal["age",
     reveal_type(person[unknown_key])  # revealed: Unknown
 ```
 
+### Writing
+
+```py
+from typing import TypedDict, Final, Literal, Any
+
+class Person(TypedDict):
+    name: str
+    surname: str
+    age: int | None
+
+NAME_FINAL: Final = "name"
+AGE_FINAL: Final[Literal["age"]] = "age"
+
+def _(person: Person):
+    person["name"] = "Alice"
+    person["age"] = 30
+
+    # error: [invalid-key] "Invalid key access on TypedDict `Person`: Unknown key "naem" - did you mean "name"?"
+    person["naem"] = "Alice"
+
+def _(person: Person):
+    person[NAME_FINAL] = "Alice"
+    person[AGE_FINAL] = 30
+
+def _(person: Person, literal_key: Literal["age"]):
+    person[literal_key] = 22
+
+def _(person: Person, union_of_keys: Literal["name", "surname"]):
+    person[union_of_keys] = "unknown"
+
+    # error: [invalid-assignment] "Can not assign value of type `Literal[1]` to key of type `Literal["name", "surname"]` on TypedDict `Person`"
+    person[union_of_keys] = 1
+
+def _(person: Person, union_of_keys: Literal["name", "age"], unknown_value: Any):
+    person[union_of_keys] = unknown_value
+
+    # error: [invalid-assignment] "Can not assign value of type `None` to key of type `Literal["name", "age"]` on TypedDict `Person`"
+    person[union_of_keys] = None
+
+def _(person: Person, str_key: str):
+    # error: [invalid-assignment] "Can not assign value of type `None` to key of type `str` on TypedDict `Person`"
+    person[str_key] = None
+
+def _(person: Person, unknown_key: Any):
+    # No error here:
+    person[unknown_key] = "Eve"
+```
+
 ## Methods on `TypedDict`
 
 ```py
@@ -398,6 +450,15 @@ def access_invalid_key(person: Person):
 
 def access_with_str_key(person: Person, str_key: str):
     person[str_key]  # error: [invalid-key]
+
+def write_to_key_with_wrong_type(person: Person):
+    person["age"] = "42"  # error: [invalid-assignment]
+
+def write_to_non_existing_key(person: Person):
+    person["naem"] = "Alice"  # error: [invalid-key]
+
+def write_to_non_literal_string_key(person: Person, str_key: str):
+    person[str_key] = "Alice"  # error: [invalid-assignment]
 ```
 
 [`typeddict`]: https://typing.python.org/en/latest/spec/typeddict.html

crates/ty_python_semantic/resources/mdtest/typed_dict.md

@@ -2022,18 +2022,30 @@ impl<'db> ClassLiteral<'db> {
                 None
             }
             (CodeGeneratorKind::TypedDict, "__setitem__") => {
-                // TODO: synthesize a set of overloads with precise types
-                let signature = Signature::new(
-                    Parameters::new([
-                        Parameter::positional_only(Some(Name::new_static("self")))
-                            .with_annotated_type(instance_ty),
-                        Parameter::positional_only(Some(Name::new_static("key"))),
-                        Parameter::positional_only(Some(Name::new_static("value"))),
-                    ]),
-                    Some(Type::none(db)),
-                );
+                let fields = self.fields(db, specialization, field_policy);
 
-                Some(CallableType::function_like(db, signature))
+                // Add (key type, value type) overloads for all TypedDict items ("fields"):
+                let overloads = fields.iter().map(|(name, field)| {
+                    let key_type = Type::StringLiteral(StringLiteralType::new(db, name.as_str()));
+
+                    Signature::new(
+                        Parameters::new([
+                            Parameter::positional_only(Some(Name::new_static("self")))
+                                .with_annotated_type(instance_ty),
+                            Parameter::positional_only(Some(Name::new_static("key")))
+                                .with_annotated_type(key_type),
+                            Parameter::positional_only(Some(Name::new_static("value")))
+                                .with_annotated_type(field.declared_ty),
+                        ]),
+                        Some(Type::none(db)),
+                    )
+                });
+
+                Some(Type::Callable(CallableType::new(
+                    db,
+                    CallableSignature::from_overloads(overloads),
+                    true,
+                )))
             }
             (CodeGeneratorKind::TypedDict, "__getitem__") => {
                 let fields = self.fields(db, specialization, field_policy);

crates/ty_python_semantic/src/types/class.rs

@@ -8,7 +8,7 @@ use super::{
 use crate::lint::{Level, LintRegistryBuilder, LintStatus};
 use crate::suppression::FileSuppressionId;
 use crate::types::LintDiagnosticGuard;
-use crate::types::class::{SolidBase, SolidBaseKind};
+use crate::types::class::{Field, SolidBase, SolidBaseKind};
 use crate::types::function::KnownFunction;
 use crate::types::string_annotation::{
     BYTE_STRING_TYPE_ANNOTATION, ESCAPE_CHARACTER_IN_FORWARD_ANNOTATION, FSTRING_TYPE_ANNOTATION,
@@ -17,9 +17,10 @@ use crate::types::string_annotation::{
 };
 use crate::types::{SpecialFormType, Type, protocol_class::ProtocolClassLiteral};
 use crate::util::diagnostics::format_enumeration;
-use crate::{Db, FxIndexMap, Module, ModuleName, Program, declare_lint};
+use crate::{Db, FxIndexMap, FxOrderMap, Module, ModuleName, Program, declare_lint};
 use itertools::Itertools;
 use ruff_db::diagnostic::{Annotation, Diagnostic, SubDiagnostic, SubDiagnosticSeverity};
+use ruff_python_ast::name::Name;
 use ruff_python_ast::{self as ast, AnyNodeRef};
 use ruff_text_size::{Ranged, TextRange};
 use rustc_hash::FxHashSet;
@@ -2570,6 +2571,53 @@ fn report_invalid_base<'ctx, 'db>(
     Some(diagnostic)
 }
 
+pub(crate) fn report_invalid_key_on_typed_dict<'db>(
+    context: &InferContext<'db, '_>,
+    value_node: AnyNodeRef,
+    slice_node: AnyNodeRef,
+    value_ty: Type<'db>,
+    slice_ty: Type<'db>,
+    items: &FxOrderMap<Name, Field<'db>>,
+) {
+    let db = context.db();
+    if let Some(builder) = context.report_lint(&INVALID_KEY, slice_node) {
+        match slice_ty {
+            Type::StringLiteral(key) => {
+                let key = key.value(db);
+                let typed_dict_name = value_ty.display(db);
+
+                let mut diagnostic = builder.into_diagnostic(format_args!(
+                    "Invalid key access on TypedDict `{typed_dict_name}`",
+                ));
+
+                diagnostic.annotate(
+                    context
+                        .secondary(value_node)
+                        .message(format_args!("TypedDict `{typed_dict_name}`")),
+                );
+
+                let existing_keys = items.iter().map(|(name, _)| name.as_str());
+
+                diagnostic.set_primary_message(format!(
+                    "Unknown key \"{key}\"{hint}",
+                    hint = if let Some(suggestion) = did_you_mean(existing_keys, key) {
+                        format!(" - did you mean \"{suggestion}\"?")
+                    } else {
+                        String::new()
+                    }
+                ));
+
+                diagnostic
+            }
+            _ => builder.into_diagnostic(format_args!(
+                "TypedDict `{}` cannot be indexed with a key of type `{}`",
+                value_ty.display(db),
+                slice_ty.display(db),
+            )),
+        };
+    }
+}
+
 /// This function receives an unresolved `from foo import bar` import,
 /// where `foo` can be resolved to a module but that module does not
 /// have a `bar` member or submodule.
@@ -2619,7 +2667,7 @@ pub(super) fn hint_if_stdlib_submodule_exists_on_other_versions(
 }
 
 /// Suggest a name from `existing_names` that is similar to `wrong_name`.
-pub(super) fn did_you_mean<S: AsRef<str>, T: AsRef<str>>(
+fn did_you_mean<S: AsRef<str>, T: AsRef<str>>(
     existing_names: impl Iterator<Item = S>,
     wrong_name: T,
 ) -> Option<String> {