diff --git a/crates/ruff_python_parser/src/token.rs b/crates/ruff_python_parser/src/token.rs index 1d9461a722..18b7648c4c 100644 --- a/crates/ruff_python_parser/src/token.rs +++ b/crates/ruff_python_parser/src/token.rs @@ -486,7 +486,7 @@ impl TokenKind { /// /// [`as_unary_operator`]: TokenKind::as_unary_operator #[inline] - pub(crate) const fn as_unary_arithmetic_operator(self) -> Option { + pub const fn as_unary_arithmetic_operator(self) -> Option { Some(match self { TokenKind::Plus => UnaryOp::UAdd, TokenKind::Minus => UnaryOp::USub, @@ -501,7 +501,7 @@ impl TokenKind { /// /// [`as_unary_arithmetic_operator`]: TokenKind::as_unary_arithmetic_operator #[inline] - pub(crate) const fn as_unary_operator(self) -> Option { + pub const fn as_unary_operator(self) -> Option { Some(match self { TokenKind::Plus => UnaryOp::UAdd, TokenKind::Minus => UnaryOp::USub, @@ -514,7 +514,7 @@ impl TokenKind { /// Returns the [`BoolOp`] that corresponds to this token kind, if it is a boolean operator, /// otherwise return [None]. #[inline] - pub(crate) const fn as_bool_operator(self) -> Option { + pub const fn as_bool_operator(self) -> Option { Some(match self { TokenKind::And => BoolOp::And, TokenKind::Or => BoolOp::Or, @@ -528,7 +528,7 @@ impl TokenKind { /// Use [`as_augmented_assign_operator`] to match against an augmented assignment token. /// /// [`as_augmented_assign_operator`]: TokenKind::as_augmented_assign_operator - pub(crate) const fn as_binary_operator(self) -> Option { + pub const fn as_binary_operator(self) -> Option { Some(match self { TokenKind::Plus => Operator::Add, TokenKind::Minus => Operator::Sub, @@ -550,7 +550,7 @@ impl TokenKind { /// Returns the [`Operator`] that corresponds to this token kind, if it is /// an augmented assignment operator, or [`None`] otherwise. #[inline] - pub(crate) const fn as_augmented_assign_operator(self) -> Option { + pub const fn as_augmented_assign_operator(self) -> Option { Some(match self { TokenKind::PlusEqual => Operator::Add, TokenKind::MinusEqual => Operator::Sub, diff --git a/crates/ty_ide/src/goto.rs b/crates/ty_ide/src/goto.rs index bda25506f2..d7a7091f94 100644 --- a/crates/ty_ide/src/goto.rs +++ b/crates/ty_ide/src/goto.rs @@ -8,19 +8,18 @@ use std::borrow::Cow; use crate::find_node::covering_node; use crate::stub_mapping::StubMapper; use ruff_db::parsed::ParsedModuleRef; -use ruff_python_ast::ExprCall; use ruff_python_ast::{self as ast, AnyNodeRef}; -use ruff_python_parser::TokenKind; +use ruff_python_parser::{TokenKind, Tokens}; use ruff_text_size::{Ranged, TextRange, TextSize}; -use ty_python_semantic::HasDefinition; -use ty_python_semantic::ImportAliasResolution; + use ty_python_semantic::ResolvedDefinition; use ty_python_semantic::types::Type; use ty_python_semantic::types::ide_support::{ call_signature_details, definitions_for_keyword_argument, }; use ty_python_semantic::{ - HasType, SemanticModel, definitions_for_imported_symbol, definitions_for_name, + HasDefinition, HasType, ImportAliasResolution, SemanticModel, definitions_for_imported_symbol, + definitions_for_name, }; #[derive(Clone, Debug)] @@ -30,6 +29,28 @@ pub(crate) enum GotoTarget<'a> { ClassDef(&'a ast::StmtClassDef), Parameter(&'a ast::Parameter), + /// Go to on the operator of a binary operation. + /// + /// ```py + /// a + b + /// ^ + /// ``` + BinOp { + expression: &'a ast::ExprBinOp, + operator_range: TextRange, + }, + + /// Go to where the operator of a unary operation is defined. + /// + /// ```py + /// -a + /// ^ + /// ``` + UnaryOp { + expression: &'a ast::ExprUnaryOp, + operator_range: TextRange, + }, + /// Multi-part module names /// Handles both `import foo.bar` and `from foo.bar import baz` cases /// ```py @@ -166,7 +187,7 @@ pub(crate) enum GotoTarget<'a> { /// The callable that can actually be selected by a cursor callable: ast::ExprRef<'a>, /// The call of the callable - call: &'a ExprCall, + call: &'a ast::ExprCall, }, } @@ -295,6 +316,16 @@ impl GotoTarget<'_> { | GotoTarget::TypeParamTypeVarTupleName(_) | GotoTarget::NonLocal { .. } | GotoTarget::Globals { .. } => return None, + GotoTarget::BinOp { expression, .. } => { + let (_, ty) = + ty_python_semantic::definitions_for_bin_op(model.db(), model, expression)?; + ty + } + GotoTarget::UnaryOp { expression, .. } => { + let (_, ty) = + ty_python_semantic::definitions_for_unary_op(model.db(), model, expression)?; + ty + } }; Some(ty) @@ -451,6 +482,23 @@ impl GotoTarget<'_> { } } + GotoTarget::BinOp { expression, .. } => { + let model = SemanticModel::new(db, file); + + let (definitions, _) = + ty_python_semantic::definitions_for_bin_op(db, &model, expression)?; + + Some(DefinitionsOrTargets::Definitions(definitions)) + } + + GotoTarget::UnaryOp { expression, .. } => { + let model = SemanticModel::new(db, file); + let (definitions, _) = + ty_python_semantic::definitions_for_unary_op(db, &model, expression)?; + + Some(DefinitionsOrTargets::Definitions(definitions)) + } + _ => None, } } @@ -524,6 +572,7 @@ impl GotoTarget<'_> { } GotoTarget::NonLocal { identifier, .. } => Some(Cow::Borrowed(identifier.as_str())), GotoTarget::Globals { identifier, .. } => Some(Cow::Borrowed(identifier.as_str())), + GotoTarget::BinOp { .. } | GotoTarget::UnaryOp { .. } => None, } } @@ -531,6 +580,7 @@ impl GotoTarget<'_> { pub(crate) fn from_covering_node<'a>( covering_node: &crate::find_node::CoveringNode<'a>, offset: TextSize, + tokens: &Tokens, ) -> Option> { tracing::trace!("Covering node is of kind {:?}", covering_node.node().kind()); @@ -690,6 +740,44 @@ impl GotoTarget<'_> { } }, + AnyNodeRef::ExprBinOp(binary) => { + if offset >= binary.left.end() && offset < binary.right.start() { + let between_operands = + tokens.in_range(TextRange::new(binary.left.end(), binary.right.start())); + if let Some(operator_token) = between_operands + .iter() + .find(|token| token.kind().as_binary_operator().is_some()) + && operator_token.range().contains_inclusive(offset) + { + return Some(GotoTarget::BinOp { + expression: binary, + operator_range: operator_token.range(), + }); + } + } + + Some(GotoTarget::Expression(binary.into())) + } + + AnyNodeRef::ExprUnaryOp(unary) => { + if offset >= unary.start() && offset < unary.operand.start() { + let before_operand = + tokens.in_range(TextRange::new(unary.start(), unary.operand.start())); + + if let Some(operator_token) = before_operand + .iter() + .find(|token| token.kind().as_unary_operator().is_some()) + && operator_token.range().contains_inclusive(offset) + { + return Some(GotoTarget::UnaryOp { + expression: unary, + operator_range: operator_token.range(), + }); + } + } + Some(GotoTarget::Expression(unary.into())) + } + node => { // Check if this is seemingly a callable being invoked (the `x` in `x(...)`) let parent = covering_node.parent(); @@ -737,6 +825,8 @@ impl Ranged for GotoTarget<'_> { GotoTarget::TypeParamTypeVarTupleName(tuple) => tuple.name.range, GotoTarget::NonLocal { identifier, .. } => identifier.range, GotoTarget::Globals { identifier, .. } => identifier.range, + GotoTarget::BinOp { operator_range, .. } + | GotoTarget::UnaryOp { operator_range, .. } => *operator_range, } } } @@ -794,7 +884,7 @@ fn definitions_for_expression<'db>( fn definitions_for_callable<'db>( db: &'db dyn crate::Db, file: ruff_db::files::File, - call: &ExprCall, + call: &ast::ExprCall, ) -> Vec> { let model = SemanticModel::new(db, file); // Attempt to refine to a specific call @@ -835,14 +925,24 @@ pub(crate) fn find_goto_target( | TokenKind::Complex | TokenKind::Float | TokenKind::Int => 1, + + TokenKind::Comment => -1, + + // if we have a+b`, prefer the `+` token (by respecting the token ordering) + // This matches VS Code's behavior where it sends the start of the clicked token as offset. + kind if kind.as_binary_operator().is_some() || kind.as_unary_operator().is_some() => 1, _ => 0, })?; + if token.kind().is_comment() { + return None; + } + let covering_node = covering_node(parsed.syntax().into(), token.range()) .find_first(|node| node.is_identifier() || node.is_expression()) .ok()?; - GotoTarget::from_covering_node(&covering_node, offset) + GotoTarget::from_covering_node(&covering_node, offset, parsed.tokens()) } /// Helper function to resolve a module name and create a navigation target. diff --git a/crates/ty_ide/src/goto_definition.rs b/crates/ty_ide/src/goto_definition.rs index fb165f61a0..6cc6d0c23d 100644 --- a/crates/ty_ide/src/goto_definition.rs +++ b/crates/ty_ide/src/goto_definition.rs @@ -798,26 +798,6 @@ my_func(my_other_func(ab=5, y=2), 0) "); } - impl CursorTest { - fn goto_definition(&self) -> String { - let Some(targets) = goto_definition(&self.db, self.cursor.file, self.cursor.offset) - else { - return "No goto target found".to_string(); - }; - - if targets.is_empty() { - return "No definitions found".to_string(); - } - - let source = targets.range; - self.render_diagnostics( - targets - .into_iter() - .map(|target| GotoDefinitionDiagnostic::new(source, &target)), - ) - } - } - #[test] fn goto_definition_overload_type_disambiguated1() { let test = CursorTest::builder() @@ -1130,6 +1110,315 @@ def ab(a: int, *, c: int): ... "#); } + #[test] + fn goto_definition_binary_operator() { + let test = CursorTest::builder() + .source( + "main.py", + " +class Test: + def __add__(self, other): + return Test() + + +a = Test() +b = Test() + +a + b +", + ) + .build(); + + assert_snapshot!(test.goto_definition(), @r" + info[goto-definition]: Definition + --> main.py:3:9 + | + 2 | class Test: + 3 | def __add__(self, other): + | ^^^^^^^ + 4 | return Test() + | + info: Source + --> main.py:10:3 + | + 8 | b = Test() + 9 | + 10 | a + b + | ^ + | + "); + } + + #[test] + fn goto_definition_binary_operator_reflected_dunder() { + let test = CursorTest::builder() + .source( + "main.py", + " +class A: + def __radd__(self, other) -> A: + return self + +class B: ... + +B() + A() +", + ) + .build(); + + assert_snapshot!(test.goto_definition(), @r" + info[goto-definition]: Definition + --> main.py:3:9 + | + 2 | class A: + 3 | def __radd__(self, other) -> A: + | ^^^^^^^^ + 4 | return self + | + info: Source + --> main.py:8:5 + | + 6 | class B: ... + 7 | + 8 | B() + A() + | ^ + | + "); + } + + #[test] + fn goto_definition_binary_operator_no_spaces_before_operator() { + let test = CursorTest::builder() + .source( + "main.py", + " +class Test: + def __add__(self, other): + return Test() + + +a = Test() +b = Test() + +a+b +", + ) + .build(); + + assert_snapshot!(test.goto_definition(), @r" + info[goto-definition]: Definition + --> main.py:3:9 + | + 2 | class Test: + 3 | def __add__(self, other): + | ^^^^^^^ + 4 | return Test() + | + info: Source + --> main.py:10:2 + | + 8 | b = Test() + 9 | + 10 | a+b + | ^ + | + "); + } + + #[test] + fn goto_definition_binary_operator_no_spaces_after_operator() { + let test = CursorTest::builder() + .source( + "main.py", + " +class Test: + def __add__(self, other): + return Test() + + +a = Test() +b = Test() + +a+b +", + ) + .build(); + + assert_snapshot!(test.goto_definition(), @r" + info[goto-definition]: Definition + --> main.py:8:1 + | + 7 | a = Test() + 8 | b = Test() + | ^ + 9 | + 10 | a+b + | + info: Source + --> main.py:10:3 + | + 8 | b = Test() + 9 | + 10 | a+b + | ^ + | + "); + } + + #[test] + fn goto_definition_binary_operator_comment() { + let test = CursorTest::builder() + .source( + "main.py", + " +class Test: + def __add__(self, other): + return Test() + + +( + Test() # comment + + Test() +) +", + ) + .build(); + + assert_snapshot!(test.goto_definition(), @"No goto target found"); + } + + #[test] + fn goto_definition_unary_operator() { + let test = CursorTest::builder() + .source( + "main.py", + " +class Test: + def __bool__(self) -> bool: ... + +a = Test() + +not a +", + ) + .build(); + + assert_snapshot!(test.goto_definition(), @r" + info[goto-definition]: Definition + --> main.py:3:9 + | + 2 | class Test: + 3 | def __bool__(self) -> bool: ... + | ^^^^^^^^ + 4 | + 5 | a = Test() + | + info: Source + --> main.py:7:1 + | + 5 | a = Test() + 6 | + 7 | not a + | ^^^ + | + "); + } + + #[test] + fn goto_definition_unary_after_operator() { + let test = CursorTest::builder() + .source( + "main.py", + " +class Test: + def __bool__(self) -> bool: ... + +a = Test() + +not a +", + ) + .build(); + + assert_snapshot!(test.goto_definition(), @r" + info[goto-definition]: Definition + --> main.py:3:9 + | + 2 | class Test: + 3 | def __bool__(self) -> bool: ... + | ^^^^^^^^ + 4 | + 5 | a = Test() + | + info: Source + --> main.py:7:1 + | + 5 | a = Test() + 6 | + 7 | not a + | ^^^ + | + "); + } + + #[test] + fn goto_definition_unary_between_operator_and_operand() { + let test = CursorTest::builder() + .source( + "main.py", + " +class Test: + def __bool__(self) -> bool: ... + +a = Test() + +-a +", + ) + .build(); + + assert_snapshot!(test.goto_definition(), @r" + info[goto-definition]: Definition + --> main.py:5:1 + | + 3 | def __bool__(self) -> bool: ... + 4 | + 5 | a = Test() + | ^ + 6 | + 7 | -a + | + info: Source + --> main.py:7:2 + | + 5 | a = Test() + 6 | + 7 | -a + | ^ + | + "); + } + + impl CursorTest { + fn goto_definition(&self) -> String { + let Some(targets) = goto_definition(&self.db, self.cursor.file, self.cursor.offset) + else { + return "No goto target found".to_string(); + }; + + if targets.is_empty() { + return "No definitions found".to_string(); + } + + let source = targets.range; + self.render_diagnostics( + targets + .into_iter() + .map(|target| GotoDefinitionDiagnostic::new(source, &target)), + ) + } + } + struct GotoDefinitionDiagnostic { source: FileRange, target: FileRange, diff --git a/crates/ty_ide/src/hover.rs b/crates/ty_ide/src/hover.rs index 241078c76f..b555295678 100644 --- a/crates/ty_ide/src/hover.rs +++ b/crates/ty_ide/src/hover.rs @@ -2514,6 +2514,125 @@ def ab(a: int, *, c: int): "); } + #[test] + fn hover_binary_operator_literal() { + let test = cursor_test( + r#" + result = 5 + 3 + "#, + ); + + assert_snapshot!(test.hover(), @r" + bound method int.__add__(value: int, /) -> int + --------------------------------------------- + Return self+value. + + --------------------------------------------- + ```python + bound method int.__add__(value: int, /) -> int + ``` + --- + ```text + Return self+value. + + ``` + --------------------------------------------- + info[hover]: Hovered content is + --> main.py:2:12 + | + 2 | result = 5 + 3 + | - + | | + | source + | Cursor offset + | + "); + } + + #[test] + fn hover_binary_operator_overload() { + let test = cursor_test( + r#" + from __future__ import annotations + from typing import overload + + class Test: + @overload + def __add__(self, other: Test, /) -> Test: ... + @overload + def __add__(self, other: Other, /) -> Test: ... + def __add__(self, other: Test | Other, /) -> Test: + return self + + class Other: ... + + Test() + Test() + "#, + ); + + // TODO: We should only show the matching overload here. + // https://github.com/astral-sh/ty/issues/73 + assert_snapshot!(test.hover(), @r" + (other: Test, /) -> Test + (other: Other, /) -> Test + --------------------------------------------- + ```python + (other: Test, /) -> Test + (other: Other, /) -> Test + ``` + --------------------------------------------- + info[hover]: Hovered content is + --> main.py:15:8 + | + 13 | class Other: ... + 14 | + 15 | Test() + Test() + | - + | | + | source + | Cursor offset + | + "); + } + + #[test] + fn hover_binary_operator_union() { + let test = cursor_test( + r#" + from __future__ import annotations + + class Test: + def __add__(self, other: Other, /) -> Other: + return other + + class Other: + def __add__(self, other: Other, /) -> Other: + return self + + def _(a: Test | Other): + a + Other() + "#, + ); + + assert_snapshot!(test.hover(), @r" + (bound method Test.__add__(other: Other, /) -> Other) | (bound method Other.__add__(other: Other, /) -> Other) + --------------------------------------------- + ```python + (bound method Test.__add__(other: Other, /) -> Other) | (bound method Other.__add__(other: Other, /) -> Other) + ``` + --------------------------------------------- + info[hover]: Hovered content is + --> main.py:13:7 + | + 12 | def _(a: Test | Other): + 13 | a + Other() + | ^- Cursor offset + | | + | source + | + "); + } + impl CursorTest { fn hover(&self) -> String { use std::fmt::Write; diff --git a/crates/ty_ide/src/references.rs b/crates/ty_ide/src/references.rs index 1dcd747418..62716fd6a5 100644 --- a/crates/ty_ide/src/references.rs +++ b/crates/ty_ide/src/references.rs @@ -18,6 +18,7 @@ use ruff_python_ast::{ self as ast, AnyNodeRef, visitor::source_order::{SourceOrderVisitor, TraversalSignal}, }; +use ruff_python_parser::Tokens; use ruff_text_size::{Ranged, TextRange}; use ty_python_semantic::ImportAliasResolution; @@ -127,6 +128,7 @@ fn references_for_file( target_definitions, references, mode, + tokens: module.tokens(), target_text, ancestors: Vec::new(), }; @@ -156,6 +158,7 @@ fn is_symbol_externally_visible(goto_target: &GotoTarget<'_>) -> bool { struct LocalReferencesFinder<'a> { db: &'a dyn Db, file: File, + tokens: &'a Tokens, target_definitions: &'a [NavigationTarget], references: &'a mut Vec, mode: ReferencesMode, @@ -282,7 +285,9 @@ impl LocalReferencesFinder<'_> { // where the identifier might be a multi-part module name. let offset = covering_node.node().start(); - if let Some(goto_target) = GotoTarget::from_covering_node(covering_node, offset) { + if let Some(goto_target) = + GotoTarget::from_covering_node(covering_node, offset, self.tokens) + { // Get the definitions for this goto target if let Some(current_definitions_nav) = goto_target .get_definition_targets(self.file, self.db, ImportAliasResolution::PreserveAliases) diff --git a/crates/ty_python_semantic/src/lib.rs b/crates/ty_python_semantic/src/lib.rs index 5f41200522..ea0b492b7b 100644 --- a/crates/ty_python_semantic/src/lib.rs +++ b/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; diff --git a/crates/ty_python_semantic/src/semantic_model.rs b/crates/ty_python_semantic/src/semantic_model.rs index beb2c7f968..a7db9d5698 100644 --- a/crates/ty_python_semantic/src/semantic_model.rs +++ b/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 } diff --git a/crates/ty_python_semantic/src/types.rs b/crates/ty_python_semantic/src/types.rs index aa13a13cd4..8fdc21e2b5 100644 --- a/crates/ty_python_semantic/src/types.rs +++ b/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); diff --git a/crates/ty_python_semantic/src/types/call.rs b/crates/ty_python_semantic/src/types/call.rs index 8c00ab3479..e2fb7dac96 100644 --- a/crates/ty_python_semantic/src/types/call.rs +++ b/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, 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> 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> for CallBinOpError { + fn from(value: CallDunderError<'_>) -> Self { + match value { + CallDunderError::CallError(_, _) => Self::CallError, + CallDunderError::MethodNotAvailable | CallDunderError::PossiblyUnbound(_) => { + CallBinOpError::NotSupported + } + } + } +} diff --git a/crates/ty_python_semantic/src/types/call/bind.rs b/crates/ty_python_semantic/src/types/call/bind.rs index 9c4df3c22c..dd1b2215c6 100644 --- a/crates/ty_python_semantic/src/types/call/bind.rs +++ b/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>), diff --git a/crates/ty_python_semantic/src/types/ide_support.rs b/crates/ty_python_semantic/src/types/ide_support.rs index d02011390b..b8dbc97a52 100644 --- a/crates/ty_python_semantic/src/types/ide_support.rs +++ b/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>, 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>, 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). diff --git a/crates/ty_python_semantic/src/types/infer/builder.rs b/crates/ty_python_semantic/src/types/infer/builder.rs index 821e650886..d67a39dab0 100644 --- a/crates/ty_python_semantic/src/types/infer/builder.rs +++ b/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(), } }