[ty] Subscript assignment diagnostics follow-up (#21452)

## Summary

Follow up from https://github.com/astral-sh/ruff/pull/21411. Again,
there are more things that could be improved here (like the diagnostics
for `lists`, or extending what we have for `dict` to `OrderedDict` etc),
but that will have to be postponed.

crates/ty_python_semantic/resources/mdtest/narrow/assignment.md

@@ -313,7 +313,7 @@ def f(c: C, s: str):
     reveal_type(c.x)  # revealed: int | None
     s = c.x  # error: [invalid-assignment]
 
-    # error: [invalid-assignment] "Method `__setitem__` of type `Overload[(key: SupportsIndex, value: int, /) -> None, (key: slice[Any, Any, Any], value: Iterable[int], /) -> None]` cannot be called with a key of type `Literal[0]` and a value of type `str` on object of type `list[int]`"
+    # error: [invalid-assignment] "Invalid subscript assignment with key of type `Literal[0]` and value of type `str` on object of type `list[int]`"
     c.l[0] = s
     reveal_type(c.l[0])  # revealed: int
 ```

crates/ty_python_semantic/resources/mdtest/snapshots/assignment_diagnosti…_-_Subscript_assignment…_-_Invalid_key_type_-_For_a_`dict`_(4aa9d1d82d07fcf1).snap

@@ -3,7 +3,7 @@ source: crates/ty_test/src/lib.rs
 expression: snapshot
 ---
 ---
-mdtest name: assignment_diagnostics.md - Subscript assignment diagnostics - Invalid key type
+mdtest name: assignment_diagnostics.md - Subscript assignment diagnostics - Invalid key type - For a `dict`
 mdtest path: crates/ty_python_semantic/resources/mdtest/subscript/assignment_diagnostics.md
 ---
 

crates/ty_python_semantic/resources/mdtest/snapshots/assignment_diagnosti…_-_Subscript_assignment…_-_Invalid_key_type_-_For_a_`list`_(752cfa73fb34c1c).snap

@@ -0,0 +1,31 @@
+---
+source: crates/ty_test/src/lib.rs
+expression: snapshot
+---
+---
+mdtest name: assignment_diagnostics.md - Subscript assignment diagnostics - Invalid key type - For a `list`
+mdtest path: crates/ty_python_semantic/resources/mdtest/subscript/assignment_diagnostics.md
+---
+
+# Python source files
+
+## mdtest_snippet.py
+
+```
+1 | numbers: list[int] = []
+2 | numbers["zero"] = 3  # error: [invalid-assignment]
+```
+
+# Diagnostics
+
+```
+error[invalid-assignment]: Invalid subscript assignment with key of type `Literal["zero"]` and value of type `Literal[3]` on object of type `list[int]`
+ --> src/mdtest_snippet.py:2:1
+  |
+1 | numbers: list[int] = []
+2 | numbers["zero"] = 3  # error: [invalid-assignment]
+  | ^^^^^^^^^^^^^^^^^^^
+  |
+info: rule `invalid-assignment` is enabled by default
+
+```

crates/ty_python_semantic/resources/mdtest/snapshots/assignment_diagnosti…_-_Subscript_assignment…_-_Invalid_value_type_-_For_a_`dict`_(177872afa1956fef).snap

@@ -3,7 +3,7 @@ source: crates/ty_test/src/lib.rs
 expression: snapshot
 ---
 ---
-mdtest name: assignment_diagnostics.md - Subscript assignment diagnostics - Invalid value type
+mdtest name: assignment_diagnostics.md - Subscript assignment diagnostics - Invalid value type - For a `dict`
 mdtest path: crates/ty_python_semantic/resources/mdtest/subscript/assignment_diagnostics.md
 ---
 

crates/ty_python_semantic/resources/mdtest/snapshots/assignment_diagnosti…_-_Subscript_assignment…_-_Invalid_value_type_-_For_a_`list`_(e7ebbd4af387837c).snap

@@ -0,0 +1,31 @@
+---
+source: crates/ty_test/src/lib.rs
+expression: snapshot
+---
+---
+mdtest name: assignment_diagnostics.md - Subscript assignment diagnostics - Invalid value type - For a `list`
+mdtest path: crates/ty_python_semantic/resources/mdtest/subscript/assignment_diagnostics.md
+---
+
+# Python source files
+
+## mdtest_snippet.py
+
+```
+1 | numbers: list[int] = []
+2 | numbers[0] = "three"  # error: [invalid-assignment]
+```
+
+# Diagnostics
+
+```
+error[invalid-assignment]: Invalid subscript assignment with key of type `Literal[0]` and value of type `Literal["three"]` on object of type `list[int]`
+ --> src/mdtest_snippet.py:2:1
+  |
+1 | numbers: list[int] = []
+2 | numbers[0] = "three"  # error: [invalid-assignment]
+  | ^^^^^^^^^^^^^^^^^^^^
+  |
+info: rule `invalid-assignment` is enabled by default
+
+```

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]: TypedDict `Person` can only be subscripted with string literal keys, got key of type `str`
+error[invalid-key]: TypedDict `Person` can only be subscripted with a string literal key, got key of type `str`
   --> src/mdtest_snippet.py:16:12
    |
 15 | def access_with_str_key(person: Person, str_key: str):

crates/ty_python_semantic/resources/mdtest/subscript/assignment_diagnostics.md

@@ -4,6 +4,15 @@
 
 ## Invalid value type
 
+### For a `list`
+
+```py
+numbers: list[int] = []
+numbers[0] = "three"  # error: [invalid-assignment]
+```
+
+### For a `dict`
+
 ```py
 config: dict[str, int] = {}
 config["retries"] = "three"  # error: [invalid-assignment]
@@ -11,6 +20,15 @@ config["retries"] = "three"  # error: [invalid-assignment]
 
 ## Invalid key type
 
+### For a `list`
+
+```py
+numbers: list[int] = []
+numbers["zero"] = 3  # error: [invalid-assignment]
+```
+
+### For a `dict`
+
 ```py
 config: dict[str, int] = {}
 config[0] = 3  # error: [invalid-assignment]

crates/ty_python_semantic/resources/mdtest/subscript/instance.md

@@ -110,6 +110,6 @@ class Identity:
         pass
 
 a = Identity()
-# error: [invalid-assignment] "Method `__setitem__` of type `bound method Identity.__setitem__(index: int, value: int) -> None` cannot be called with a key of type `Literal["a"]` and a value of type `Literal[0]` on object of type `Identity`"
+# error: [invalid-assignment] "Invalid subscript assignment with key of type `Literal["a"]` and value of type `Literal[0]` on object of type `Identity`"
 a["a"] = 0
 ```

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] "TypedDict `Person` can only be subscripted with string literal keys, got key of type `str`"
+# error: [invalid-key] "TypedDict `Person` can only be subscripted with a string literal key, got key of type `str`"
 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] "TypedDict `Person` can only be subscripted with string literal keys, got key of type `str`"
+    # error: [invalid-key] "TypedDict `Person` can only be subscripted with a string literal key, got key of type `str`"
     reveal_type(person[str_key])  # revealed: Unknown
 
     # No error here:

crates/ty_python_semantic/src/types/diagnostic.rs

@@ -3179,7 +3179,7 @@ pub(crate) fn report_invalid_key_on_typed_dict<'db>(
             }
             _ => {
                 let mut diagnostic = builder.into_diagnostic(format_args!(
-                    "TypedDict `{}` can only be subscripted with string literal keys, \
+                    "TypedDict `{}` can only be subscripted with a string literal key, \
                      got key of type `{}`",
                     typed_dict_ty.display(db),
                     key_ty.display(db),

crates/ty_python_semantic/src/types/infer/builder.rs

@@ -3691,7 +3691,7 @@ impl<'db, 'ast> TypeInferenceBuilder<'db, 'ast> {
                     // 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 "can only be subscripted with string literal keys" hint in that case.
+                    // fail to provide the "can only be subscripted with a string literal key" hint in that case.
 
                     if slice_ty.is_dynamic() {
                         return true;
@@ -3813,18 +3813,18 @@ impl<'db, 'ast> TypeInferenceBuilder<'db, 'ast> {
                                             let assigned_d = rhs_value_ty.display(db);
                                             let object_d = object_ty.display(db);
 
+                                            let mut diagnostic = builder.into_diagnostic(format_args!(
+                                                    "Invalid subscript assignment with key of type `{}` and value of \
+                                                     type `{assigned_d}` on object of type `{object_d}`",
+                                                    slice_ty.display(db),
+                                                ));
+
                                             // Special diagnostic for dictionaries
                                             if let Some([expected_key_ty, expected_value_ty]) =
                                                 object_ty
                                                     .known_specialization(db, KnownClass::Dict)
                                                     .map(|s| s.types(db))
                                             {
-                                                let mut diagnostic = builder.into_diagnostic(format_args!(
-                                                    "Invalid subscript assignment with key of type `{}` and value of \
-                                                     type `{assigned_d}` on object of type `{object_d}`",
-                                                    slice_ty.display(db),
-                                                ));
-
                                                 if !slice_ty.is_assignable_to(db, *expected_key_ty)
                                                 {
                                                     diagnostic.annotate(
@@ -3851,17 +3851,9 @@ impl<'db, 'ast> TypeInferenceBuilder<'db, 'ast> {
                                                             )),
                                                     );
                                                 }
-
-                                                attach_original_type_info(&mut diagnostic);
-                                            } else {
-                                                let mut 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 `{object_d}`",
-                                                    bindings.callable_type().display(db),
-                                                    slice_ty.display(db),
-                                                ));
-                                                attach_original_type_info(&mut diagnostic);
                                             }
+
+                                            attach_original_type_info(&mut diagnostic);
                                         }
                                     }
                                 }
@@ -3892,19 +3884,25 @@ impl<'db, 'ast> TypeInferenceBuilder<'db, 'ast> {
                                 ));
                                 attach_original_type_info(&mut diagnostic);
 
-                                // Use `KnownClass` as a crude proxy for checking if this is not a user-defined class. Otherwise,
-                                // we end up suggesting things like "Consider adding a `__setitem__` method to `None`".
+                                // If it's a user-defined class, suggest adding a `__setitem__` method.
                                 if object_ty
                                     .as_nominal_instance()
-                                    .is_some_and(|instance| instance.class(db).known(db).is_some())
+                                    .and_then(|instance| {
+                                        file_to_module(
+                                            db,
+                                            instance.class(db).class_literal(db).0.file(db),
+                                        )
+                                    })
+                                    .and_then(|module| module.search_path(db))
+                                    .is_some_and(crate::SearchPath::is_first_party)
                                 {
-                                    diagnostic.info(format_args!(
-                                        "`{}` does not have a `__setitem__` method.",
+                                    diagnostic.help(format_args!(
+                                        "Consider adding a `__setitem__` method to `{}`.",
                                         object_ty.display(db),
                                     ));
                                 } else {
-                                    diagnostic.help(format_args!(
-                                        "Consider adding a `__setitem__` method to `{}`.",
+                                    diagnostic.info(format_args!(
+                                        "`{}` does not have a `__setitem__` method.",
                                         object_ty.display(db),
                                     ));
                                 }