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

Co-authored-by: Alex Waygood <Alex.Waygood@Gmail.com>
This commit is contained in:
Micha Reiser 2025-10-21 19:25:41 +02:00 committed by GitHub
parent 2dbca6370b
commit 9d1ffd605c
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
12 changed files with 774 additions and 121 deletions

View file

@ -486,7 +486,7 @@ impl TokenKind {
///
/// [`as_unary_operator`]: TokenKind::as_unary_operator
#[inline]
pub(crate) const fn as_unary_arithmetic_operator(self) -> Option<UnaryOp> {
pub const fn as_unary_arithmetic_operator(self) -> Option<UnaryOp> {
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<UnaryOp> {
pub const fn as_unary_operator(self) -> Option<UnaryOp> {
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<BoolOp> {
pub const fn as_bool_operator(self) -> Option<BoolOp> {
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<Operator> {
pub const fn as_binary_operator(self) -> Option<Operator> {
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<Operator> {
pub const fn as_augmented_assign_operator(self) -> Option<Operator> {
Some(match self {
TokenKind::PlusEqual => Operator::Add,
TokenKind::MinusEqual => Operator::Sub,

View file

@ -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<GotoTarget<'a>> {
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<ResolvedDefinition<'db>> {
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<CURSOR>+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.

View file

@ -798,26 +798,6 @@ my_func(my_other_func(a<CURSOR>b=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 <CURSOR>+ 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() <CURSOR>+ 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<CURSOR>+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+<CURSOR>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() <CURSOR># 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()
<CURSOR>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<CURSOR> 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()
-<CURSOR>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,

View file

@ -2514,6 +2514,125 @@ def ab(a: int, *, c: int):
");
}
#[test]
fn hover_binary_operator_literal() {
let test = cursor_test(
r#"
result = 5 <CURSOR>+ 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() <CURSOR>+ 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 +<CURSOR> 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;

View file

@ -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<ReferenceTarget>,
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)

View file

@ -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;

View file

@ -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
}

View file

@ -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);

View file

@ -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
}
}
}
}

View file

@ -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>),

View file

@ -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).

View file

@ -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(),
}
}