[ty] Faster subscript assignment checks for (unions of) TypedDicts (#21378)

## Summary

We synthesize a (potentially large) set of `__setitem__` overloads for
every item in a `TypedDict`. Previously, validation of subscript
assignments on `TypedDict`s relied on actually calling `__setitem__`
with the provided key and value types, which implied that we needed to
do the full overload call evaluation for this large set of overloads.
This PR improves the performance of subscript assignment checks on
`TypedDict`s by validating the assignment directly instead of calling
`__setitem__`.

This PR also adds better handling for assignments to subscripts on union
and intersection types (but does not attempt to make it perfect). It
achieves this by distributing the check over unions and intersections,
instead of calling `__setitem__` on the union/intersection directly. We
already do something similar when validating *attribute* assignments.

## Ecosystem impact

* A lot of diagnostics change their rule type, and/or split into
multiple diagnostics. The new version is more verbose, but easier to
understand, in my opinion
* Almost all of the invalid-key diagnostics come from pydantic, and they
should all go away (including many more) when we implement
https://github.com/astral-sh/ty/issues/1479
* Everything else looks correct to me. There may be some new diagnostics
due to the fact that we now check intersections.

## Test Plan

New Markdown tests.
This commit is contained in:
David Peter 2025-11-12 20:16:38 +01:00 committed by GitHub
parent 9dd666d677
commit 2f6f3e1042
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
15 changed files with 440 additions and 180 deletions

View file

@ -181,7 +181,7 @@ static PYDANTIC: Benchmark = Benchmark::new(
max_dep_date: "2025-06-17",
python_version: PythonVersion::PY39,
},
3000,
5000,
);
static SYMPY: Benchmark = Benchmark::new(

View file

@ -23,7 +23,7 @@ mdtest path: crates/ty_python_semantic/resources/mdtest/subscript/assignment_dia
# Diagnostics
```
error[invalid-assignment]: Cannot assign to object of type `ReadOnlyDict` with no `__setitem__` method
error[invalid-assignment]: Cannot assign to a subscript on an object of type `ReadOnlyDict` with no `__setitem__` method
--> src/mdtest_snippet.py:6:1
|
5 | config = ReadOnlyDict()

View file

@ -13,19 +13,20 @@ mdtest path: crates/ty_python_semantic/resources/mdtest/subscript/assignment_dia
```
1 | def _(config: dict[str, int] | None) -> None:
2 | config["retries"] = 3 # error: [possibly-missing-implicit-call]
2 | config["retries"] = 3 # error: [invalid-assignment]
```
# Diagnostics
```
warning[possibly-missing-implicit-call]: Method `__setitem__` of type `dict[str, int] | None` may be missing
error[invalid-assignment]: Cannot assign to a subscript on an object of type `None` with no `__setitem__` method
--> src/mdtest_snippet.py:2:5
|
1 | def _(config: dict[str, int] | None) -> None:
2 | config["retries"] = 3 # error: [possibly-missing-implicit-call]
2 | config["retries"] = 3 # error: [invalid-assignment]
| ^^^^^^
|
info: rule `possibly-missing-implicit-call` is enabled by default
info: The full type of the subscripted object is `dict[str, int] | None`
info: rule `invalid-assignment` is enabled by default
```

View file

@ -22,19 +22,39 @@ mdtest path: crates/ty_python_semantic/resources/mdtest/subscript/assignment_dia
8 | legs: int
9 |
10 | def _(being: Person | Animal) -> None:
11 | being["surname"] = "unknown" # error: [invalid-assignment]
11 | # error: [invalid-key]
12 | # error: [invalid-key]
13 | being["surname"] = "unknown"
```
# Diagnostics
```
error[invalid-assignment]: Method `__setitem__` of type `(key: Literal["name"], value: str, /) -> None` cannot be called with a key of type `Literal["surname"]` and a value of type `Literal["unknown"]` on object of type `Person | Animal`
--> src/mdtest_snippet.py:11:5
error[invalid-key]: Invalid key for TypedDict `Person`
--> src/mdtest_snippet.py:13:5
|
10 | def _(being: Person | Animal) -> None:
11 | being["surname"] = "unknown" # error: [invalid-assignment]
| ^^^^^
11 | # error: [invalid-key]
12 | # error: [invalid-key]
13 | being["surname"] = "unknown"
| ----- ^^^^^^^^^ Unknown key "surname" - did you mean "name"?
| |
| TypedDict `Person` in union type `Person | Animal`
|
info: rule `invalid-assignment` is enabled by default
info: rule `invalid-key` is enabled by default
```
```
error[invalid-key]: Invalid key for TypedDict `Animal`
--> src/mdtest_snippet.py:13:5
|
11 | # error: [invalid-key]
12 | # error: [invalid-key]
13 | being["surname"] = "unknown"
| ----- ^^^^^^^^^ Unknown key "surname" - did you mean "name"?
| |
| TypedDict `Animal` in union type `Person | Animal`
|
info: rule `invalid-key` is enabled by default
```

View file

@ -22,19 +22,21 @@ mdtest path: crates/ty_python_semantic/resources/mdtest/subscript/assignment_dia
8 | legs: int
9 |
10 | def _(being: Person | Animal) -> None:
11 | being["legs"] = 4 # error: [invalid-assignment]
11 | being["legs"] = 4 # error: [invalid-key]
```
# Diagnostics
```
error[invalid-assignment]: Method `__setitem__` of type `(key: Literal["name"], value: str, /) -> None` cannot be called with a key of type `Literal["legs"]` and a value of type `Literal[4]` on object of type `Person | Animal`
error[invalid-key]: Invalid key for TypedDict `Person`
--> src/mdtest_snippet.py:11:5
|
10 | def _(being: Person | Animal) -> None:
11 | being["legs"] = 4 # error: [invalid-assignment]
| ^^^^^
11 | being["legs"] = 4 # error: [invalid-key]
| ----- ^^^^^^ Unknown key "legs"
| |
| TypedDict `Person` in union type `Person | Animal`
|
info: rule `invalid-assignment` is enabled by default
info: rule `invalid-key` is enabled by default
```

View file

@ -19,13 +19,14 @@ mdtest path: crates/ty_python_semantic/resources/mdtest/subscript/assignment_dia
# Diagnostics
```
error[invalid-assignment]: Method `__setitem__` of type `(bound method dict[str, int].__setitem__(key: str, value: int, /) -> None) | (bound method dict[str, str].__setitem__(key: str, value: str, /) -> None)` cannot be called with a key of type `Literal["retries"]` and a value of type `Literal[3]` on object of type `dict[str, int] | dict[str, str]`
error[invalid-assignment]: Method `__setitem__` of type `bound method dict[str, str].__setitem__(key: str, value: str, /) -> None` cannot be called with a key of type `Literal["retries"]` and a value of type `Literal[3]` on object of type `dict[str, str]`
--> src/mdtest_snippet.py:2:5
|
1 | def _(config: dict[str, int] | dict[str, str]) -> None:
2 | config["retries"] = 3 # error: [invalid-assignment]
| ^^^^^^
|
info: The full type of the subscripted object is `dict[str, int] | dict[str, str]`
info: rule `invalid-assignment` is enabled by default
```

View file

@ -13,19 +13,37 @@ mdtest path: crates/ty_python_semantic/resources/mdtest/subscript/assignment_dia
```
1 | def _(config: dict[str, int] | dict[str, str]) -> None:
2 | config["retries"] = 3.0 # error: [invalid-assignment]
2 | # error: [invalid-assignment]
3 | # error: [invalid-assignment]
4 | config["retries"] = 3.0
```
# Diagnostics
```
error[invalid-assignment]: Method `__setitem__` of type `(bound method dict[str, int].__setitem__(key: str, value: int, /) -> None) | (bound method dict[str, str].__setitem__(key: str, value: str, /) -> None)` cannot be called with a key of type `Literal["retries"]` and a value of type `float` on object of type `dict[str, int] | dict[str, str]`
--> src/mdtest_snippet.py:2:5
error[invalid-assignment]: Method `__setitem__` of type `bound method dict[str, int].__setitem__(key: str, value: int, /) -> None` cannot be called with a key of type `Literal["retries"]` and a value of type `float` on object of type `dict[str, int]`
--> src/mdtest_snippet.py:4:5
|
1 | def _(config: dict[str, int] | dict[str, str]) -> None:
2 | config["retries"] = 3.0 # error: [invalid-assignment]
2 | # error: [invalid-assignment]
3 | # error: [invalid-assignment]
4 | config["retries"] = 3.0
| ^^^^^^
|
info: The full type of the subscripted object is `dict[str, int] | dict[str, str]`
info: rule `invalid-assignment` is enabled by default
```
```
error[invalid-assignment]: Method `__setitem__` of type `bound method dict[str, str].__setitem__(key: str, value: str, /) -> None` cannot be called with a key of type `Literal["retries"]` and a value of type `float` on object of type `dict[str, str]`
--> src/mdtest_snippet.py:4:5
|
2 | # error: [invalid-assignment]
3 | # error: [invalid-assignment]
4 | config["retries"] = 3.0
| ^^^^^^
|
info: The full type of the subscripted object is `dict[str, int] | dict[str, str]`
info: rule `invalid-assignment` is enabled by default
```

View file

@ -89,7 +89,7 @@ info: rule `invalid-key` is enabled by default
```
```
error[invalid-key]: Invalid key for TypedDict `Person` of type `str`
error[invalid-key]: Invalid key of type `str` for TypedDict `Person`
--> src/mdtest_snippet.py:16:12
|
15 | def access_with_str_key(person: Person, str_key: str):

View file

@ -67,7 +67,7 @@ config["retries"] = 3 # error: [invalid-assignment]
```py
def _(config: dict[str, int] | None) -> None:
config["retries"] = 3 # error: [possibly-missing-implicit-call]
config["retries"] = 3 # error: [invalid-assignment]
```
## Unknown key for one element of a union
@ -83,7 +83,7 @@ class Animal(TypedDict):
legs: int
def _(being: Person | Animal) -> None:
being["legs"] = 4 # error: [invalid-assignment]
being["legs"] = 4 # error: [invalid-key]
```
## Unknown key for all elemens of a union
@ -99,7 +99,9 @@ class Animal(TypedDict):
legs: int
def _(being: Person | Animal) -> None:
being["surname"] = "unknown" # error: [invalid-assignment]
# error: [invalid-key]
# error: [invalid-key]
being["surname"] = "unknown"
```
## Wrong value type for one element of a union
@ -113,5 +115,7 @@ def _(config: dict[str, int] | dict[str, str]) -> None:
```py
def _(config: dict[str, int] | dict[str, str]) -> None:
config["retries"] = 3.0 # error: [invalid-assignment]
# error: [invalid-assignment]
# error: [invalid-assignment]
config["retries"] = 3.0
```

View file

@ -76,7 +76,7 @@ a[0] = 0
class NoSetitem: ...
a = NoSetitem()
a[0] = 0 # error: "Cannot assign to object of type `NoSetitem` with no `__setitem__` method"
a[0] = 0 # error: "Cannot assign to a subscript on an object of type `NoSetitem` with no `__setitem__` method"
```
## `__setitem__` not callable

View file

@ -69,7 +69,7 @@ def name_or_age() -> Literal["name", "age"]:
carol: Person = {NAME: "Carol", AGE: 20}
reveal_type(carol[NAME]) # revealed: str
# error: [invalid-key] "Invalid key for TypedDict `Person` of type `str`"
# error: [invalid-key] "Invalid key of type `str` for TypedDict `Person`"
reveal_type(carol[non_literal()]) # revealed: Unknown
reveal_type(carol[name_or_age()]) # revealed: str | int | None
@ -553,7 +553,7 @@ def _(
# error: [invalid-key] "Invalid key for TypedDict `Person`: Unknown key "non_existing""
reveal_type(person["non_existing"]) # revealed: Unknown
# error: [invalid-key] "Invalid key for TypedDict `Person` of type `str`"
# error: [invalid-key] "Invalid key of type `str` for TypedDict `Person`"
reveal_type(person[str_key]) # revealed: Unknown
# No error here:
@ -602,16 +602,18 @@ def _(person: Person, literal_key: Literal["age"]):
def _(person: Person, union_of_keys: Literal["name", "surname"]):
person[union_of_keys] = "unknown"
# error: [invalid-assignment] "Cannot assign value of type `Literal[1]` to key of type `Literal["name", "surname"]` on TypedDict `Person`"
# error: [invalid-assignment] "Invalid assignment to key "name" with declared type `str` on TypedDict `Person`: value of type `Literal[1]`"
# error: [invalid-assignment] "Invalid assignment to key "surname" with declared type `str` on TypedDict `Person`: value of type `Literal[1]`"
person[union_of_keys] = 1
def _(being: Person | Animal):
being["name"] = "Being"
# error: [invalid-assignment] "Method `__setitem__` of type `(Overload[(key: Literal["name"], value: str, /) -> None, (key: Literal["surname"], value: str, /) -> None, (key: Literal["age"], value: int | None, /) -> None]) | (Overload[(key: Literal["name"], value: str, /) -> None, (key: Literal["legs"], value: int, /) -> None])` cannot be called with a key of type `Literal["name"]` and a value of type `Literal[1]` on object of type `Person | Animal`"
# error: [invalid-assignment] "Invalid assignment to key "name" with declared type `str` on TypedDict `Person`: value of type `Literal[1]`"
# error: [invalid-assignment] "Invalid assignment to key "name" with declared type `str` on TypedDict `Animal`: value of type `Literal[1]`"
being["name"] = 1
# error: [invalid-assignment] "Method `__setitem__` of type `(Overload[(key: Literal["name"], value: str, /) -> None, (key: Literal["surname"], value: str, /) -> None, (key: Literal["age"], value: int | None, /) -> None]) | (Overload[(key: Literal["name"], value: str, /) -> None, (key: Literal["legs"], value: int, /) -> None])` cannot be called with a key of type `Literal["surname"]` and a value of type `Literal["unknown"]` on object of type `Person | Animal`"
# error: [invalid-key] "Invalid key for TypedDict `Animal`: Unknown key "surname" - did you mean "name"?"
being["surname"] = "unknown"
def _(centaur: Intersection[Person, Animal]):
@ -619,13 +621,13 @@ def _(centaur: Intersection[Person, Animal]):
centaur["age"] = 100
centaur["legs"] = 4
# TODO: This should be an `invalid-key` error
# error: [invalid-key] "Invalid key for TypedDict `Person`: Unknown key "unknown""
centaur["unknown"] = "value"
def _(person: Person, union_of_keys: Literal["name", "age"], unknown_value: Any):
person[union_of_keys] = unknown_value
# error: [invalid-assignment] "Cannot assign value of type `None` to key of type `Literal["name", "age"]` on TypedDict `Person`"
# error: [invalid-assignment] "Invalid assignment to key "name" with declared type `str` on TypedDict `Person`: value of type `None`"
person[union_of_keys] = None
def _(person: Person, str_key: str, literalstr_key: LiteralString):

View file

@ -1163,6 +1163,10 @@ impl<'db> Type<'db> {
}
}
pub(crate) const fn is_union(&self) -> bool {
matches!(self, Type::Union(_))
}
pub(crate) const fn as_union(self) -> Option<UnionType<'db>> {
match self {
Type::Union(union_type) => Some(union_type),

View file

@ -3063,6 +3063,7 @@ pub(crate) fn report_invalid_key_on_typed_dict<'db>(
typed_dict_node: AnyNodeRef,
key_node: AnyNodeRef,
typed_dict_ty: Type<'db>,
full_object_ty: Option<Type<'db>>,
key_ty: Type<'db>,
items: &FxOrderMap<Name, Field<'db>>,
) {
@ -3077,11 +3078,21 @@ pub(crate) fn report_invalid_key_on_typed_dict<'db>(
"Invalid key for TypedDict `{typed_dict_name}`",
));
diagnostic.annotate(
diagnostic.annotate(if let Some(full_object_ty) = full_object_ty {
context.secondary(typed_dict_node).message(format_args!(
"TypedDict `{typed_dict_name}` in {kind} type `{full_object_ty}`",
kind = if full_object_ty.is_union() {
"union"
} else {
"intersection"
},
full_object_ty = full_object_ty.display(db)
))
} else {
context
.secondary(typed_dict_node)
.message(format_args!("TypedDict `{typed_dict_name}`")),
);
.message(format_args!("TypedDict `{typed_dict_name}`"))
});
let existing_keys = items.iter().map(|(name, _)| name.as_str());
@ -3093,15 +3104,22 @@ pub(crate) fn report_invalid_key_on_typed_dict<'db>(
String::new()
}
));
diagnostic
}
_ => builder.into_diagnostic(format_args!(
"Invalid key for TypedDict `{}` of type `{}`",
typed_dict_ty.display(db),
key_ty.display(db),
)),
};
_ => {
let mut diagnostic = builder.into_diagnostic(format_args!(
"Invalid key of type `{}` for TypedDict `{}`",
key_ty.display(db),
typed_dict_ty.display(db),
));
if let Some(full_object_ty) = full_object_ty {
diagnostic.info(format_args!(
"The full type of the subscripted object is `{}`",
full_object_ty.display(db)
));
}
}
}
}
}

View file

@ -102,10 +102,10 @@ use crate::types::visitor::any_over_type;
use crate::types::{
CallDunderError, CallableBinding, CallableType, ClassLiteral, ClassType, DataclassParams,
DynamicType, InferredAs, InternedType, InternedTypes, IntersectionBuilder, IntersectionType,
KnownClass, KnownInstanceType, MemberLookupPolicy, MetaclassCandidate, PEP695TypeAliasType,
Parameter, ParameterForm, Parameters, SpecialFormType, SubclassOfType, TrackedConstraintSet,
Truthiness, Type, TypeAliasType, TypeAndQualifiers, TypeContext, TypeQualifiers,
TypeVarBoundOrConstraintsEvaluation, TypeVarDefaultEvaluation, TypeVarIdentity,
KnownClass, KnownInstanceType, LintDiagnosticGuard, MemberLookupPolicy, MetaclassCandidate,
PEP695TypeAliasType, Parameter, ParameterForm, Parameters, SpecialFormType, SubclassOfType,
TrackedConstraintSet, Truthiness, Type, TypeAliasType, TypeAndQualifiers, TypeContext,
TypeQualifiers, TypeVarBoundOrConstraintsEvaluation, TypeVarDefaultEvaluation, TypeVarIdentity,
TypeVarInstance, TypeVarKind, TypeVarVariance, TypedDictType, UnionBuilder, UnionType,
binding_type, todo_type,
};
@ -3538,142 +3538,305 @@ impl<'db, 'ast> TypeInferenceBuilder<'db, 'ast> {
}
}
/// Make sure that the subscript assignment `obj[slice] = value` is valid.
/// Validate a subscript assignment of the form `object[key] = rhs_value`.
fn validate_subscript_assignment(
&mut self,
target: &ast::ExprSubscript,
rhs: &ast::Expr,
assigned_ty: Type<'db>,
rhs_value: &ast::Expr,
rhs_value_ty: Type<'db>,
) -> bool {
let ast::ExprSubscript {
range: _,
node_index: _,
value,
value: object,
slice,
ctx: _,
} = target;
let value_ty = self.infer_expression(value, TypeContext::default());
let object_ty = self.infer_expression(object, TypeContext::default());
let slice_ty = self.infer_expression(slice, TypeContext::default());
self.validate_subscript_assignment_impl(
object.as_ref(),
None,
object_ty,
slice.as_ref(),
slice_ty,
rhs_value,
rhs_value_ty,
true,
)
}
#[expect(clippy::too_many_arguments)]
fn validate_subscript_assignment_impl(
&self,
object_node: &'ast ast::Expr,
full_object_ty: Option<Type<'db>>,
object_ty: Type<'db>,
slice_node: &'ast ast::Expr,
slice_ty: Type<'db>,
rhs_value_node: &'ast ast::Expr,
rhs_value_ty: Type<'db>,
emit_diagnostic: bool,
) -> bool {
/// Given a string literal or a union of string literals, return an iterator over the contained
/// strings, or `None`, if the type is neither.
fn key_literals<'db>(
db: &'db dyn Db,
slice_ty: Type<'db>,
) -> Option<impl Iterator<Item = &'db str> + 'db> {
if let Some(literal) = slice_ty.as_string_literal() {
Some(Either::Left(std::iter::once(literal.value(db))))
} else {
slice_ty.as_union().map(|union| {
Either::Right(
union
.elements(db)
.iter()
.filter_map(|ty| ty.as_string_literal().map(|lit| lit.value(db))),
)
})
}
}
let db = self.db();
let context = &self.context;
match value_ty.try_call_dunder(
db,
"__setitem__",
CallArguments::positional([slice_ty, assigned_ty]),
TypeContext::default(),
) {
Ok(_) => true,
Err(err) => match err {
CallDunderError::PossiblyUnbound { .. } => {
if let Some(builder) =
context.report_lint(&POSSIBLY_MISSING_IMPLICIT_CALL, &**value)
{
builder.into_diagnostic(format_args!(
"Method `__setitem__` of type `{}` may be missing",
value_ty.display(db),
));
}
false
let attach_original_type_info = |mut diagnostic: LintDiagnosticGuard| {
if let Some(full_object_ty) = full_object_ty {
diagnostic.info(format_args!(
"The full type of the subscripted object is `{}`",
full_object_ty.display(db)
));
}
};
match object_ty {
Type::Union(union) => {
// Note that we use a loop here instead of .all(…) to avoid short-circuiting.
// We need to keep iterating to emit all diagnostics.
let mut valid = true;
for element_ty in union.elements(db) {
valid &= self.validate_subscript_assignment_impl(
object_node,
full_object_ty.or(Some(object_ty)),
*element_ty,
slice_node,
slice_ty,
rhs_value_node,
rhs_value_ty,
emit_diagnostic,
);
}
CallDunderError::CallError(call_error_kind, bindings) => {
match call_error_kind {
CallErrorKind::NotCallable => {
if let Some(builder) = context.report_lint(&CALL_NON_CALLABLE, &**value)
{
builder.into_diagnostic(format_args!(
"Method `__setitem__` of type `{}` is not callable \
on object of type `{}`",
bindings.callable_type().display(db),
value_ty.display(db),
));
}
}
CallErrorKind::BindingError => {
let assigned_d = assigned_ty.display(db);
let value_d = value_ty.display(db);
valid
}
if let Some(typed_dict) = value_ty.as_typed_dict() {
if let Some(key) = slice_ty.as_string_literal() {
let key = key.value(self.db());
validate_typed_dict_key_assignment(
&self.context,
typed_dict,
key,
assigned_ty,
value.as_ref(),
slice.as_ref(),
rhs,
TypedDictAssignmentKind::Subscript,
);
} 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)
Type::Intersection(intersection) => {
let check_positive_elements = |emit_diagnostic_and_short_circuit| {
let mut valid = false;
for element_ty in intersection.positive(db) {
valid |= self.validate_subscript_assignment_impl(
object_node,
full_object_ty.or(Some(object_ty)),
*element_ty,
slice_node,
slice_ty,
rhs_value_node,
rhs_value_ty,
emit_diagnostic_and_short_circuit,
);
if !valid && emit_diagnostic_and_short_circuit {
break;
}
}
valid
};
// Perform an initial check of all elements. If the assignment is valid
// for at least one element, we do not emit any diagnostics. Otherwise,
// we re-run the check and emit a diagnostic on the first failing element.
let valid = check_positive_elements(false);
if !valid {
check_positive_elements(true);
}
valid
}
Type::TypedDict(typed_dict) => {
// As an optimization, prevent calling `__setitem__` on (unions of) large `TypedDict`s, and
// validate the assignment ourselves. This also allows us to emit better diagnostics.
let mut valid = true;
let Some(keys) = key_literals(db, slice_ty) 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_dynamic() {
return true;
}
let assigned_d = rhs_value_ty.display(db);
let value_d = object_ty.display(db);
if slice_ty.is_assignable_to(db, Type::LiteralString)
&& !slice_ty.is_equivalent_to(db, Type::LiteralString)
{
if let Some(builder) =
self.context.report_lint(&INVALID_ASSIGNMENT, slice_node)
{
let diagnostic = builder.into_diagnostic(format_args!(
"Cannot assign value of type `{assigned_d}` to key of type `{}` on TypedDict `{value_d}`",
slice_ty.display(db)
));
attach_original_type_info(diagnostic);
}
} else {
if let Some(builder) = self.context.report_lint(&INVALID_KEY, slice_node) {
let diagnostic = 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)
));
attach_original_type_info(diagnostic);
}
}
return false;
};
for key in keys {
valid &= validate_typed_dict_key_assignment(
&self.context,
typed_dict,
full_object_ty,
key,
rhs_value_ty,
object_node,
slice_node,
rhs_value_node,
TypedDictAssignmentKind::Subscript,
emit_diagnostic,
);
}
valid
}
_ => {
match object_ty.try_call_dunder(
db,
"__setitem__",
CallArguments::positional([slice_ty, rhs_value_ty]),
TypeContext::default(),
) {
Ok(_) => true,
Err(err) => match err {
CallDunderError::PossiblyUnbound { .. } => {
if emit_diagnostic
&& let Some(builder) = self
.context
.report_lint(&POSSIBLY_MISSING_IMPLICIT_CALL, rhs_value_node)
{
let diagnostic = builder.into_diagnostic(format_args!(
"Method `__setitem__` of type `{}` may be missing",
object_ty.display(db),
));
attach_original_type_info(diagnostic);
}
false
}
CallDunderError::CallError(call_error_kind, bindings) => {
match call_error_kind {
CallErrorKind::NotCallable => {
if emit_diagnostic
&& let Some(builder) = self
.context
.report_lint(&CALL_NON_CALLABLE, object_node)
{
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)
));
let diagnostic = builder.into_diagnostic(format_args!(
"Method `__setitem__` of type `{}` is not callable \
on object of type `{}`",
bindings.callable_type().display(db),
object_ty.display(db),
));
attach_original_type_info(diagnostic);
}
}
CallErrorKind::BindingError => {
if let Some(typed_dict) = object_ty.as_typed_dict() {
if let Some(key) = slice_ty.as_string_literal() {
let key = key.value(db);
validate_typed_dict_key_assignment(
&self.context,
typed_dict,
full_object_ty,
key,
rhs_value_ty,
object_node,
slice_node,
rhs_value_node,
TypedDictAssignmentKind::Subscript,
true,
);
}
} else {
if let Some(builder) =
context.report_lint(&INVALID_KEY, &**slice)
if emit_diagnostic
&& let Some(builder) = self
.context
.report_lint(&INVALID_ASSIGNMENT, object_node)
{
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)
let assigned_d = rhs_value_ty.display(db);
let value_d = object_ty.display(db);
let diagnostic = 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),
));
attach_original_type_info(diagnostic);
}
}
}
} 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 => {
if emit_diagnostic
&& let Some(builder) = self
.context
.report_lint(&CALL_NON_CALLABLE, object_node)
{
let diagnostic = builder.into_diagnostic(format_args!(
"Method `__setitem__` of type `{}` may not be callable on object of type `{}`",
bindings.callable_type().display(db),
object_ty.display(db),
));
attach_original_type_info(diagnostic);
}
}
}
false
}
CallErrorKind::PossiblyNotCallable => {
if let Some(builder) = context.report_lint(&CALL_NON_CALLABLE, &**value)
CallDunderError::MethodNotAvailable => {
if emit_diagnostic
&& let Some(builder) =
self.context.report_lint(&INVALID_ASSIGNMENT, object_node)
{
builder.into_diagnostic(format_args!(
"Method `__setitem__` of type `{}` may not be \
callable on object of type `{}`",
bindings.callable_type().display(db),
value_ty.display(db),
let diagnostic = builder.into_diagnostic(format_args!(
"Cannot assign to a subscript on an object of type `{}` with no `__setitem__` method",
object_ty.display(db),
));
attach_original_type_info(diagnostic);
}
false
}
}
false
},
}
CallDunderError::MethodNotAvailable => {
if let Some(builder) = context.report_lint(&INVALID_ASSIGNMENT, &**value) {
builder.into_diagnostic(format_args!(
"Cannot assign to object of type `{}` with no `__setitem__` method",
value_ty.display(db),
));
}
false
}
},
}
}
}
@ -7682,6 +7845,7 @@ impl<'db, 'ast> TypeInferenceBuilder<'db, 'ast> {
first_arg.into(),
first_arg.into(),
Type::TypedDict(typed_dict_ty),
None,
key_ty,
&items,
);
@ -10908,6 +11072,7 @@ impl<'db, 'ast> TypeInferenceBuilder<'db, 'ast> {
value_node.into(),
slice_node.into(),
value_ty,
None,
slice_ty,
&typed_dict.items(db),
);

View file

@ -143,30 +143,57 @@ impl TypedDictAssignmentKind {
pub(super) fn validate_typed_dict_key_assignment<'db, 'ast>(
context: &InferContext<'db, 'ast>,
typed_dict: TypedDictType<'db>,
full_object_ty: Option<Type<'db>>,
key: &str,
value_ty: Type<'db>,
typed_dict_node: impl Into<AnyNodeRef<'ast>>,
typed_dict_node: impl Into<AnyNodeRef<'ast>> + Copy,
key_node: impl Into<AnyNodeRef<'ast>>,
value_node: impl Into<AnyNodeRef<'ast>>,
assignment_kind: TypedDictAssignmentKind,
emit_diagnostic: bool,
) -> bool {
let db = context.db();
let items = typed_dict.items(db);
// Check if key exists in `TypedDict`
let Some((_, item)) = items.iter().find(|(name, _)| *name == key) else {
report_invalid_key_on_typed_dict(
context,
typed_dict_node.into(),
key_node.into(),
Type::TypedDict(typed_dict),
Type::string_literal(db, key),
&items,
);
if emit_diagnostic {
report_invalid_key_on_typed_dict(
context,
typed_dict_node.into(),
key_node.into(),
Type::TypedDict(typed_dict),
full_object_ty,
Type::string_literal(db, key),
&items,
);
}
return false;
};
let add_object_type_annotation =
|diagnostic: &mut Diagnostic| {
if let Some(full_object_ty) = full_object_ty {
diagnostic.annotate(context.secondary(typed_dict_node.into()).message(
format_args!(
"TypedDict `{}` in {kind} type `{}`",
Type::TypedDict(typed_dict).display(db),
full_object_ty.display(db),
kind = if full_object_ty.is_union() {
"union"
} else {
"intersection"
},
),
));
} else {
diagnostic.annotate(context.secondary(typed_dict_node.into()).message(
format_args!("TypedDict `{}`", Type::TypedDict(typed_dict).display(db)),
));
}
};
let add_item_definition_subdiagnostic = |diagnostic: &mut Diagnostic, message| {
if let Some(declaration) = item.single_declaration {
let file = declaration.file(db);
@ -184,8 +211,9 @@ pub(super) fn validate_typed_dict_key_assignment<'db, 'ast>(
};
if assignment_kind.is_subscript() && item.is_read_only() {
if let Some(builder) =
context.report_lint(assignment_kind.diagnostic_type(), key_node.into())
if emit_diagnostic
&& let Some(builder) =
context.report_lint(assignment_kind.diagnostic_type(), key_node.into())
{
let typed_dict_ty = Type::TypedDict(typed_dict);
let typed_dict_d = typed_dict_ty.display(db);
@ -195,13 +223,7 @@ pub(super) fn validate_typed_dict_key_assignment<'db, 'ast>(
));
diagnostic.set_primary_message(format_args!("key is marked read-only"));
diagnostic.annotate(
context
.secondary(typed_dict_node.into())
.message(format_args!("TypedDict `{typed_dict_d}`")),
);
add_object_type_annotation(&mut diagnostic);
add_item_definition_subdiagnostic(&mut diagnostic, "Read-only item declared here");
}
@ -219,7 +241,9 @@ pub(super) fn validate_typed_dict_key_assignment<'db, 'ast>(
}
// Invalid assignment - emit diagnostic
if let Some(builder) = context.report_lint(assignment_kind.diagnostic_type(), value_node) {
if emit_diagnostic
&& let Some(builder) = context.report_lint(assignment_kind.diagnostic_type(), value_node)
{
let typed_dict_ty = Type::TypedDict(typed_dict);
let typed_dict_d = typed_dict_ty.display(db);
let value_d = value_ty.display(db);
@ -232,12 +256,6 @@ pub(super) fn validate_typed_dict_key_assignment<'db, 'ast>(
diagnostic.set_primary_message(format_args!("value of type `{value_d}`"));
diagnostic.annotate(
context
.secondary(typed_dict_node.into())
.message(format_args!("TypedDict `{typed_dict_d}`")),
);
diagnostic.annotate(
context
.secondary(key_node.into())
@ -245,6 +263,7 @@ pub(super) fn validate_typed_dict_key_assignment<'db, 'ast>(
);
add_item_definition_subdiagnostic(&mut diagnostic, "Item declared here");
add_object_type_annotation(&mut diagnostic);
}
false
@ -343,12 +362,14 @@ fn validate_from_dict_literal<'db, 'ast>(
validate_typed_dict_key_assignment(
context,
typed_dict,
None,
key_str,
value_type,
error_node,
key_expr,
&dict_item.value,
TypedDictAssignmentKind::Constructor,
true,
);
}
}
@ -380,12 +401,14 @@ fn validate_from_keywords<'db, 'ast>(
validate_typed_dict_key_assignment(
context,
typed_dict,
None,
arg_name.as_str(),
arg_type,
error_node,
keyword,
&keyword.value,
TypedDictAssignmentKind::Constructor,
true,
);
}
}
@ -418,12 +441,14 @@ pub(super) fn validate_typed_dict_dict_literal<'db>(
valid &= validate_typed_dict_key_assignment(
context,
typed_dict,
None,
key_str,
value_type,
error_node,
key_expr,
&item.value,
TypedDictAssignmentKind::Constructor,
true,
);
}
}