[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

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