[ty] Handle overloads in rename

crates/ty_ide/src/goto.rs

@@ -392,11 +392,30 @@ impl GotoTarget<'_> {
             GotoTarget::Expression(expression) => {
                 definitions_for_expression(model, *expression, alias_resolution)
             }
-            // For already-defined symbols, they are their own definitions
-            GotoTarget::FunctionDef(function) => Some(vec![ResolvedDefinition::Definition(
-                function.definition(model),
-            )]),
 
+            GotoTarget::FunctionDef(function) => {
+                let self_definition = function.definition(model);
+
+                let mut definitions = Vec::new();
+
+                // TODO: Skip the implementation in go to definition
+                // TODO: Only take the implementation in go to declaration
+                if let Some(ty) = function.inferred_type(model).as_function_literal() {
+                    let overload = ty.resolve_overload(model.db());
+
+                    definitions.extend(overload.iter_overloads_and_implementation(model.db()).map(
+                        |literal| ResolvedDefinition::Definition(literal.definition(model.db())),
+                    ))
+                }
+
+                if definitions.is_empty() {
+                    definitions.push(ResolvedDefinition::Definition(self_definition));
+                }
+
+                Some(definitions)
+            }
+
+            // For already-defined symbols, they are their own definitions
             GotoTarget::ClassDef(class) => Some(vec![ResolvedDefinition::Definition(
                 class.definition(model),
             )]),

crates/ty_ide/src/goto_declaration.rs

@@ -2785,8 +2785,9 @@ def ab(a: int, *, c: int): ...
 
     impl CursorTest {
         fn goto_declaration(&self) -> String {
-            let Some(targets) = goto_declaration(&self.db, self.cursor.file, self.cursor.offset)
+            let Some(targets) = salsa::attach(&self.db, || {
-            else {
+                goto_declaration(&self.db, self.cursor.file, self.cursor.offset)
+            }) else {
                 return "No goto target found".to_string();
             };
 

crates/ty_ide/src/goto_definition.rs

@@ -1069,6 +1069,54 @@ def ab(a: int, *, c: int): ...
         "#);
     }
 
+    #[test]
+    fn goto_definition_on_overloaded_function_literal() {
+        let test = CursorTest::builder()
+            .source(
+                "main.py",
+                "
+from typing import overload, Any
+
+@overload
+def ab(a: int): ...
+
+@overload
+def ab<CURSOR>(a: int, b: int): ...
+
+def ab(a: int, c: Any):
+    return 1
+",
+            )
+            .build();
+
+        assert_snapshot!(test.goto_definition(), @r"
+        info[goto-definition]: Go to definition
+          --> main.py:8:5
+           |
+         7 | @overload
+         8 | def ab(a: int, b: int): ...
+           |     ^^ Clicking here
+         9 |
+        10 | def ab(a: int, c: Any):
+           |
+        info: Found 3 definitions
+          --> main.py:5:5
+           |
+         4 | @overload
+         5 | def ab(a: int): ...
+           |     --
+         6 |
+         7 | @overload
+         8 | def ab(a: int, b: int): ...
+           |     --
+         9 |
+        10 | def ab(a: int, c: Any):
+           |     --
+        11 |     return 1
+           |
+        ");
+    }
+
     #[test]
     fn goto_definition_binary_operator() {
         let test = CursorTest::builder()
@@ -1697,8 +1745,9 @@ Traceb<CURSOR>ackType
 
     impl CursorTest {
         fn goto_definition(&self) -> String {
-            let Some(targets) = goto_definition(&self.db, self.cursor.file, self.cursor.offset)
+            let Some(targets) = salsa::attach(&self.db, || {
-            else {
+                goto_definition(&self.db, self.cursor.file, self.cursor.offset)
+            }) else {
                 return "No goto target found".to_string();
             };
 

crates/ty_ide/src/goto_type_definition.rs

@@ -1900,9 +1900,9 @@ def function():
 
     impl CursorTest {
         fn goto_type_definition(&self) -> String {
-            let Some(targets) =
+            let Some(targets) = salsa::attach(&self.db, || {
                 goto_type_definition(&self.db, self.cursor.file, self.cursor.offset)
-            else {
+            }) else {
                 return "No goto target found".to_string();
             };
 

crates/ty_ide/src/rename.rs

@@ -106,14 +106,23 @@ mod tests {
         }
 
         fn rename(&self, new_name: &str) -> String {
-            let Some(_) = can_rename(&self.db, self.cursor.file, self.cursor.offset) else {
+            let rename_results = salsa::attach(&self.db, || {
-                return "Cannot rename".to_string();
+                let Some(_) = can_rename(&self.db, self.cursor.file, self.cursor.offset) else {
-            };
+                    return Err("Cannot rename".to_string());
+                };
 
-            let Some(rename_results) =
+                let Some(rename_results) =
-                rename(&self.db, self.cursor.file, self.cursor.offset, new_name)
+                    rename(&self.db, self.cursor.file, self.cursor.offset, new_name)
-            else {
+                else {
-                return "Cannot rename".to_string();
+                    return Err("Cannot rename".to_string());
+                };
+
+                Ok(rename_results)
+            });
+
+            let rename_results = match rename_results {
+                Ok(rename_results) => rename_results,
+                Err(err) => return err,
             };
 
             if rename_results.is_empty() {
@@ -1533,24 +1542,283 @@ result = func(10, y=20)
             )
             .build();
 
+        assert_snapshot!(test.rename("better_name"), @r###"
+        info[rename]: Rename symbol (found 6 locations)
+          --> lib.py:5:5
+           |
+         4 | @overload
+         5 | def test() -> None: ...
+           |     ^^^^
+         6 | @overload
+         7 | def test(a: str) -> str: ...
+           |     ----
+         8 | @overload
+         9 | def test(a: int) -> int: ...
+           |     ----
+        10 |
+        11 | def test(a: Any) -> Any:
+           |     ----
+        12 |     return a
+           |
+          ::: main.py:2:17
+           |
+         2 | from lib import test
+           |                 ----
+         3 |
+         4 | test("test")
+           | ----
+           |
+        "###);
+    }
+
+    #[test]
+    fn rename_non_first_overloaded_function_declaration() {
+        let test = CursorTest::builder()
+            .source(
+                "lib.py",
+                r#"
+                from typing import overload, Any
+
+                @overload
+                def test() -> None: ...
+                @overload
+                def test<CURSOR>(a: str) -> str: ...
+                @overload
+                def test(a: int) -> int: ...
+
+                def test(a: Any) -> Any:
+                    return a
+                "#,
+            )
+            .source(
+                "main.py",
+                r#"
+                from lib import test
+
+                test("test")
+                "#,
+            )
+            .build();
+
+        assert_snapshot!(test.rename("better_name"), @r###"
+        info[rename]: Rename symbol (found 6 locations)
+          --> lib.py:5:5
+           |
+         4 | @overload
+         5 | def test() -> None: ...
+           |     ^^^^
+         6 | @overload
+         7 | def test(a: str) -> str: ...
+           |     ----
+         8 | @overload
+         9 | def test(a: int) -> int: ...
+           |     ----
+        10 |
+        11 | def test(a: Any) -> Any:
+           |     ----
+        12 |     return a
+           |
+          ::: main.py:2:17
+           |
+         2 | from lib import test
+           |                 ----
+         3 |
+         4 | test("test")
+           | ----
+           |
+        "###);
+    }
+
+    #[test]
+    fn rename_overloaded_function_definition() {
+        let test = CursorTest::builder()
+            .source(
+                "lib.py",
+                r#"
+                from typing import overload, Any
+
+                @overload
+                def test() -> None: ...
+                @overload
+                def test(a: str) -> str: ...
+                @overload
+                def test(a: int) -> int: ...
+
+                def test<CURSOR>(a: Any) -> Any:
+                    return a
+                "#,
+            )
+            .source(
+                "main.py",
+                r#"
+                from lib import test
+
+                test("test")
+                "#,
+            )
+            .build();
+
+        assert_snapshot!(test.rename("better_name"), @r###"
+        info[rename]: Rename symbol (found 6 locations)
+          --> lib.py:5:5
+           |
+         4 | @overload
+         5 | def test() -> None: ...
+           |     ^^^^
+         6 | @overload
+         7 | def test(a: str) -> str: ...
+           |     ----
+         8 | @overload
+         9 | def test(a: int) -> int: ...
+           |     ----
+        10 |
+        11 | def test(a: Any) -> Any:
+           |     ----
+        12 |     return a
+           |
+          ::: main.py:2:17
+           |
+         2 | from lib import test
+           |                 ----
+         3 |
+         4 | test("test")
+           | ----
+           |
+        "###);
+    }
+
+    #[test]
+    fn rename_overloaded_function_with_conditional_definitions() {
+        let test = CursorTest::builder()
+            .source(
+                "lib.py",
+                r#"
+                from typing import overload, Any
+                def foo() -> bool: ...
+
+                @overload
+                def test() -> None: ...
+
+                if foo():
+                    @overload
+                    def test<CURSOR>(a: str) -> str: ...
+                else:
+                    @overload
+                    def test(a: int) -> int: ...
+
+                def test(a: Any) -> Any:
+                    return a
+                "#,
+            )
+            .source(
+                "main.py",
+                r#"
+                from lib import test
+
+                test("test")
+                "#,
+            )
+            .build();
+
         assert_snapshot!(test.rename("better_name"), @r#"
-        info[rename]: Rename symbol (found 3 locations)
+        info[rename]: Rename symbol (found 6 locations)
-         --> lib.py:5:5
+          --> lib.py:6:5
-          |
+           |
-        4 | @overload
+         5 | @overload
-        5 | def test() -> None: ...
+         6 | def test() -> None: ...
-          |     ^^^^
+           |     ^^^^
-        6 | @overload
+         7 |
-        7 | def test(a: str) -> str: ...
+         8 | if foo():
-          |
+         9 |     @overload
-         ::: main.py:2:17
+        10 |     def test(a: str) -> str: ...
-          |
+           |         ----
-        2 | from lib import test
+        11 | else:
-          |                 ----
+        12 |     @overload
-        3 |
+        13 |     def test(a: int) -> int: ...
-        4 | test("test")
+           |         ----
-          | ----
+        14 |
-          |
+        15 | def test(a: Any) -> Any:
+           |     ----
+        16 |     return a
+           |
+          ::: main.py:2:17
+           |
+         2 | from lib import test
+           |                 ----
+         3 |
+         4 | test("test")
+           | ----
+           |
+        "#);
+    }
+
+    #[test]
+    fn rename_overloaded_function_with_always_true_conditional_definitions() {
+        // Ideally, the overload in the `else` branch would be renamed too
+        // but it gets inferred as `Never`.
+        // The alternative would be to use `definitions_for_name` and simply rename all
+        // symbols (or functions) with the same name, ignoring whether they are indeed overloads
+        // of the same function. However, renaming all symbols comes at the risk that we
+        // rename more symbols than we should, that's why we're erroring on the side of caution
+        // here and only rename the "reachable" sybmols.
+        let test = CursorTest::builder()
+            .source(
+                "lib.py",
+                r#"
+                from typing import overload, Any
+
+                @overload
+                def test() -> None: ...
+
+                if True:
+                    @overload
+                    def test<CURSOR>(a: str) -> str: ...
+                else:
+                    @overload
+                    def test(a: int) -> int: ...
+
+                def test(a: Any) -> Any:
+                    return a
+                "#,
+            )
+            .source(
+                "main.py",
+                r#"
+                from lib import test
+
+                test("test")
+                "#,
+            )
+            .build();
+
+        assert_snapshot!(test.rename("better_name"), @r#"
+        info[rename]: Rename symbol (found 5 locations)
+          --> lib.py:5:5
+           |
+         4 | @overload
+         5 | def test() -> None: ...
+           |     ^^^^
+         6 |
+         7 | if True:
+         8 |     @overload
+         9 |     def test(a: str) -> str: ...
+           |         ----
+        10 | else:
+        11 |     @overload
+        12 |     def test(a: int) -> int: ...
+        13 |
+        14 | def test(a: Any) -> Any:
+           |     ----
+        15 |     return a
+           |
+          ::: main.py:2:17
+           |
+         2 | from lib import test
+           |                 ----
+         3 |
+         4 | test("test")
+           | ----
+           |
         "#);
     }
 
@@ -1585,25 +1853,33 @@ result = func(10, y=20)
             )
             .build();
 
-        assert_snapshot!(test.rename("better_name"), @r#"
+        assert_snapshot!(test.rename("better_name"), @r###"
-        info[rename]: Rename symbol (found 2 locations)
+        info[rename]: Rename symbol (found 5 locations)
-         --> lib.py:6:9
+          --> lib.py:6:9
-          |
+           |
-        4 | class Test:
+         4 | class Test:
-        5 |     @overload
+         5 |     @overload
-        6 |     def test() -> None: ...
+         6 |     def test() -> None: ...
-          |         ^^^^
+           |         ^^^^
-        7 |     @overload
+         7 |     @overload
-        8 |     def test(a: str) -> str: ...
+         8 |     def test(a: str) -> str: ...
-          |
+           |         ----
-         ::: main.py:4:8
+         9 |     @overload
-          |
+        10 |     def test(a: int) -> int: ...
-        2 | from lib import Test
+           |         ----
-        3 |
+        11 |
-        4 | Test().test("test")
+        12 |     def test(a: Any) -> Any:
-          |        ----
+           |         ----
-          |
+        13 |         return a
-        "#);
+           |
+          ::: main.py:4:8
+           |
+         2 | from lib import Test
+         3 |
+         4 | Test().test("test")
+           |        ----
+           |
+        "###);
     }
 
     #[test]
@@ -1635,25 +1911,33 @@ result = func(10, y=20)
             )
             .build();
 
-        assert_snapshot!(test.rename("better_name"), @r#"
+        assert_snapshot!(test.rename("better_name"), @r###"
-        info[rename]: Rename symbol (found 3 locations)
+        info[rename]: Rename symbol (found 6 locations)
-         --> main.py:2:17
+          --> main.py:2:17
-          |
+           |
-        2 | from lib import test
+         2 | from lib import test
-          |                 ^^^^
+           |                 ^^^^
-        3 |
+         3 |
-        4 | test("test")
+         4 | test("test")
-          | ----
+           | ----
-          |
+           |
-         ::: lib.py:5:5
+          ::: lib.py:5:5
-          |
+           |
-        4 | @overload
+         4 | @overload
-        5 | def test() -> None: ...
+         5 | def test() -> None: ...
-          |     ----
+           |     ----
-        6 | @overload
+         6 | @overload
-        7 | def test(a: str) -> str: ...
+         7 | def test(a: str) -> str: ...
-          |
+           |     ----
-        "#);
+         8 | @overload
+         9 | def test(a: int) -> int: ...
+           |     ----
+        10 |
+        11 | def test(a: Any) -> Any:
+           |     ----
+        12 |     return a
+           |
+        "###);
     }
 
     #[test]

crates/ty_python_semantic/src/semantic_model.rs

@@ -241,6 +241,27 @@ impl<'db> SemanticModel<'db> {
         let index = semantic_index(self.db, self.file);
         match self.node_in_ast(node) {
             ast::AnyNodeRef::Identifier(identifier) => index.try_expression_scope_id(identifier),
+            ast::AnyNodeRef::StmtFunctionDef(def) => {
+                Some(def.definition(self).scope(self.db).file_scope_id(self.db))
+            }
+            ast::AnyNodeRef::StmtClassDef(class) => {
+                Some(class.definition(self).scope(self.db).file_scope_id(self.db))
+            }
+            ast::AnyNodeRef::Parameter(param) => {
+                Some(param.definition(self).scope(self.db).file_scope_id(self.db))
+            }
+            ast::AnyNodeRef::ParameterWithDefault(param) => {
+                Some(param.definition(self).scope(self.db).file_scope_id(self.db))
+            }
+            ast::AnyNodeRef::ExceptHandlerExceptHandler(handler) => Some(
+                handler
+                    .definition(self)
+                    .scope(self.db)
+                    .file_scope_id(self.db),
+            ),
+            ast::AnyNodeRef::TypeParamTypeVar(var) => {
+                Some(var.definition(self).scope(self.db).file_scope_id(self.db))
+            }
             node => match node.as_expr_ref() {
                 // If we couldn't identify a specific
                 // expression that we're in, then just
@@ -395,12 +416,12 @@ impl<'db> Completion<'db> {
     }
 }
 
-pub trait HasType {
+pub trait HasType<'db> {
     /// Returns the inferred type of `self`.
     ///
     /// ## Panics
     /// May panic if `self` is from another file than `model`.
-    fn inferred_type<'db>(&self, model: &SemanticModel<'db>) -> Type<'db>;
+    fn inferred_type(&self, model: &SemanticModel<'db>) -> Type<'db>;
 }
 
 pub trait HasDefinition {
@@ -411,8 +432,8 @@ pub trait HasDefinition {
     fn definition<'db>(&self, model: &SemanticModel<'db>) -> Definition<'db>;
 }
 
-impl HasType for ast::ExprRef<'_> {
+impl<'db> HasType<'db> for ast::ExprRef<'_> {
-    fn inferred_type<'db>(&self, model: &SemanticModel<'db>) -> Type<'db> {
+    fn inferred_type(&self, model: &SemanticModel<'db>) -> Type<'db> {
         let index = semantic_index(model.db, model.file);
         // TODO(#1637): semantic tokens is making this crash even with
         // `try_expr_ref_in_ast` guarding this, for now just use `try_expression_scope_id`.
@@ -429,9 +450,9 @@ impl HasType for ast::ExprRef<'_> {
 
 macro_rules! impl_expression_has_type {
     ($ty: ty) => {
-        impl HasType for $ty {
+        impl<'db> HasType<'db> for $ty {
             #[inline]
-            fn inferred_type<'db>(&self, model: &SemanticModel<'db>) -> Type<'db> {
+            fn inferred_type(&self, model: &SemanticModel<'db>) -> Type<'db> {
                 let expression_ref = ExprRef::from(self);
                 expression_ref.inferred_type(model)
             }
@@ -473,8 +494,8 @@ impl_expression_has_type!(ast::ExprTuple);
 impl_expression_has_type!(ast::ExprSlice);
 impl_expression_has_type!(ast::ExprIpyEscapeCommand);
 
-impl HasType for ast::Expr {
+impl<'db> HasType<'db> for ast::Expr {
-    fn inferred_type<'db>(&self, model: &SemanticModel<'db>) -> Type<'db> {
+    fn inferred_type(&self, model: &SemanticModel<'db>) -> Type<'db> {
         match self {
             Expr::BoolOp(inner) => inner.inferred_type(model),
             Expr::Named(inner) => inner.inferred_type(model),
@@ -523,11 +544,11 @@ macro_rules! impl_binding_has_ty_def {
             }
         }
 
-        impl HasType for $ty {
+        impl<'db> HasType<'db> for $ty {
             #[inline]
-            fn inferred_type<'db>(&self, model: &SemanticModel<'db>) -> Type<'db> {
+            fn inferred_type(&self, model: &SemanticModel<'db>) -> Type<'db> {
                 let binding = HasDefinition::definition(self, model);
-                binding_type(model.db, binding)
+                binding.inferred_type(model)
             }
         }
     };
@@ -540,8 +561,8 @@ impl_binding_has_ty_def!(ast::ParameterWithDefault);
 impl_binding_has_ty_def!(ast::ExceptHandlerExceptHandler);
 impl_binding_has_ty_def!(ast::TypeParamTypeVar);
 
-impl HasType for ast::Alias {
+impl<'db> HasType<'db> for ast::Alias {
-    fn inferred_type<'db>(&self, model: &SemanticModel<'db>) -> Type<'db> {
+    fn inferred_type(&self, model: &SemanticModel<'db>) -> Type<'db> {
         if &self.name == "*" {
             return Type::Never;
         }
@@ -550,6 +571,12 @@ impl HasType for ast::Alias {
     }
 }
 
+impl<'db> HasType<'db> for Definition<'db> {
+    fn inferred_type(&self, model: &SemanticModel<'db>) -> Type<'db> {
+        binding_type(model.db, *self)
+    }
+}
+
 /// Implemented by types for which the semantic index tracks their scope.
 pub(crate) trait HasTrackedScope: HasNodeIndex {}
 

crates/ty_python_semantic/src/types.rs

@@ -1278,7 +1278,7 @@ impl<'db> Type<'db> {
         }
     }
 
-    pub(crate) const fn as_function_literal(self) -> Option<FunctionType<'db>> {
+    pub const fn as_function_literal(self) -> Option<FunctionType<'db>> {
         match self {
             Type::FunctionLiteral(function_type) => Some(function_type),
             _ => None,

crates/ty_python_semantic/src/types/function.rs

@@ -52,6 +52,7 @@
 use std::str::FromStr;
 
 use bitflags::bitflags;
+use itertools::Itertools;
 use ruff_db::diagnostic::{Annotation, DiagnosticId, Severity, Span};
 use ruff_db::files::{File, FileRange};
 use ruff_db::parsed::{ParsedModuleRef, parsed_module};
@@ -63,7 +64,7 @@ use crate::place::{Definedness, Place, place_from_bindings};
 use crate::semantic_index::ast_ids::HasScopedUseId;
 use crate::semantic_index::definition::Definition;
 use crate::semantic_index::scope::ScopeId;
-use crate::semantic_index::{FileScopeId, SemanticIndex, semantic_index};
+use crate::semantic_index::{FileScopeId, SemanticIndex, place_table, semantic_index, use_def_map};
 use crate::types::call::{Binding, CallArguments};
 use crate::types::constraints::ConstraintSet;
 use crate::types::context::InferContext;
@@ -260,7 +261,7 @@ impl<'db> OverloadLiteral<'db> {
         self.decorators(db).contains(decorator)
     }
 
-    pub(crate) fn is_overload(self, db: &dyn Db) -> bool {
+    pub fn is_overload(self, db: &dyn Db) -> bool {
         self.has_known_decorator(db, FunctionDecorators::OVERLOAD)
     }
 
@@ -350,7 +351,7 @@ impl<'db> OverloadLiteral<'db> {
     /// calling query is not in the same file as this function is defined in, then this will create
     /// a cross-module dependency directly on the full AST which will lead to cache
     /// over-invalidation.
-    fn definition(self, db: &'db dyn Db) -> Definition<'db> {
+    pub fn definition(self, db: &'db dyn Db) -> Definition<'db> {
         let body_scope = self.body_scope(db);
         let index = semantic_index(db, body_scope.file(db));
         index.expect_single_definition(body_scope.node(db).expect_function())
@@ -363,7 +364,7 @@ impl<'db> OverloadLiteral<'db> {
         // here to get the previous function definition with the same name.
         let scope = self.definition(db).scope(db);
         let module = parsed_module(db, self.file(db)).load(db);
-        let use_def = semantic_index(db, scope.file(db)).use_def_map(scope.file_scope_id(db));
+        let use_def = use_def_map(db, scope);
         let use_id = self
             .body_scope(db)
             .node(db)
@@ -563,7 +564,7 @@ impl<'db> OverloadLiteral<'db> {
 #[salsa::interned(debug, heap_size=ruff_memory_usage::heap_size)]
 #[derive(PartialOrd, Ord)]
 pub struct FunctionLiteral<'db> {
-    pub(crate) last_definition: OverloadLiteral<'db>,
+    pub last_definition: OverloadLiteral<'db>,
 }
 
 // The Salsa heap is tracked separately.
@@ -721,7 +722,7 @@ impl<'db> FunctionLiteral<'db> {
 #[salsa::interned(debug, heap_size=ruff_memory_usage::heap_size)]
 #[derive(PartialOrd, Ord)]
 pub struct FunctionType<'db> {
-    pub(crate) literal: FunctionLiteral<'db>,
+    pub literal: FunctionLiteral<'db>,
 
     /// Contains a potentially modified signature for this function literal, in case certain operations
     /// (like type mappings) have been applied to it.
@@ -945,13 +946,44 @@ impl<'db> FunctionType<'db> {
 
     /// Returns an iterator of all of the definitions of this function, including both overload
     /// signatures and any implementation, all in source order.
-    pub(crate) fn iter_overloads_and_implementation(
+    pub fn iter_overloads_and_implementation(
         self,
         db: &'db dyn Db,
     ) -> impl Iterator<Item = OverloadLiteral<'db>> + 'db {
         self.literal(db).iter_overloads_and_implementation(db)
     }
 
+    pub fn resolve_overload(self, db: &'db dyn Db) -> FunctionType<'db> {
+        let scope = self.definition(db).scope(db);
+        let module = parsed_module(db, self.file(db)).load(db);
+        let use_def = use_def_map(db, scope);
+        let place_table = place_table(db, scope);
+
+        let function = self.node(db, scope.file(db), &module);
+        let symbol_id = place_table.symbol_id(&function.name.id).unwrap();
+
+        let overload_literal = self.literal(db).last_definition(db);
+
+        // Not overloaded
+        if !overload_literal.is_overload(db) {
+            return self;
+        }
+
+        // Find the last binding in the outer scope that contains self as one possible overload.
+        use_def
+            .end_of_scope_symbol_bindings(symbol_id.into())
+            .filter_map(|binding| {
+                let ty = binding_type(db, binding.binding.definition()?).as_function_literal()?;
+
+                ty.literal(db)
+                    .iter_overloads_and_implementation(db)
+                    .contains(&overload_literal)
+                    .then_some(ty)
+            })
+            .last()
+            .unwrap_or(self)
+    }
+
     pub(crate) fn first_overload_or_implementation(self, db: &'db dyn Db) -> OverloadLiteral<'db> {
         self.iter_overloads_and_implementation(db)
             .next()
@@ -971,7 +1003,7 @@ impl<'db> FunctionType<'db> {
     /// Were this not a salsa query, then the calling query
     /// would depend on the function's AST and rerun for every change in that file.
     #[salsa::tracked(returns(ref), cycle_initial=signature_cycle_initial, heap_size=ruff_memory_usage::heap_size)]
-    pub(crate) fn signature(self, db: &'db dyn Db) -> CallableSignature<'db> {
+    pub fn signature(self, db: &'db dyn Db) -> CallableSignature<'db> {
         self.updated_signature(db)
             .cloned()
             .unwrap_or_else(|| self.literal(db).signature(db))
@@ -990,7 +1022,7 @@ impl<'db> FunctionType<'db> {
         returns(ref), cycle_initial=last_definition_signature_cycle_initial,
         heap_size=ruff_memory_usage::heap_size,
     )]
-    pub(crate) fn last_definition_signature(self, db: &'db dyn Db) -> Signature<'db> {
+    pub fn last_definition_signature(self, db: &'db dyn Db) -> Signature<'db> {
         self.updated_last_definition_signature(db)
             .cloned()
             .unwrap_or_else(|| self.literal(db).last_definition_signature(db))

crates/ty_python_semantic/src/types/signatures.rs

@@ -147,7 +147,7 @@ impl<'db> CallableSignature<'db> {
         }
     }
 
-    pub(crate) fn iter(&self) -> std::slice::Iter<'_, Signature<'db>> {
+    pub fn iter(&self) -> std::slice::Iter<'_, Signature<'db>> {
         self.overloads.iter()
     }
 
@@ -772,7 +772,7 @@ impl<'db> Signature<'db> {
     }
 
     /// Return the definition associated with this signature, if any.
-    pub(crate) fn definition(&self) -> Option<Definition<'db>> {
+    pub fn definition(&self) -> Option<Definition<'db>> {
         self.definition
     }