From 2f6f3e10425515f6c13208efb377fd398341d63e Mon Sep 17 00:00:00 2001 From: David Peter Date: Wed, 12 Nov 2025 20:16:38 +0100 Subject: [PATCH] [ty] Faster subscript assignment checks for (unions of) `TypedDict`s (#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. --- crates/ruff_benchmark/benches/ty_walltime.rs | 2 +- ..._No_`__setitem__`_met…_(468f62a3bdd1d60c).snap | 2 +- ..._Possibly_missing_`__…_(efd3f0c02e9b89e9).snap | 9 +- ..._Unknown_key_for_all_…_(1c685d9d10678263).snap | 34 +- ..._Unknown_key_for_one_…_(b515711c0a451a86).snap | 12 +- ..._Wrong_value_type_for…_(57372b65e30392a8).snap | 3 +- ..._Wrong_value_type_for…_(ffe39a3bae68cfe4).snap | 28 +- ...ict`_-_Diagnostics_(e5289abf5c570c29).snap | 2 +- .../subscript/assignment_diagnostics.md | 12 +- .../resources/mdtest/subscript/instance.md | 2 +- .../resources/mdtest/typed_dict.md | 16 +- crates/ty_python_semantic/src/types.rs | 4 + .../src/types/diagnostic.rs | 40 +- .../src/types/infer/builder.rs | 379 +++++++++++++----- .../src/types/typed_dict.rs | 75 ++-- 15 files changed, 440 insertions(+), 180 deletions(-) diff --git a/crates/ruff_benchmark/benches/ty_walltime.rs b/crates/ruff_benchmark/benches/ty_walltime.rs index 8f13ab7ca7..697a0c989d 100644 --- a/crates/ruff_benchmark/benches/ty_walltime.rs +++ b/crates/ruff_benchmark/benches/ty_walltime.rs @@ -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( diff --git a/crates/ty_python_semantic/resources/mdtest/snapshots/assignment_diagnosti…_-_Subscript_assignment…_-_No_`__setitem__`_met…_(468f62a3bdd1d60c).snap b/crates/ty_python_semantic/resources/mdtest/snapshots/assignment_diagnosti…_-_Subscript_assignment…_-_No_`__setitem__`_met…_(468f62a3bdd1d60c).snap index dec0ab3417..f9c43e5882 100644 --- a/crates/ty_python_semantic/resources/mdtest/snapshots/assignment_diagnosti…_-_Subscript_assignment…_-_No_`__setitem__`_met…_(468f62a3bdd1d60c).snap +++ b/crates/ty_python_semantic/resources/mdtest/snapshots/assignment_diagnosti…_-_Subscript_assignment…_-_No_`__setitem__`_met…_(468f62a3bdd1d60c).snap @@ -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() diff --git a/crates/ty_python_semantic/resources/mdtest/snapshots/assignment_diagnosti…_-_Subscript_assignment…_-_Possibly_missing_`__…_(efd3f0c02e9b89e9).snap b/crates/ty_python_semantic/resources/mdtest/snapshots/assignment_diagnosti…_-_Subscript_assignment…_-_Possibly_missing_`__…_(efd3f0c02e9b89e9).snap index ced810cf72..a12bb7c666 100644 --- a/crates/ty_python_semantic/resources/mdtest/snapshots/assignment_diagnosti…_-_Subscript_assignment…_-_Possibly_missing_`__…_(efd3f0c02e9b89e9).snap +++ b/crates/ty_python_semantic/resources/mdtest/snapshots/assignment_diagnosti…_-_Subscript_assignment…_-_Possibly_missing_`__…_(efd3f0c02e9b89e9).snap @@ -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 ``` diff --git a/crates/ty_python_semantic/resources/mdtest/snapshots/assignment_diagnosti…_-_Subscript_assignment…_-_Unknown_key_for_all_…_(1c685d9d10678263).snap b/crates/ty_python_semantic/resources/mdtest/snapshots/assignment_diagnosti…_-_Subscript_assignment…_-_Unknown_key_for_all_…_(1c685d9d10678263).snap index 6444c84f36..2e7bbcfe4d 100644 --- a/crates/ty_python_semantic/resources/mdtest/snapshots/assignment_diagnosti…_-_Subscript_assignment…_-_Unknown_key_for_all_…_(1c685d9d10678263).snap +++ b/crates/ty_python_semantic/resources/mdtest/snapshots/assignment_diagnosti…_-_Subscript_assignment…_-_Unknown_key_for_all_…_(1c685d9d10678263).snap @@ -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 ``` diff --git a/crates/ty_python_semantic/resources/mdtest/snapshots/assignment_diagnosti…_-_Subscript_assignment…_-_Unknown_key_for_one_…_(b515711c0a451a86).snap b/crates/ty_python_semantic/resources/mdtest/snapshots/assignment_diagnosti…_-_Subscript_assignment…_-_Unknown_key_for_one_…_(b515711c0a451a86).snap index 2b840a6783..6c919e6937 100644 --- a/crates/ty_python_semantic/resources/mdtest/snapshots/assignment_diagnosti…_-_Subscript_assignment…_-_Unknown_key_for_one_…_(b515711c0a451a86).snap +++ b/crates/ty_python_semantic/resources/mdtest/snapshots/assignment_diagnosti…_-_Subscript_assignment…_-_Unknown_key_for_one_…_(b515711c0a451a86).snap @@ -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 ``` diff --git a/crates/ty_python_semantic/resources/mdtest/snapshots/assignment_diagnosti…_-_Subscript_assignment…_-_Wrong_value_type_for…_(57372b65e30392a8).snap b/crates/ty_python_semantic/resources/mdtest/snapshots/assignment_diagnosti…_-_Subscript_assignment…_-_Wrong_value_type_for…_(57372b65e30392a8).snap index 37ea1c111a..0f603931aa 100644 --- a/crates/ty_python_semantic/resources/mdtest/snapshots/assignment_diagnosti…_-_Subscript_assignment…_-_Wrong_value_type_for…_(57372b65e30392a8).snap +++ b/crates/ty_python_semantic/resources/mdtest/snapshots/assignment_diagnosti…_-_Subscript_assignment…_-_Wrong_value_type_for…_(57372b65e30392a8).snap @@ -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 ``` diff --git a/crates/ty_python_semantic/resources/mdtest/snapshots/assignment_diagnosti…_-_Subscript_assignment…_-_Wrong_value_type_for…_(ffe39a3bae68cfe4).snap b/crates/ty_python_semantic/resources/mdtest/snapshots/assignment_diagnosti…_-_Subscript_assignment…_-_Wrong_value_type_for…_(ffe39a3bae68cfe4).snap index dfd0136536..635a402c9b 100644 --- a/crates/ty_python_semantic/resources/mdtest/snapshots/assignment_diagnosti…_-_Subscript_assignment…_-_Wrong_value_type_for…_(ffe39a3bae68cfe4).snap +++ b/crates/ty_python_semantic/resources/mdtest/snapshots/assignment_diagnosti…_-_Subscript_assignment…_-_Wrong_value_type_for…_(ffe39a3bae68cfe4).snap @@ -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 ``` diff --git a/crates/ty_python_semantic/resources/mdtest/snapshots/typed_dict.md_-_`TypedDict`_-_Diagnostics_(e5289abf5c570c29).snap b/crates/ty_python_semantic/resources/mdtest/snapshots/typed_dict.md_-_`TypedDict`_-_Diagnostics_(e5289abf5c570c29).snap index a5b9456acd..51b0f0ce69 100644 --- a/crates/ty_python_semantic/resources/mdtest/snapshots/typed_dict.md_-_`TypedDict`_-_Diagnostics_(e5289abf5c570c29).snap +++ b/crates/ty_python_semantic/resources/mdtest/snapshots/typed_dict.md_-_`TypedDict`_-_Diagnostics_(e5289abf5c570c29).snap @@ -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): diff --git a/crates/ty_python_semantic/resources/mdtest/subscript/assignment_diagnostics.md b/crates/ty_python_semantic/resources/mdtest/subscript/assignment_diagnostics.md index ed23208eb1..e4959e3627 100644 --- a/crates/ty_python_semantic/resources/mdtest/subscript/assignment_diagnostics.md +++ b/crates/ty_python_semantic/resources/mdtest/subscript/assignment_diagnostics.md @@ -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 ``` diff --git a/crates/ty_python_semantic/resources/mdtest/subscript/instance.md b/crates/ty_python_semantic/resources/mdtest/subscript/instance.md index 7d1ad7f183..b15ec4abc1 100644 --- a/crates/ty_python_semantic/resources/mdtest/subscript/instance.md +++ b/crates/ty_python_semantic/resources/mdtest/subscript/instance.md @@ -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 diff --git a/crates/ty_python_semantic/resources/mdtest/typed_dict.md b/crates/ty_python_semantic/resources/mdtest/typed_dict.md index 8b8fcfffa3..422711b4c1 100644 --- a/crates/ty_python_semantic/resources/mdtest/typed_dict.md +++ b/crates/ty_python_semantic/resources/mdtest/typed_dict.md @@ -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): diff --git a/crates/ty_python_semantic/src/types.rs b/crates/ty_python_semantic/src/types.rs index b49b8e2d23..4284b15278 100644 --- a/crates/ty_python_semantic/src/types.rs +++ b/crates/ty_python_semantic/src/types.rs @@ -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> { match self { Type::Union(union_type) => Some(union_type), diff --git a/crates/ty_python_semantic/src/types/diagnostic.rs b/crates/ty_python_semantic/src/types/diagnostic.rs index ccb0c82472..25c1efa27d 100644 --- a/crates/ty_python_semantic/src/types/diagnostic.rs +++ b/crates/ty_python_semantic/src/types/diagnostic.rs @@ -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>, key_ty: Type<'db>, items: &FxOrderMap>, ) { @@ -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) + )); + } + } + } } } diff --git a/crates/ty_python_semantic/src/types/infer/builder.rs b/crates/ty_python_semantic/src/types/infer/builder.rs index 53c1c1dc96..2c445f92ea 100644 --- a/crates/ty_python_semantic/src/types/infer/builder.rs +++ b/crates/ty_python_semantic/src/types/infer/builder.rs @@ -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>, + 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 + '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), ); diff --git a/crates/ty_python_semantic/src/types/typed_dict.rs b/crates/ty_python_semantic/src/types/typed_dict.rs index 83b4ae946e..e07dbe6e60 100644 --- a/crates/ty_python_semantic/src/types/typed_dict.rs +++ b/crates/ty_python_semantic/src/types/typed_dict.rs @@ -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>, key: &str, value_ty: Type<'db>, - typed_dict_node: impl Into>, + typed_dict_node: impl Into> + Copy, key_node: impl Into>, value_node: impl Into>, 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, ); } }