Allow users to provide custom diagnostic messages when unwrapping calls (#13597)

## Summary

You can now call `return_ty_result` to operate on a `Result` directly
thereby using your own diagnostics, as in:

```rust
return dunder_getitem_method
    .call(self.db, &[slice_ty])
    .return_ty_result(self.db, value.as_ref().into(), self)
    .unwrap_or_else(|err| {
        self.add_diagnostic(
            (&**value).into(),
            "call-non-callable",
            format_args!(
                "Method `__getitem__` is not callable on object of type '{}'.",
                value_ty.display(self.db),
            ),
        );
        err.return_ty()
    });
```

crates/red_knot_python_semantic/src/types.rs

@@ -221,9 +221,9 @@ fn declarations_ty<'db>(
         first
     };
     if conflicting.is_empty() {
-        DeclaredTypeResult::Ok(declared_ty)
+        Ok(declared_ty)
     } else {
-        DeclaredTypeResult::Err((
+        Err((
             declared_ty,
             [first].into_iter().chain(conflicting).collect(),
         ))
@@ -900,15 +900,73 @@ impl<'db> CallOutcome<'db> {
         }
     }
 
-    /// Get the return type of the call, emitting diagnostics if needed.
+    /// Get the return type of the call, emitting default diagnostics if needed.
     fn unwrap_with_diagnostic<'a>(
         &self,
         db: &'db dyn Db,
         node: ast::AnyNodeRef,
         builder: &'a mut TypeInferenceBuilder<'db>,
     ) -> Type<'db> {
+        match self.return_ty_result(db, node, builder) {
+            Ok(return_ty) => return_ty,
+            Err(NotCallableError::Type {
+                not_callable_ty,
+                return_ty,
+            }) => {
+                builder.add_diagnostic(
+                    node,
+                    "call-non-callable",
+                    format_args!(
+                        "Object of type '{}' is not callable.",
+                        not_callable_ty.display(db)
+                    ),
+                );
+                return_ty
+            }
+            Err(NotCallableError::UnionElement {
+                not_callable_ty,
+                called_ty,
+                return_ty,
+            }) => {
+                builder.add_diagnostic(
+                    node,
+                    "call-non-callable",
+                    format_args!(
+                        "Object of type '{}' is not callable (due to union element '{}').",
+                        called_ty.display(db),
+                        not_callable_ty.display(db),
+                    ),
+                );
+                return_ty
+            }
+            Err(NotCallableError::UnionElements {
+                not_callable_tys,
+                called_ty,
+                return_ty,
+            }) => {
+                builder.add_diagnostic(
+                    node,
+                    "call-non-callable",
+                    format_args!(
+                        "Object of type '{}' is not callable (due to union elements {}).",
+                        called_ty.display(db),
+                        not_callable_tys.display(db),
+                    ),
+                );
+                return_ty
+            }
+        }
+    }
+
+    /// Get the return type of the call as a result.
+    fn return_ty_result<'a>(
+        &self,
+        db: &'db dyn Db,
+        node: ast::AnyNodeRef,
+        builder: &'a mut TypeInferenceBuilder<'db>,
+    ) -> Result<Type<'db>, NotCallableError<'db>> {
         match self {
-            Self::Callable { return_ty } => *return_ty,
+            Self::Callable { return_ty } => Ok(*return_ty),
             Self::RevealType {
                 return_ty,
                 revealed_ty,
@@ -918,19 +976,12 @@ impl<'db> CallOutcome<'db> {
                     "revealed-type",
                     format_args!("Revealed type is '{}'.", revealed_ty.display(db)),
                 );
-                *return_ty
+                Ok(*return_ty)
-            }
-            Self::NotCallable { not_callable_ty } => {
-                builder.add_diagnostic(
-                    node,
-                    "call-non-callable",
-                    format_args!(
-                        "Object of type '{}' is not callable.",
-                        not_callable_ty.display(db)
-                    ),
-                );
-                Type::Unknown
             }
+            Self::NotCallable { not_callable_ty } => Err(NotCallableError::Type {
+                not_callable_ty: *not_callable_ty,
+                return_ty: Type::Unknown,
+            }),
             Self::Union {
                 outcomes,
                 called_ty,
@@ -959,41 +1010,75 @@ impl<'db> CallOutcome<'db> {
                     };
                     union_builder = union_builder.add(return_ty);
                 }
+                let return_ty = union_builder.build();
                 match not_callable[..] {
-                    [] => {}
+                    [] => Ok(return_ty),
-                    [elem] => builder.add_diagnostic(
+                    [elem] => Err(NotCallableError::UnionElement {
-                        node,
+                        not_callable_ty: elem,
-                        "call-non-callable",
+                        called_ty: *called_ty,
-                        format_args!(
+                        return_ty,
-                            "Object of type '{}' is not callable (due to union element '{}').",
+                    }),
-                            called_ty.display(db),
+                    _ if not_callable.len() == outcomes.len() => Err(NotCallableError::Type {
-                            elem.display(db),
+                        not_callable_ty: *called_ty,
-                        ),
+                        return_ty,
-                    ),
+                    }),
-                    _ if not_callable.len() == outcomes.len() => builder.add_diagnostic(
+                    _ => Err(NotCallableError::UnionElements {
-                        node,
+                        not_callable_tys: not_callable.into_boxed_slice(),
-                        "call-non-callable",
+                        called_ty: *called_ty,
-                        format_args!(
+                        return_ty,
-                            "Object of type '{}' is not callable.",
+                    }),
-                            called_ty.display(db)
-                        ),
-                    ),
-                    _ => builder.add_diagnostic(
-                        node,
-                        "call-non-callable",
-                        format_args!(
-                            "Object of type '{}' is not callable (due to union elements {}).",
-                            called_ty.display(db),
-                            not_callable.display(db),
-                        ),
-                    ),
                 }
-                union_builder.build()
             }
         }
     }
 }
 
+#[derive(Debug, Clone, PartialEq, Eq)]
+enum NotCallableError<'db> {
+    /// The type is not callable.
+    Type {
+        not_callable_ty: Type<'db>,
+        return_ty: Type<'db>,
+    },
+    /// A single union element is not callable.
+    UnionElement {
+        not_callable_ty: Type<'db>,
+        called_ty: Type<'db>,
+        return_ty: Type<'db>,
+    },
+    /// Multiple (but not all) union elements are not callable.
+    UnionElements {
+        not_callable_tys: Box<[Type<'db>]>,
+        called_ty: Type<'db>,
+        return_ty: Type<'db>,
+    },
+}
+
+impl<'db> NotCallableError<'db> {
+    /// The return type that should be used when a call is not callable.
+    fn return_ty(&self) -> Type<'db> {
+        match self {
+            Self::Type { return_ty, .. } => *return_ty,
+            Self::UnionElement { return_ty, .. } => *return_ty,
+            Self::UnionElements { return_ty, .. } => *return_ty,
+        }
+    }
+
+    /// The resolved type that was not callable.
+    ///
+    /// For unions, returns the union type itself, which may contain a mix of callable and
+    /// non-callable types.
+    fn called_ty(&self) -> Type<'db> {
+        match self {
+            Self::Type {
+                not_callable_ty, ..
+            } => *not_callable_ty,
+            Self::UnionElement { called_ty, .. } => *called_ty,
+            Self::UnionElements { called_ty, .. } => *called_ty,
+        }
+    }
+}
+
 #[derive(Debug, Clone, Copy, PartialEq, Eq)]
 enum IterationOutcome<'db> {
     Iterable { element_ty: Type<'db> },

crates/red_knot_python_semantic/src/types/infer.rs

@@ -2616,7 +2616,19 @@ impl<'db> TypeInferenceBuilder<'db> {
                 if !dunder_getitem_method.is_unbound() {
                     return dunder_getitem_method
                         .call(self.db, &[slice_ty])
-                        .unwrap_with_diagnostic(self.db, value.as_ref().into(), self);
+                        .return_ty_result(self.db, value.as_ref().into(), self)
+                        .unwrap_or_else(|err| {
+                            self.add_diagnostic(
+                                (&**value).into(),
+                                "call-non-callable",
+                                format_args!(
+                                    "Method `__getitem__` of type '{}' is not callable on object of type '{}'.",
+                                    err.called_ty().display(self.db),
+                                    value_ty.display(self.db),
+                                ),
+                            );
+                            err.return_ty()
+                        });
                 }
 
                 // Otherwise, if the value is itself a class and defines `__class_getitem__`,
@@ -2626,7 +2638,19 @@ impl<'db> TypeInferenceBuilder<'db> {
                     if !dunder_class_getitem_method.is_unbound() {
                         return dunder_class_getitem_method
                             .call(self.db, &[slice_ty])
-                            .unwrap_with_diagnostic(self.db, value.as_ref().into(), self);
+                            .return_ty_result(self.db, value.as_ref().into(), self)
+                            .unwrap_or_else(|err| {
+                                self.add_diagnostic(
+                                    (&**value).into(),
+                                    "call-non-callable",
+                                    format_args!(
+                                        "Method `__class_getitem__` of type '{}' is not callable on object of type '{}'.",
+                                        err.called_ty().display(self.db),
+                                        value_ty.display(self.db),
+                                    ),
+                                );
+                                err.return_ty()
+                            });
                     }
 
                     self.non_subscriptable_diagnostic(
@@ -6840,7 +6864,7 @@ mod tests {
         assert_file_diagnostics(
             &db,
             "/src/a.py",
-            &["Object of type 'None' is not callable."],
+            &["Method `__getitem__` of type 'None' is not callable on object of type 'NotSubscriptable'."],
         );
 
         Ok(())
@@ -7015,7 +7039,7 @@ mod tests {
         assert_file_diagnostics(
             &db,
             "/src/a.py",
-            &["Object of type 'Literal[__class_getitem__] | Unbound' is not callable (due to union element 'Unbound')."],
+            &["Method `__class_getitem__` of type 'Literal[__class_getitem__] | Unbound' is not callable on object of type 'Literal[Identity, Identity]'."],
         );
 
         Ok(())