[ty] Implement go-to for binary and unary operators (#21001)

Co-authored-by: Alex Waygood <Alex.Waygood@Gmail.com>

crates/ty_python_semantic/src/lib.rs

@@ -27,8 +27,9 @@ pub use semantic_model::{
 pub use site_packages::{PythonEnvironment, SitePackagesPaths, SysPrefixPathOrigin};
 pub use types::DisplaySettings;
 pub use types::ide_support::{
-    ImportAliasResolution, ResolvedDefinition, definitions_for_attribute,
-    definitions_for_imported_symbol, definitions_for_name, map_stub_definition,
+    ImportAliasResolution, ResolvedDefinition, definitions_for_attribute, definitions_for_bin_op,
+    definitions_for_imported_symbol, definitions_for_name, definitions_for_unary_op,
+    map_stub_definition,
 };
 
 pub mod ast_node_ref;

crates/ty_python_semantic/src/semantic_model.rs

@@ -27,7 +27,7 @@ impl<'db> SemanticModel<'db> {
 
     // TODO we don't actually want to expose the Db directly to lint rules, but we need to find a
     // solution for exposing information from types
-    pub fn db(&self) -> &dyn Db {
+    pub fn db(&self) -> &'db dyn Db {
         self.db
     }
 

crates/ty_python_semantic/src/types.rs

@@ -9987,6 +9987,14 @@ impl<'db> BoundMethodType<'db> {
         self_instance
     }
 
+    pub(crate) fn map_self_type(
+        self,
+        db: &'db dyn Db,
+        f: impl FnOnce(Type<'db>) -> Type<'db>,
+    ) -> Self {
+        Self::new(db, self.function(db), f(self.self_instance(db)))
+    }
+
     #[salsa::tracked(cycle_fn=into_callable_type_cycle_recover, cycle_initial=into_callable_type_cycle_initial, heap_size=ruff_memory_usage::heap_size)]
     pub(crate) fn into_callable_type(self, db: &'db dyn Db) -> CallableType<'db> {
         let function = self.function(db);

crates/ty_python_semantic/src/types/call.rs

@@ -1,14 +1,87 @@
 use super::context::InferContext;
-use super::{Signature, Type};
+use super::{Signature, Type, TypeContext};
 use crate::Db;
 use crate::types::PropertyInstanceType;
 use crate::types::call::bind::BindingError;
+use ruff_python_ast as ast;
 
 mod arguments;
 pub(crate) mod bind;
 pub(super) use arguments::{Argument, CallArguments};
 pub(super) use bind::{Binding, Bindings, CallableBinding, MatchedArgument};
 
+impl<'db> Type<'db> {
+    pub(crate) fn try_call_bin_op(
+        db: &'db dyn Db,
+        left_ty: Type<'db>,
+        op: ast::Operator,
+        right_ty: Type<'db>,
+    ) -> Result<Bindings<'db>, CallBinOpError> {
+        // We either want to call lhs.__op__ or rhs.__rop__. The full decision tree from
+        // the Python spec [1] is:
+        //
+        //   - If rhs is a (proper) subclass of lhs, and it provides a different
+        //     implementation of __rop__, use that.
+        //   - Otherwise, if lhs implements __op__, use that.
+        //   - Otherwise, if lhs and rhs are different types, and rhs implements __rop__,
+        //     use that.
+        //
+        // [1] https://docs.python.org/3/reference/datamodel.html#object.__radd__
+
+        // Technically we don't have to check left_ty != right_ty here, since if the types
+        // are the same, they will trivially have the same implementation of the reflected
+        // dunder, and so we'll fail the inner check. But the type equality check will be
+        // faster for the common case, and allow us to skip the (two) class member lookups.
+        let left_class = left_ty.to_meta_type(db);
+        let right_class = right_ty.to_meta_type(db);
+        if left_ty != right_ty && right_ty.is_subtype_of(db, left_ty) {
+            let reflected_dunder = op.reflected_dunder();
+            let rhs_reflected = right_class.member(db, reflected_dunder).place;
+            // TODO: if `rhs_reflected` is possibly unbound, we should union the two possible
+            // Bindings together
+            if !rhs_reflected.is_undefined()
+                && rhs_reflected != left_class.member(db, reflected_dunder).place
+            {
+                return Ok(right_ty
+                    .try_call_dunder(
+                        db,
+                        reflected_dunder,
+                        CallArguments::positional([left_ty]),
+                        TypeContext::default(),
+                    )
+                    .or_else(|_| {
+                        left_ty.try_call_dunder(
+                            db,
+                            op.dunder(),
+                            CallArguments::positional([right_ty]),
+                            TypeContext::default(),
+                        )
+                    })?);
+            }
+        }
+
+        let call_on_left_instance = left_ty.try_call_dunder(
+            db,
+            op.dunder(),
+            CallArguments::positional([right_ty]),
+            TypeContext::default(),
+        );
+
+        call_on_left_instance.or_else(|_| {
+            if left_ty == right_ty {
+                Err(CallBinOpError::NotSupported)
+            } else {
+                Ok(right_ty.try_call_dunder(
+                    db,
+                    op.reflected_dunder(),
+                    CallArguments::positional([left_ty]),
+                    TypeContext::default(),
+                )?)
+            }
+        })
+    }
+}
+
 /// Wraps a [`Bindings`] for an unsuccessful call with information about why the call was
 /// unsuccessful.
 ///
@@ -26,7 +99,7 @@ impl<'db> CallError<'db> {
             return None;
         }
         self.1
-            .into_iter()
+            .iter()
             .flatten()
             .flat_map(bind::Binding::errors)
             .find_map(|error| match error {
@@ -89,3 +162,24 @@ impl<'db> From<CallError<'db>> for CallDunderError<'db> {
         Self::CallError(kind, bindings)
     }
 }
+
+#[derive(Debug)]
+pub(crate) enum CallBinOpError {
+    /// The dunder attribute exists but it can't be called with the given arguments.
+    ///
+    /// This includes non-callable dunder attributes that are possibly unbound.
+    CallError,
+
+    NotSupported,
+}
+
+impl From<CallDunderError<'_>> for CallBinOpError {
+    fn from(value: CallDunderError<'_>) -> Self {
+        match value {
+            CallDunderError::CallError(_, _) => Self::CallError,
+            CallDunderError::MethodNotAvailable | CallDunderError::PossiblyUnbound(_) => {
+                CallBinOpError::NotSupported
+            }
+        }
+    }
+}

crates/ty_python_semantic/src/types/call/bind.rs

@@ -96,6 +96,10 @@ impl<'db> Bindings<'db> {
         &self.argument_forms.values
     }
 
+    pub(crate) fn iter(&self) -> std::slice::Iter<'_, CallableBinding<'db>> {
+        self.elements.iter()
+    }
+
     /// Match the arguments of a call site against the parameters of a collection of possibly
     /// unioned, possibly overloaded signatures.
     ///
@@ -1178,7 +1182,16 @@ impl<'a, 'db> IntoIterator for &'a Bindings<'db> {
     type IntoIter = std::slice::Iter<'a, CallableBinding<'db>>;
 
     fn into_iter(self) -> Self::IntoIter {
-        self.elements.iter()
+        self.iter()
+    }
+}
+
+impl<'db> IntoIterator for Bindings<'db> {
+    type Item = CallableBinding<'db>;
+    type IntoIter = smallvec::IntoIter<[CallableBinding<'db>; 1]>;
+
+    fn into_iter(self) -> Self::IntoIter {
+        self.elements.into_iter()
     }
 }
 
@@ -2106,6 +2119,15 @@ impl<'a, 'db> IntoIterator for &'a CallableBinding<'db> {
     }
 }
 
+impl<'db> IntoIterator for CallableBinding<'db> {
+    type Item = Binding<'db>;
+    type IntoIter = smallvec::IntoIter<[Binding<'db>; 1]>;
+
+    fn into_iter(self) -> Self::IntoIter {
+        self.overloads.into_iter()
+    }
+}
+
 #[derive(Debug, Copy, Clone)]
 enum OverloadCallReturnType<'db> {
     ArgumentTypeExpansion(Type<'db>),

crates/ty_python_semantic/src/types/ide_support.rs

@@ -13,7 +13,7 @@ use crate::semantic_index::{
 use crate::types::call::{CallArguments, MatchedArgument};
 use crate::types::signatures::Signature;
 use crate::types::{
-    ClassBase, ClassLiteral, DynamicType, KnownClass, KnownInstanceType, Type,
+    ClassBase, ClassLiteral, DynamicType, KnownClass, KnownInstanceType, Type, TypeContext,
     TypeVarBoundOrConstraints, class::CodeGeneratorKind,
 };
 use crate::{Db, HasType, NameKind, SemanticModel};
@@ -908,18 +908,19 @@ pub fn call_signature_details<'db>(
             .into_iter()
             .flat_map(std::iter::IntoIterator::into_iter)
             .map(|binding| {
-                let signature = &binding.signature;
+                let argument_to_parameter_mapping = binding.argument_matches().to_vec();
+                let signature = binding.signature;
                 let display_details = signature.display(db).to_string_parts();
-                let parameter_label_offsets = display_details.parameter_ranges.clone();
-                let parameter_names = display_details.parameter_names.clone();
+                let parameter_label_offsets = display_details.parameter_ranges;
+                let parameter_names = display_details.parameter_names;
 
                 CallSignatureDetails {
-                    signature: signature.clone(),
+                    definition: signature.definition(),
+                    signature,
                     label: display_details.label,
                     parameter_label_offsets,
                     parameter_names,
-                    definition: signature.definition(),
-                    argument_to_parameter_mapping: binding.argument_matches().to_vec(),
+                    argument_to_parameter_mapping,
                 }
             })
             .collect()
@@ -929,6 +930,91 @@ pub fn call_signature_details<'db>(
     }
 }
 
+/// Returns the definitions of the binary operation along with its callable type.
+pub fn definitions_for_bin_op<'db>(
+    db: &'db dyn Db,
+    model: &SemanticModel<'db>,
+    binary_op: &ast::ExprBinOp,
+) -> Option<(Vec<ResolvedDefinition<'db>>, Type<'db>)> {
+    let left_ty = binary_op.left.inferred_type(model);
+    let right_ty = binary_op.right.inferred_type(model);
+
+    let Ok(bindings) = Type::try_call_bin_op(db, left_ty, binary_op.op, right_ty) else {
+        return None;
+    };
+
+    let callable_type = promote_literals_for_self(db, bindings.callable_type());
+
+    let definitions: Vec<_> = bindings
+        .into_iter()
+        .flat_map(std::iter::IntoIterator::into_iter)
+        .filter_map(|binding| {
+            Some(ResolvedDefinition::Definition(
+                binding.signature.definition?,
+            ))
+        })
+        .collect();
+
+    Some((definitions, callable_type))
+}
+
+/// Returns the definitions for an unary operator along with their callable types.
+pub fn definitions_for_unary_op<'db>(
+    db: &'db dyn Db,
+    model: &SemanticModel<'db>,
+    unary_op: &ast::ExprUnaryOp,
+) -> Option<(Vec<ResolvedDefinition<'db>>, Type<'db>)> {
+    let operand_ty = unary_op.operand.inferred_type(model);
+
+    let unary_dunder_method = match unary_op.op {
+        ast::UnaryOp::Invert => "__invert__",
+        ast::UnaryOp::UAdd => "__pos__",
+        ast::UnaryOp::USub => "__neg__",
+        ast::UnaryOp::Not => "__bool__",
+    };
+
+    let Ok(bindings) = operand_ty.try_call_dunder(
+        db,
+        unary_dunder_method,
+        CallArguments::none(),
+        TypeContext::default(),
+    ) else {
+        return None;
+    };
+
+    let callable_type = promote_literals_for_self(db, bindings.callable_type());
+
+    let definitions = bindings
+        .into_iter()
+        .flat_map(std::iter::IntoIterator::into_iter)
+        .filter_map(|binding| {
+            Some(ResolvedDefinition::Definition(
+                binding.signature.definition?,
+            ))
+        })
+        .collect();
+
+    Some((definitions, callable_type))
+}
+
+/// Promotes literal types in `self` positions to their fallback instance types.
+///
+/// This is so that we show e.g. `int.__add__` instead of `Literal[4].__add__`.
+fn promote_literals_for_self<'db>(db: &'db dyn Db, ty: Type<'db>) -> Type<'db> {
+    match ty {
+        Type::BoundMethod(method) => Type::BoundMethod(method.map_self_type(db, |self_ty| {
+            self_ty.literal_fallback_instance(db).unwrap_or(self_ty)
+        })),
+        Type::Union(elements) => elements.map(db, |ty| match ty {
+            Type::BoundMethod(method) => Type::BoundMethod(method.map_self_type(db, |self_ty| {
+                self_ty.literal_fallback_instance(db).unwrap_or(self_ty)
+            })),
+            _ => *ty,
+        }),
+        ty => ty,
+    }
+}
+
 /// Find the active signature index from `CallSignatureDetails`.
 /// The active signature is the first signature where all arguments present in the call
 /// have valid mappings to parameters (i.e., none of the mappings are None).

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

@@ -8216,80 +8216,9 @@ impl<'db, 'ast> TypeInferenceBuilder<'db, 'ast> {
                 | Type::TypeIs(_)
                 | Type::TypedDict(_),
                 op,
-            ) => {
-                // We either want to call lhs.__op__ or rhs.__rop__. The full decision tree from
-                // the Python spec [1] is:
-                //
-                //   - If rhs is a (proper) subclass of lhs, and it provides a different
-                //     implementation of __rop__, use that.
-                //   - Otherwise, if lhs implements __op__, use that.
-                //   - Otherwise, if lhs and rhs are different types, and rhs implements __rop__,
-                //     use that.
-                //
-                // [1] https://docs.python.org/3/reference/datamodel.html#object.__radd__
-
-                // Technically we don't have to check left_ty != right_ty here, since if the types
-                // are the same, they will trivially have the same implementation of the reflected
-                // dunder, and so we'll fail the inner check. But the type equality check will be
-                // faster for the common case, and allow us to skip the (two) class member lookups.
-                let left_class = left_ty.to_meta_type(self.db());
-                let right_class = right_ty.to_meta_type(self.db());
-                if left_ty != right_ty && right_ty.is_subtype_of(self.db(), left_ty) {
-                    let reflected_dunder = op.reflected_dunder();
-                    let rhs_reflected = right_class.member(self.db(), reflected_dunder).place;
-                    // TODO: if `rhs_reflected` is possibly unbound, we should union the two possible
-                    // Bindings together
-                    if !rhs_reflected.is_undefined()
-                        && rhs_reflected != left_class.member(self.db(), reflected_dunder).place
-                    {
-                        return right_ty
-                            .try_call_dunder(
-                                self.db(),
-                                reflected_dunder,
-                                CallArguments::positional([left_ty]),
-                                TypeContext::default(),
-                            )
-                            .map(|outcome| outcome.return_type(self.db()))
-                            .or_else(|_| {
-                                left_ty
-                                    .try_call_dunder(
-                                        self.db(),
-                                        op.dunder(),
-                                        CallArguments::positional([right_ty]),
-                                        TypeContext::default(),
-                                    )
-                                    .map(|outcome| outcome.return_type(self.db()))
-                            })
-                            .ok();
-                    }
-                }
-
-                let call_on_left_instance = left_ty
-                    .try_call_dunder(
-                        self.db(),
-                        op.dunder(),
-                        CallArguments::positional([right_ty]),
-                        TypeContext::default(),
-                    )
-                    .map(|outcome| outcome.return_type(self.db()))
-                    .ok();
-
-                call_on_left_instance.or_else(|| {
-                    if left_ty == right_ty {
-                        None
-                    } else {
-                        right_ty
-                            .try_call_dunder(
-                                self.db(),
-                                op.reflected_dunder(),
-                                CallArguments::positional([left_ty]),
-                                TypeContext::default(),
-                            )
-                            .map(|outcome| outcome.return_type(self.db()))
-                            .ok()
-                    }
-                })
-            }
+            ) => Type::try_call_bin_op(self.db(), left_ty, op, right_ty)
+                .map(|outcome| outcome.return_type(self.db()))
+                .ok(),
         }
     }