go-to-definition on calls jumps to __init__/__new__/__call__

Summary:
When go-to-definition is invoked on a call expression, jump to the
underlying method definition instead of the class or variable:

- Constructor calls (`Foo(1)`, `mod.Foo(1)`): jump to `__init__` and/or
  `__new__` on the class
- Callable instances (`adder(5)` where `adder` has `__call__`): jump to
  `__call__` on the instance's class
- Regular function calls: unchanged (still jumps to the function def)

The implementation detects call context via `covering_nodes`, resolves the
callee's type, and dispatches to `find_call_target_definitions` which uses
the existing `find_attribute_definition_for_base_type` machinery. Only
`ClassDef` types get constructor lookup; only `ClassType` (instances) get
`__call__` lookup; functions and other callables fall through to normal
go-to-definition.

Note: this stack is a different implementation from D87102818. D87102818 reused the find-references queue which ended up being more complicated and required storing a reverse-mapping in state. it was more correct (it would only go to the constructor pyrefly thought was used). instead, this implementation goes to all `__init__`, `__new__`, `__call__`. it might be more useful to see all of them as separate definitions and is much cleaner to implement.

fixes #1877
fixes #1728

Reviewed By: yangdanny97

Differential Revision: D96921976

fbshipit-source-id: c9bfad0699a421720a75ad146b0db92b33b97d36

Author: kinto0

pyrefly/lib/state/lsp.rs

@@ -1402,6 +1402,40 @@ impl<'a> Transaction<'a> {
         })
     }
 
+    /// When a name or attribute in a call position resolves to a class, find
+    /// `__init__` and `__new__` definitions. When it resolves to a class
+    /// instance, find `__call__`. Returns all found definitions, or empty if
+    /// neither case applies. Does not match functions/callables — those should
+    /// use the normal go-to-definition path.
+    fn find_call_target_definitions(
+        &self,
+        handle: &Handle,
+        preference: FindPreference,
+        ty: Type,
+    ) -> Vec<FindDefinitionItemWithDocstring> {
+        match &ty {
+            Type::ClassDef(_) => {
+                let mut defs = self.find_attribute_definition_for_base_type(
+                    handle,
+                    preference,
+                    ty.clone(),
+                    &dunder::INIT,
+                );
+                defs.extend(self.find_attribute_definition_for_base_type(
+                    handle,
+                    preference,
+                    ty,
+                    &dunder::NEW,
+                ));
+                defs
+            }
+            Type::ClassType(_) => {
+                self.find_attribute_definition_for_base_type(handle, preference, ty, &dunder::CALL)
+            }
+            _ => vec![],
+        }
+    }
+
     pub(crate) fn find_definition_for_base_type(
         &self,
         handle: &Handle,
@@ -1740,6 +1774,23 @@ impl<'a> Transaction<'a> {
                             .map_or(vec![], |item| vec![item])
                     }
                     ExprContext::Load | ExprContext::Del | ExprContext::Invalid => {
+                        // If this name is the callee of a call expression, jump
+                        // to constructor or __call__ definitions when applicable.
+                        if let Some(AnyNodeRef::ExprCall(call)) = covering_nodes.get(1)
+                            && call.func.range() == id.range
+                            && let Some(bindings) = self.get_bindings(handle)
+                        {
+                            let key = Key::BoundName(ShortIdentifier::new(&id));
+                            if bindings.is_valid_key(&key)
+                                && let Some(ty) = self.get_type(handle, &key)
+                            {
+                                let defs =
+                                    self.find_call_target_definitions(handle, preference, ty);
+                                if !defs.is_empty() {
+                                    return defs;
+                                }
+                            }
+                        }
                         // This is a usage of the variable
                         self.find_definition_for_name_use(handle, &id, preference)
                             .map_or(vec![], |item| vec![item])
@@ -1900,6 +1951,18 @@ impl<'a> Transaction<'a> {
                 identifier,
                 context: IdentifierContext::Attribute { base_range, .. },
             }) => {
+                // If this attribute is the callee of a call expression, jump
+                // to constructor or __call__ definitions when applicable.
+                if let Some(AnyNodeRef::ExprAttribute(attr)) = covering_nodes.get(1)
+                    && let Some(AnyNodeRef::ExprCall(call)) = covering_nodes.get(2)
+                    && call.func.range() == attr.range()
+                    && let Some(ty) = self.get_type_trace(handle, attr.range())
+                {
+                    let defs = self.find_call_target_definitions(handle, preference, ty);
+                    if !defs.is_empty() {
+                        return defs;
+                    }
+                }
                 self.find_definition_for_attribute(handle, base_range, identifier.id(), preference)
             }
             Some(IdentifierWithContext {

pyrefly/lib/test/lsp/definition.rs

@@ -2238,10 +2238,6 @@ Definition Result:
     );
 }
 
-// TODO: go-to-definition on constructor calls should jump to __init__/__new__,
-// not the class definition. These tests document the current (wrong) behavior
-// and will be updated when the feature is implemented.
-
 #[test]
 fn goto_def_constructor_call_same_module() {
     let code = r#"
@@ -2259,8 +2255,8 @@ Bar("hello")
 6 | Bar("hello")
       ^
 Definition Result:
-2 | class Bar:
-          ^^^
+3 |     def __init__(self, name: str) -> None:
+            ^^^^^^^^
 "#
         .trim(),
         report.trim(),
@@ -2289,8 +2285,8 @@ Foo(1)
 3 | Foo(1)
       ^
 Definition Result:
-2 | class Foo:
-          ^^^
+3 |     def __init__(self, x: int) -> None:
+            ^^^^^^^^
 
 
 # foo_mod.py
@@ -2326,8 +2322,8 @@ Child(1)
 7 | Child(1)
        ^
 Definition Result:
-4 | class Child(Base):
-          ^^^^^
+3 |     def __init__(self, x: int) -> None:
+            ^^^^^^^^
 
 
 # base_mod.py
@@ -2357,8 +2353,8 @@ Singleton()
 9 | Singleton()
          ^
 Definition Result:
-2 | class Singleton:
-          ^^^^^^^^^
+4 |     def __new__(cls) -> "Singleton":
+            ^^^^^^^
 "#
         .trim(),
         report.trim(),
@@ -2384,8 +2380,11 @@ MyClass()
 8 | MyClass()
        ^
 Definition Result:
-2 | class MyClass:
-          ^^^^^^^
+5 |     def __init__(self) -> None:
+            ^^^^^^^^
+Definition Result:
+3 |     def __new__(cls) -> "MyClass":
+            ^^^^^^^
 "#
         .trim(),
         report.trim(),
@@ -2414,8 +2413,8 @@ foo_mod.Foo(1)
 3 | foo_mod.Foo(1)
             ^
 Definition Result:
-2 | class Foo:
-          ^^^
+3 |     def __init__(self, x: int) -> None:
+            ^^^^^^^^
 
 
 # foo_mod.py
@@ -2424,3 +2423,78 @@ Definition Result:
         report.trim(),
     );
 }
+
+#[test]
+fn goto_def_callable_instance() {
+    let code = r#"
+class Adder:
+    def __call__(self, x: int) -> int:
+        return x + 1
+
+adder = Adder()
+adder(5)
+# ^
+"#;
+    let report = get_batched_lsp_operations_report(&[("main", code)], get_test_report);
+    assert_eq!(
+        r#"
+# main.py
+7 | adder(5)
+      ^
+Definition Result:
+3 |     def __call__(self, x: int) -> int:
+            ^^^^^^^^
+"#
+        .trim(),
+        report.trim(),
+    );
+}
+
+#[test]
+fn goto_def_non_constructor_call_goes_to_function() {
+    let code = r#"
+def foo(x: int) -> int:
+    return x
+
+foo(1)
+# ^
+"#;
+    let report = get_batched_lsp_operations_report(&[("main", code)], get_test_report);
+    assert_eq!(
+        r#"
+# main.py
+5 | foo(1)
+      ^
+Definition Result:
+2 | def foo(x: int) -> int:
+        ^^^
+"#
+        .trim(),
+        report.trim(),
+    );
+}
+
+#[test]
+fn goto_def_class_name_without_call_goes_to_class() {
+    let code = r#"
+class Baz:
+    def __init__(self) -> None:
+        pass
+
+x = Baz
+#   ^
+"#;
+    let report = get_batched_lsp_operations_report(&[("main", code)], get_test_report);
+    assert_eq!(
+        r#"
+# main.py
+6 | x = Baz
+        ^
+Definition Result:
+2 | class Baz:
+          ^^^
+"#
+        .trim(),
+        report.trim(),
+    );
+}

pyrefly/lib/test/lsp/hover.rs

@@ -1020,7 +1020,7 @@ from mymod.submod.deep import Bar
 }
 
 #[test]
-fn hover_on_constructor_shows_instance_type() {
+fn hover_on_constructor() {
     let code = r#"
 class Person:
     def __init__(self, name: str, age: int) -> None: ...
@@ -1030,8 +1030,9 @@ Person()
 "#;
     let report = get_batched_lsp_operations_report_allow_error(&[("main", code)], get_test_report);
     assert!(
-        report
-            .contains("def Person(\n    self: Person,\n    name: str,\n    age: int\n) -> Person"),
+        report.contains(
+            "def __init__(\n    self: Person,\n    name: str,\n    age: int\n) -> Person"
+        ),
         "Expected constructor hover to show complete signature with -> Person, got: {report}"
     );
 }

pyrefly/lib/test/lsp/lsp_interaction/definition.rs

@@ -489,18 +489,18 @@ fn goto_type_def_on_list_of_primitives_shows_selector() {
 
 #[test]
 fn test_go_to_def_constructor_calls() {
-    // Note: go-to-definition currently goes to the class definition, not __init__.
+    // go-to-definition on constructor calls should go to __init__
     let root = get_test_files_root();
     let constructor_root = root.path().join("constructor_references");
     test_go_to_def(
         constructor_root,
         None,
         "usage.py",
         vec![
-            // Person("Alice", 30) - goes to class Person definition
-            (7, 7, "person.py", 6, 6, 6, 12),
-            // Person("Bob", 25) - goes to class Person definition
-            (8, 7, "person.py", 6, 6, 6, 12),
+            // Person("Alice", 30) - goes to Person.__init__ definition
+            (7, 7, "person.py", 7, 8, 7, 16),
+            // Person("Bob", 25) - goes to Person.__init__ definition
+            (8, 7, "person.py", 7, 8, 7, 16),
         ],
     );
 }