[ty] Validate writes to TypedDict keys (#19782)

## Summary

Validates writes to `TypedDict` keys, for example:

```py
class Person(TypedDict):
    name: str
    age: int | None


def f(person: Person):
    person["naem"] = "Alice"  # error: [invalid-key]

    person["age"] = "42"  # error: [invalid-assignment]
```

The new specialized `invalid-assignment` diagnostic looks like this:

<img width="1160" height="279" alt="image"
src="https://github.com/user-attachments/assets/51259455-3501-4829-a84e-df26ff90bd89"
/>

## Ecosystem analysis

As far as I can tell, all true positives!

There are some extremely long diagnostic messages. We should truncate
our display of overload sets somehow.

## Test Plan

New Markdown tests
This commit is contained in:
David Peter 2025-08-07 00:19:13 +02:00 committed by GitHub
parent 65b39f2ca9
commit 98df62db79
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
6 changed files with 365 additions and 136 deletions

View file

@ -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);

View file

@ -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> {

View file

@ -99,11 +99,11 @@ use crate::types::diagnostic::{
INVALID_TYPE_FORM, INVALID_TYPE_GUARD_CALL, INVALID_TYPE_VARIABLE_CONSTRAINTS,
IncompatibleBases, POSSIBLY_UNBOUND_IMPLICIT_CALL, POSSIBLY_UNBOUND_IMPORT,
TypeCheckDiagnostics, UNDEFINED_REVEAL, UNRESOLVED_ATTRIBUTE, UNRESOLVED_GLOBAL,
UNRESOLVED_IMPORT, UNRESOLVED_REFERENCE, UNSUPPORTED_OPERATOR, did_you_mean,
report_implicit_return_type, report_instance_layout_conflict,
report_invalid_argument_number_to_special_form, report_invalid_arguments_to_annotated,
report_invalid_arguments_to_callable, report_invalid_assignment,
report_invalid_attribute_assignment, report_invalid_generator_function_return_type,
UNRESOLVED_IMPORT, UNRESOLVED_REFERENCE, UNSUPPORTED_OPERATOR, report_implicit_return_type,
report_instance_layout_conflict, report_invalid_argument_number_to_special_form,
report_invalid_arguments_to_annotated, report_invalid_arguments_to_callable,
report_invalid_assignment, report_invalid_attribute_assignment,
report_invalid_generator_function_return_type, report_invalid_key_on_typed_dict,
report_invalid_return_type, report_possibly_unbound_attribute,
};
use crate::types::enums::is_enum_class;
@ -3702,13 +3702,14 @@ impl<'db, 'ast> TypeInferenceBuilder<'db, 'ast> {
ast::Expr::Name(_) => None,
_ => Some(infer_value_expr(self, value)),
};
self.infer_target_impl(target, assigned_ty);
self.infer_target_impl(target, value, assigned_ty);
}
/// Make sure that the subscript assignment `obj[slice] = value` is valid.
fn validate_subscript_assignment(
&mut self,
target: &ast::ExprSubscript,
rhs: &ast::Expr,
assigned_ty: Type<'db>,
) -> bool {
let ast::ExprSubscript {
@ -3757,17 +3758,92 @@ impl<'db, 'ast> TypeInferenceBuilder<'db, 'ast> {
}
}
CallErrorKind::BindingError => {
if let Some(builder) =
context.report_lint(&INVALID_ASSIGNMENT, &**value)
{
builder.into_diagnostic(format_args!(
"Method `__setitem__` of type `{}` cannot be called with \
a key of type `{}` and a value of type `{}` on object of type `{}`",
bindings.callable_type().display(db),
slice_ty.display(db),
assigned_ty.display(db),
value_ty.display(db),
));
let assigned_d = assigned_ty.display(db);
let value_d = value_ty.display(db);
if let Some(typed_dict) = value_ty.into_typed_dict() {
if let Some(key) = slice_ty.into_string_literal() {
let key = key.value(self.db());
let items = typed_dict.items(self.db());
if let Some((_, item)) =
items.iter().find(|(name, _)| *name == key)
{
if let Some(builder) =
context.report_lint(&INVALID_ASSIGNMENT, rhs)
{
let mut diagnostic = builder.into_diagnostic(format_args!(
"Invalid assignment to key \"{key}\" with declared type `{}` on TypedDict `{value_d}`",
item.declared_ty.display(db),
));
diagnostic.set_primary_message(format_args!(
"value of type `{assigned_d}`"
));
diagnostic.annotate(
self.context
.secondary(value.as_ref())
.message(format_args!("TypedDict `{value_d}`")),
);
diagnostic.annotate(
self.context.secondary(slice.as_ref()).message(
format_args!(
"key has declared type `{}`",
item.declared_ty.display(db),
),
),
);
}
} else {
report_invalid_key_on_typed_dict(
&self.context,
value.as_ref().into(),
slice.as_ref().into(),
value_ty,
slice_ty,
&items,
);
}
} else {
// Check if the key has a valid type. We only allow string literals, a union of string literals,
// or a dynamic type like `Any`. We can do this by checking assignability to `LiteralString`,
// but we need to exclude `LiteralString` itself. This check would technically allow weird key
// types like `LiteralString & Any` to pass, but it does not need to be perfect. We would just
// fail to provide the "Only string literals are allowed" hint in that case.
if slice_ty.is_assignable_to(db, Type::LiteralString)
&& !slice_ty.is_equivalent_to(db, Type::LiteralString)
{
if let Some(builder) =
context.report_lint(&INVALID_ASSIGNMENT, &**slice)
{
builder.into_diagnostic(format_args!(
"Cannot assign value of type `{assigned_d}` to key of type `{}` on TypedDict `{value_d}`",
slice_ty.display(db)
));
}
} else {
if let Some(builder) =
context.report_lint(&INVALID_KEY, &**slice)
{
builder.into_diagnostic(format_args!(
"Cannot access `{value_d}` with a key of type `{}`. Only string literals are allowed as keys on TypedDicts.",
slice_ty.display(db)
));
}
}
}
} else {
if let Some(builder) =
context.report_lint(&INVALID_ASSIGNMENT, &**value)
{
builder.into_diagnostic(format_args!(
"Method `__setitem__` of type `{}` cannot be called with \
a key of type `{}` and a value of type `{assigned_d}` on object of type `{value_d}`",
bindings.callable_type().display(db),
slice_ty.display(db),
));
}
}
}
CallErrorKind::PossiblyNotCallable => {
@ -4333,7 +4409,12 @@ impl<'db, 'ast> TypeInferenceBuilder<'db, 'ast> {
}
}
fn infer_target_impl(&mut self, target: &ast::Expr, assigned_ty: Option<Type<'db>>) {
fn infer_target_impl(
&mut self,
target: &ast::Expr,
value: &ast::Expr,
assigned_ty: Option<Type<'db>>,
) {
match target {
ast::Expr::Name(name) => self.infer_definition(name),
ast::Expr::List(ast::ExprList { elts, .. })
@ -4346,7 +4427,7 @@ impl<'db, 'ast> TypeInferenceBuilder<'db, 'ast> {
};
for element in elts {
self.infer_target_impl(element, assigned_tys.next());
self.infer_target_impl(element, value, assigned_tys.next());
}
}
ast::Expr::Attribute(
@ -4375,7 +4456,7 @@ impl<'db, 'ast> TypeInferenceBuilder<'db, 'ast> {
self.store_expression_type(target, assigned_ty.unwrap_or(Type::unknown()));
if let Some(assigned_ty) = assigned_ty {
self.validate_subscript_assignment(subscript_expr, assigned_ty);
self.validate_subscript_assignment(subscript_expr, value, assigned_ty);
}
}
_ => {
@ -8879,46 +8960,14 @@ impl<'db, 'ast> TypeInferenceBuilder<'db, 'ast> {
if let Some(typed_dict) = value_ty.into_typed_dict() {
let slice_node = subscript.slice.as_ref();
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(
self.context.secondary(value_node).message(
format_args!("TypedDict `{typed_dict_name}`"),
),
);
let items = typed_dict.items(db);
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),
)),
};
}
report_invalid_key_on_typed_dict(
context,
value_node.into(),
slice_node.into(),
value_ty,
slice_ty,
&typed_dict.items(db),
);
} else {
if let Some(builder) =
context.report_lint(&INVALID_ARGUMENT_TYPE, value_node)