Hover type with chosen overload

Summary: In this diff, I setup the infra to record the chosen overload, and use it in hover. When we hover on callees, we will show the chosen overload instead. In the next diff, we will use the same infra to power signature help.

Reviewed By: kinto0

Differential Revision: D76222224

fbshipit-source-id: 9ff8578bb387ad4aec2b4292150882b2d39cbd13

Author: SamChou19815

pyrefly/lib/alt/answers.rs

@@ -61,6 +61,7 @@ use crate::table;
 use crate::table_for_each;
 use crate::table_mut_for_each;
 use crate::table_try_for_each;
+use crate::types::callable::Callable;
 use crate::types::class::Class;
 use crate::types::equality::TypeEq;
 use crate::types::equality::TypeEqCtx;
@@ -81,9 +82,18 @@ pub struct Index {
     pub externally_defined_attribute_references: SmallMap<ModulePath, Vec<(TextRange, TextRange)>>,
 }
 
+#[derive(Debug)]
+struct OverloadedCallee {
+    all_overloads: Vec<Callable>,
+    closest_overload: Callable,
+    is_closest_overload_chosen: bool,
+}
+
 #[derive(Debug, Default)]
 pub struct Traces {
     types: SmallMap<TextRange, Arc<Type>>,
+    /// A map from (range of callee, overload information)
+    overloaded_callees: SmallMap<TextRange, OverloadedCallee>,
 }
 
 /// Invariants:
@@ -535,6 +545,42 @@ impl Answers {
         let lock = self.trace.as_ref()?.lock();
         lock.types.get(&range).duped()
     }
+
+    pub fn get_chosen_overload_trace(&self, range: TextRange) -> Option<Callable> {
+        let lock = self.trace.as_ref()?.lock();
+        let overloaded_callee = lock.overloaded_callees.get(&range)?;
+        if overloaded_callee.is_closest_overload_chosen {
+            Some(overloaded_callee.closest_overload.clone())
+        } else {
+            None
+        }
+    }
+
+    /// Returns all the overload, and the index of a chosen one
+    #[allow(dead_code)]
+    pub fn get_all_overload_trace(
+        &self,
+        range: TextRange,
+    ) -> Option<(Vec<Callable>, Option<usize>)> {
+        let lock = self.trace.as_ref()?.lock();
+        let overloaded_callee = lock.overloaded_callees.get(&range)?;
+        let chosen_overload_index =
+            overloaded_callee
+                .all_overloads
+                .iter()
+                .enumerate()
+                .find_map(|(index, signature)| {
+                    if signature == &overloaded_callee.closest_overload {
+                        Some(index)
+                    } else {
+                        None
+                    }
+                });
+        Some((
+            overloaded_callee.all_overloads.clone(),
+            chosen_overload_index,
+        ))
+    }
 }
 
 impl<'a, Ans: LookupAnswer> AnswersSolver<'a, Ans> {
@@ -699,6 +745,27 @@ impl<'a, Ans: LookupAnswer> AnswersSolver<'a, Ans> {
         }
     }
 
+    /// Record all the overloads and the chosen overload.
+    /// The trace will be used to power signature help and hover for overloaded functions.
+    pub fn record_overload_trace(
+        &self,
+        loc: TextRange,
+        all_overloads: &[Callable],
+        closest_overload: &Callable,
+        is_closest_overload_chosen: bool,
+    ) {
+        if let Some(trace) = &self.current.trace {
+            trace.lock().overloaded_callees.insert(
+                loc,
+                OverloadedCallee {
+                    all_overloads: all_overloads.to_vec(),
+                    closest_overload: closest_overload.clone(),
+                    is_closest_overload_chosen,
+                },
+            );
+        }
+    }
+
     /// Check if `want` matches `got` returning `want` if the check fails.
     pub fn check_and_return_type_info(
         &self,

pyrefly/lib/alt/call.rs

@@ -670,6 +670,8 @@ impl<'a, Ans: LookupAnswer> AnswersSolver<'a, Ans> {
                 None,
             );
             if arg_errors.is_empty() && call_errors.is_empty() {
+                // An overload is chosen, we should record it to power IDE services.
+                self.record_overload_trace(range, overloads.as_slice(), callable, true);
                 // It's only safe to return immediately if both arg_errors and call_errors are
                 // empty, as parameter types from the overload signature may be used as hints when
                 // evaluating arguments, producing arg_errors for some overloads but not others.
@@ -692,6 +694,12 @@ impl<'a, Ans: LookupAnswer> AnswersSolver<'a, Ans> {
         }
         // We're guaranteed to have at least one overload.
         let closest_overload = closest_overload.unwrap();
+        self.record_overload_trace(
+            range,
+            overloads.as_slice(),
+            &closest_overload.signature,
+            false,
+        );
         errors.extend(closest_overload.arg_errors);
         if closest_overload.call_errors.is_empty() {
             // No overload evaluated completely successfully, but we still say we found a match if

pyrefly/lib/state/lsp.rs

@@ -5,6 +5,8 @@
  * LICENSE file in the root directory of this source tree.
  */
 
+use std::sync::Arc;
+
 use dupe::Dupe;
 use itertools::Itertools;
 use lsp_types::CompletionItem;
@@ -17,7 +19,9 @@ use pyrefly_util::visit::Visit;
 use ruff_python_ast::AnyNodeRef;
 use ruff_python_ast::Expr;
 use ruff_python_ast::ExprAttribute;
+use ruff_python_ast::ExprCall;
 use ruff_python_ast::Identifier;
+use ruff_python_ast::ModModule;
 use ruff_python_ast::Stmt;
 use ruff_python_ast::name::Name;
 use ruff_text_size::Ranged;
@@ -242,13 +246,45 @@ impl<'a> Transaction<'a> {
         if let Some(key) = self.definition_at(handle, position) {
             return self.get_type(handle, &key);
         }
+        fn callee_at(mod_module: Arc<ModModule>, position: TextSize) -> Option<ExprCall> {
+            fn f(x: &Expr, find: TextSize, res: &mut Option<ExprCall>) {
+                if let Expr::Call(call) = x
+                    && call.func.range().contains_inclusive(find)
+                {
+                    f(call.func.as_ref(), find, res);
+                    if res.is_some() {
+                        return;
+                    }
+                    *res = Some(call.clone());
+                } else {
+                    x.recurse(&mut |x| f(x, find, res));
+                }
+            }
+            let mut res = None;
+            mod_module.visit(&mut |x| f(x, position, &mut res));
+            res
+        }
+        let callee = callee_at(self.get_ast(handle)?, position);
         match self.identifier_at(handle, position) {
             Some(IdentifierWithContext {
                 identifier: id,
                 context: IdentifierContext::Expr,
             }) => {
                 if self.get_bindings(handle)?.is_valid_usage(&id) {
-                    return self.get_type(handle, &Key::BoundName(ShortIdentifier::new(&id)));
+                    if let Some(ExprCall {
+                        range: _,
+                        func,
+                        arguments,
+                    }) = &callee
+                        && func.range() == id.range
+                        && let Some(chosen_overload) = self
+                            .get_answers(handle)
+                            .and_then(|answers| answers.get_chosen_overload_trace(arguments.range))
+                    {
+                        return Some(Type::Callable(Box::new(chosen_overload)));
+                    } else {
+                        return self.get_type(handle, &Key::BoundName(ShortIdentifier::new(&id)));
+                    }
                 } else {
                     return None;
                 }
@@ -269,7 +305,20 @@ impl<'a> Transaction<'a> {
             None => {}
         }
         let attribute = self.attribute_at(handle, position)?;
-        self.get_type_trace(handle, attribute.range)
+        if let Some(ExprCall {
+            range: _,
+            func,
+            arguments,
+        }) = &callee
+            && func.range() == attribute.range
+            && let Some(chosen_overload) = self
+                .get_answers(handle)
+                .and_then(|answers| answers.get_chosen_overload_trace(arguments.range))
+        {
+            Some(Type::Callable(Box::new(chosen_overload)))
+        } else {
+            self.get_type_trace(handle, attribute.range)
+        }
     }
 
     fn resolve_named_import(

pyrefly/lib/test/lsp/hover_type.rs

@@ -11,6 +11,7 @@ use ruff_text_size::TextSize;
 use crate::state::handle::Handle;
 use crate::state::state::State;
 use crate::test::util::get_batched_lsp_operations_report;
+use crate::test::util::get_batched_lsp_operations_report_allow_error;
 
 fn get_test_report(state: &State, handle: &Handle, position: TextSize) -> String {
     if let Some(t) = state.transaction().get_type_at(handle, position) {
@@ -282,3 +283,86 @@ Hover Result: `Literal[5] | int`
         report.trim(),
     );
 }
+
+#[test]
+fn overloaded_functions_test() {
+    let code = r#"
+from typing import overload
+
+@overload
+def overloaded_func(a: str) -> bool: ...
+@overload
+def overloaded_func(a: int, b: bool) -> str: ...
+def overloaded_func():
+    pass
+
+overloaded_func("")
+# ^
+overloaded_func(1, True)
+# ^
+overloaded_func(False)
+# ^
+
+"#;
+    let report = get_batched_lsp_operations_report_allow_error(&[("main", code)], get_test_report);
+    assert_eq!(
+        r#"
+# main.py
+11 | overloaded_func("")
+       ^
+Hover Result: `(a: str) -> bool`
+
+13 | overloaded_func(1, True)
+       ^
+Hover Result: `(a: int, b: bool) -> str`
+
+15 | overloaded_func(False)
+       ^
+Hover Result: `Overload[(a: str) -> bool, (a: int, b: bool) -> str]`
+"#
+        .trim(),
+        report.trim(),
+    );
+}
+
+#[test]
+fn overloaded_methods_test() {
+    let code = r#"
+from typing import overload
+
+class Foo:
+    @overload
+    def overloaded_meth(self, a: str) -> bool: ...
+    @overload
+    def overloaded_meth(self, a: int, b: bool) -> str: ...
+    def overloaded_meth(self):
+        pass
+
+foo = Foo()
+foo.overloaded_meth("")
+#       ^
+foo.overloaded_meth(1, True)
+#       ^
+foo.overloaded_meth(False)
+#       ^
+"#;
+    let report = get_batched_lsp_operations_report_allow_error(&[("main", code)], get_test_report);
+    assert_eq!(
+        r#"
+# main.py
+13 | foo.overloaded_meth("")
+             ^
+Hover Result: `(self: Foo, a: str) -> bool`
+
+15 | foo.overloaded_meth(1, True)
+             ^
+Hover Result: `(self: Foo, a: int, b: bool) -> str`
+
+17 | foo.overloaded_meth(False)
+             ^
+Hover Result: `BoundMethod[Foo, Overload[(self: Self@Foo, a: str) -> bool, (self: Self@Foo, a: int, b: bool) -> str]]`
+"#
+        .trim(),
+        report.trim(),
+    );
+}