[ty] Support __setitem__ and improve __getitem__ related diagnostics (#19578)

## Summary

Adds validation to subscript assignment expressions.

```py
class Foo: ...

class Bar:
    __setattr__ = None

class Baz:
    def __setitem__(self, index: str, value: int) -> None:
        pass

# We now emit a diagnostic on these statements
Foo()[1] = 2
Bar()[1] = 2
Baz()[1] = 2

```

Also improves error messages on invalid `__getitem__` expressions

## Test Plan

Update mdtests and add more to `subscript/instance.md`

---------

Co-authored-by: David Peter <sharkdp@users.noreply.github.com>
Co-authored-by: David Peter <mail@david-peter.de>

crates/ty_python_semantic/resources/mdtest/del.md

@@ -181,7 +181,7 @@ def f(l: list[int]):
     # but if it was greater than that, it will not be an error.
     reveal_type(l[0])  # revealed: int
 
-    # error: [call-non-callable]
+    # error: [invalid-argument-type]
     del l["string"]
 
     l[0] = 1

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

@@ -253,6 +253,7 @@ does["not"]["exist"] = 0
 reveal_type(does["not"]["exist"])  # revealed: Unknown
 
 non_subscriptable = 1
+# error: [invalid-assignment]
 non_subscriptable[0] = 0
 # error: [non-subscriptable]
 reveal_type(non_subscriptable[0])  # revealed: Unknown
@@ -317,7 +318,7 @@ def f(c: C, s: str):
     reveal_type(c.x)  # revealed: int | None
     s = c.x  # error: [invalid-assignment]
 
-    # TODO: This assignment is invalid and should result in an error.
+    # 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]`"
     c.l[0] = s
     reveal_type(c.l[0])  # revealed: int
 ```

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

@@ -14,7 +14,7 @@ a = NotSubscriptable()[0]  # error: "Cannot subscript object of type `NotSubscri
 class NotSubscriptable:
     __getitem__ = None
 
-# error: "Method `__getitem__` of type `Unknown | None` is not callable on object of type `NotSubscriptable`"
+# error: "Method `__getitem__` of type `Unknown | None` is possibly not callable on object of type `NotSubscriptable`"
 a = NotSubscriptable()[0]
 ```
 
@@ -43,6 +43,18 @@ def _(flag: bool):
     reveal_type(Identity()[0])  # revealed: int | str
 ```
 
+## `__getitem__` with invalid index argument
+
+```py
+class Identity:
+    def __getitem__(self, index: int) -> int:
+        return index
+
+a = Identity()
+# error: [invalid-argument-type] "Method `__getitem__` of type `bound method Identity.__getitem__(index: int) -> int` cannot be called with key of type `Literal["a"]` on object of type `Identity`"
+a["a"]
+```
+
 ## `__setitem__` with no `__getitem__`
 
 ```py
@@ -53,3 +65,45 @@ class NoGetitem:
 a = NoGetitem()
 a[0] = 0
 ```
+
+## Subscript store with no `__setitem__`
+
+```py
+class NoSetitem: ...
+
+a = NoSetitem()
+a[0] = 0  # error: "Cannot assign to object of type `NoSetitem` with no `__setitem__` method"
+```
+
+## `__setitem__` not callable
+
+```py
+class NoSetitem:
+    __setitem__ = None
+
+a = NoSetitem()
+a[0] = 0  # error: "Method `__setitem__` of type `Unknown | None` is possibly not callable on object of type `NoSetitem`"
+```
+
+## Valid `__setitem__` method
+
+```py
+class Identity:
+    def __setitem__(self, index: int, value: int) -> None:
+        pass
+
+a = Identity()
+a[0] = 0
+```
+
+## `__setitem__` with invalid index argument
+
+```py
+class Identity:
+    def __setitem__(self, index: int, value: int) -> None:
+        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`"
+a["a"] = 0
+```

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

@@ -17,7 +17,7 @@ reveal_type(x[0])  # revealed: Unknown
 # TODO reveal list[int]
 reveal_type(x[0:1])  # revealed: list[Unknown]
 
-# error: [call-non-callable]
+# error: [invalid-argument-type]
 reveal_type(x["a"])  # revealed: Unknown
 ```
 
@@ -29,11 +29,9 @@ In assignment, we might also have a named assignment. This should also get type
 x = [1, 2, 3]
 x[0 if (y := 2) else 1] = 5
 
-# TODO: better error than "method `__getitem__` not callable on type `list`"
+# error: [invalid-assignment]
-# error: [call-non-callable]
 x["a" if (y := 2) else 1] = 6
 
-# TODO: better error than "method `__getitem__` not callable on type `list`"
+# error: [invalid-assignment]
-# error: [call-non-callable]
 x["a" if (y := 2) else "b"] = 6
 ```

crates/ty_python_semantic/src/types/infer.rs

@@ -89,7 +89,7 @@ use crate::semantic_index::symbol::ScopedSymbolId;
 use crate::semantic_index::{
     ApplicableConstraints, EnclosingSnapshotResult, SemanticIndex, place_table, semantic_index,
 };
-use crate::types::call::{Binding, Bindings, CallArguments, CallError};
+use crate::types::call::{Binding, Bindings, CallArguments, CallError, CallErrorKind};
 use crate::types::class::{CodeGeneratorKind, DataclassField, MetaclassErrorKind, SliceLiteral};
 use crate::types::diagnostic::{
     self, CALL_NON_CALLABLE, CONFLICTING_DECLARATIONS, CONFLICTING_METACLASS,
@@ -3619,6 +3619,99 @@ impl<'db, 'ast> TypeInferenceBuilder<'db, 'ast> {
         self.infer_target_impl(target, assigned_ty);
     }
 
+    /// Make sure that the subscript assignment `obj[slice] = value` is valid.
+    fn validate_subscript_assignment(
+        &mut self,
+        target: &ast::ExprSubscript,
+        assigned_ty: Type<'db>,
+    ) -> bool {
+        let ast::ExprSubscript {
+            range: _,
+            node_index: _,
+            value,
+            slice,
+            ctx: _,
+        } = target;
+
+        let value_ty = self.infer_expression(value);
+        let slice_ty = self.infer_expression(slice);
+
+        let db = self.db();
+        let context = &self.context;
+
+        match value_ty.try_call_dunder(
+            db,
+            "__setitem__",
+            CallArguments::positional([slice_ty, assigned_ty]),
+        ) {
+            Ok(_) => true,
+            Err(err) => match err {
+                CallDunderError::PossiblyUnbound { .. } => {
+                    if let Some(builder) =
+                        context.report_lint(&POSSIBLY_UNBOUND_IMPLICIT_CALL, &**value)
+                    {
+                        builder.into_diagnostic(format_args!(
+                            "Method `__setitem__` of type `{}` is possibly unbound",
+                            value_ty.display(db),
+                        ));
+                    }
+                    false
+                }
+                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 => {
+                            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 `{}` on object of type `{}`",
+                                    bindings.callable_type().display(db),
+                                    slice_ty.display(db),
+                                    assigned_ty.display(db),
+                                    value_ty.display(db),
+                                ));
+                            }
+                        }
+                        CallErrorKind::PossiblyNotCallable => {
+                            if let Some(builder) = context.report_lint(&CALL_NON_CALLABLE, &**value)
+                            {
+                                builder.into_diagnostic(format_args!(
+                                    "Method `__setitem__` of type `{}` is possibly not \
+                                    callable on object of type `{}`",
+                                    bindings.callable_type().display(db),
+                                    value_ty.display(db),
+                                ));
+                            }
+                        }
+                    }
+                    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
+                }
+            },
+        }
+    }
+
     /// Make sure that the attribute assignment `obj.attribute = value` is valid.
     ///
     /// `target` is the node for the left-hand side, `object_ty` is the type of `obj`, `attribute` is
@@ -4191,6 +4284,13 @@ impl<'db, 'ast> TypeInferenceBuilder<'db, 'ast> {
                     );
                 }
             }
+            ast::Expr::Subscript(subscript_expr) => {
+                self.store_expression_type(target, assigned_ty.unwrap_or(Type::unknown()));
+
+                if let Some(assigned_ty) = assigned_ty {
+                    self.validate_subscript_assignment(subscript_expr, assigned_ty);
+                }
+            }
             _ => {
                 // TODO: Remove this once we handle all possible assignment targets.
                 self.infer_expression(target);
@@ -8616,14 +8716,40 @@ impl<'db, 'ast> TypeInferenceBuilder<'db, 'ast> {
 
                 return err.fallback_return_type(db);
             }
-            Err(CallDunderError::CallError(_, bindings)) => {
+            Err(CallDunderError::CallError(call_error_kind, bindings)) => {
-                if let Some(builder) = context.report_lint(&CALL_NON_CALLABLE, value_node) {
+                match call_error_kind {
-                    builder.into_diagnostic(format_args!(
+                    CallErrorKind::NotCallable => {
-                        "Method `__getitem__` of type `{}` \
+                        if let Some(builder) = context.report_lint(&CALL_NON_CALLABLE, value_node) {
-                            is not callable on object of type `{}`",
+                            builder.into_diagnostic(format_args!(
-                        bindings.callable_type().display(db),
+                                "Method `__getitem__` of type `{}` \
-                        value_ty.display(db),
+                                is not callable on object of type `{}`",
-                    ));
+                                bindings.callable_type().display(db),
+                                value_ty.display(db),
+                            ));
+                        }
+                    }
+                    CallErrorKind::BindingError => {
+                        if let Some(builder) =
+                            context.report_lint(&INVALID_ARGUMENT_TYPE, value_node)
+                        {
+                            builder.into_diagnostic(format_args!(
+                                "Method `__getitem__` of type `{}` cannot be called with key of \
+                                type `{}` on object of type `{}`",
+                                bindings.callable_type().display(db),
+                                slice_ty.display(db),
+                                value_ty.display(db),
+                            ));
+                        }
+                    }
+                    CallErrorKind::PossiblyNotCallable => {
+                        if let Some(builder) = context.report_lint(&CALL_NON_CALLABLE, value_node) {
+                            builder.into_diagnostic(format_args!(
+                                "Method `__getitem__` of type `{}` is possibly not callable on object of type `{}`",
+                                bindings.callable_type().display(db),
+                                value_ty.display(db),
+                            ));
+                        }
+                    }
                 }
 
                 return bindings.return_type(db);