Skip to content

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

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
sharkdp merged 4 commits into from
Nov 12, 2025
Merged

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

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
4 commits
Select commit Hold shift + click to select a range
726e465
[ty] Faster subscript assignment checks on `TypedDict`s
sharkdp Nov 11, 2025
6443797
Keep the overloads
sharkdp Nov 11, 2025
1299bc5
Always distribute over unions/intersections
sharkdp Nov 12, 2025
ccb7a0c
Add original type as a info hint
sharkdp Nov 12, 2025
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
2 changes: 1 addition & 1 deletion crates/ruff_benchmark/benches/ty_walltime.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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(
Expand Down
2 changes: 1 addition & 1 deletion ...gnment_diagnosti…_-_Subscript_assignment…_-_No_`__setitem__`_met…_(468f62a3bdd1d60c).snap
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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()
Expand Down
9 changes: 5 additions & 4 deletions ...gnment_diagnosti…_-_Subscript_assignment…_-_Possibly_missing_`__…_(efd3f0c02e9b89e9).snap
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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

```
34 changes: 27 additions & 7 deletions ...gnment_diagnosti…_-_Subscript_assignment…_-_Unknown_key_for_all_…_(1c685d9d10678263).snap
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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-key` is enabled by default
Comment on lines +33 to +43
Copy link
Member

@AlexWaygood AlexWaygood Nov 12, 2025

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

What about including the name of the unknown key in the summary line here and taking it out of the primary-annotation message?

error[invalid-key]: Invalid key "surname" for TypedDict `Person`
  --> src/mdtest_snippet.py:13:5
   |
11 |     # error: [invalid-key]
12 |     # error: [invalid-key]
13 |     being["surname"] = "unknown"
   |     ----- ^^^^^^^^^ Did you mean "name"?
   |     |
   |     TypedDict `Person` in union type `Person | Animal`
   |
info: rule `invalid-key` is enabled by default

We could also consider linking back to the definition of the Person typeddict in a subdiagnostic, though I guess that would be very verbose if you had a union of typeddicts.

Copy link
Contributor Author

@sharkdp sharkdp Nov 12, 2025

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

What about including the name of the unknown key in the summary line here and taking it out of the primary-annotation message?

I can try that, but it's always a bit tricky to get both the full diagnostic and the concise diagnostic to look good. I still wish we had an API to set a separate concise message.


```

```
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-assignment` is enabled by default
info: rule `invalid-key` is enabled by default

```
12 changes: 7 additions & 5 deletions ...gnment_diagnosti…_-_Subscript_assignment…_-_Unknown_key_for_one_…_(b515711c0a451a86).snap
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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

```
3 changes: 2 additions & 1 deletion ...gnment_diagnosti…_-_Subscript_assignment…_-_Wrong_value_type_for…_(57372b65e30392a8).snap
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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

```
28 changes: 23 additions & 5 deletions ...gnment_diagnosti…_-_Subscript_assignment…_-_Wrong_value_type_for…_(ffe39a3bae68cfe4).snap
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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

```
2 changes: 1 addition & 1 deletion ...ources/mdtest/snapshots/typed_dict.md_-_`TypedDict`_-_Diagnostics_(e5289abf5c570c29).snap
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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):
Expand Down
12 changes: 8 additions & 4 deletions crates/ty_python_semantic/resources/mdtest/subscript/assignment_diagnostics.md
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand All @@ -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
Expand All @@ -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
Expand All @@ -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
```
2 changes: 1 addition & 1 deletion crates/ty_python_semantic/resources/mdtest/subscript/instance.md
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down
16 changes: 9 additions & 7 deletions crates/ty_python_semantic/resources/mdtest/typed_dict.md
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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

Expand Down Expand Up @@ -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:
Expand Down Expand Up @@ -602,30 +602,32 @@ 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]):
centaur["name"] = "Chiron"
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):
Expand Down
4 changes: 4 additions & 0 deletions crates/ty_python_semantic/src/types.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -1164,6 +1164,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),
Expand Down
38 changes: 28 additions & 10 deletions crates/ty_python_semantic/src/types/diagnostic.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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>>,
) {
Expand All @@ -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());

Expand All @@ -3093,15 +3104,22 @@ pub(crate) fn report_invalid_key_on_typed_dict<'db>(
String::new()
}
));
}
_ => {
let mut diagnostic = builder.into_diagnostic(format_args!(
"Invalid key of type `{}` for TypedDict `{}`",
key_ty.display(db),
typed_dict_ty.display(db),
));

diagnostic
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)
));
}
}
_ => builder.into_diagnostic(format_args!(
"Invalid key for TypedDict `{}` of type `{}`",
typed_dict_ty.display(db),
key_ty.display(db),
)),
};
}
}
}

Expand Down
Loading
Loading