From 79093f3b6dadbf8edfa6dc168fbb723dbbd968b5 Mon Sep 17 00:00:00 2001
From: WilliamK112 <164879897+WilliamK112@users.noreply.github.com>
Date: Sat, 4 Jul 2026 09:04:23 -0500
Subject: [PATCH] Fix TSP definition-site inferred returns

---
 pyrefly/lib/state/lsp.rs                      | 29 +++++---
 .../tsp/tsp_interaction/get_type_queries.rs   | 66 +++++++++++++++++++
 2 files changed, 86 insertions(+), 9 deletions(-)

diff --git a/pyrefly/lib/state/lsp.rs b/pyrefly/lib/state/lsp.rs
index 8dd776cc7b..9b73dc28fd 100644
--- a/pyrefly/lib/state/lsp.rs
+++ b/pyrefly/lib/state/lsp.rs
@@ -1248,15 +1248,26 @@ impl<'a> Transaction<'a> {
     ///
     /// TSP prefers raw bound types of identifiers since it re-resolves declarations.
     pub fn get_computed_type_at_range(&self, handle: &Handle, range: TextRange) -> Option<Type> {
-        // Bare names need declaration-preserving lookup to avoid coercing
-        // overloaded functions into synthesized callables.
-        if range.is_empty()
-            || self.get_ast(handle).is_some_and(|module| {
-                Ast::locate_node(&module, range.start())
-                    .into_iter()
-                    .any(|node| node.range() == range && matches!(node, AnyNodeRef::ExprName(_)))
-            })
-        {
+        // Bare names and definition-name ranges need declaration-preserving
+        // lookup to avoid coercing overloaded functions into synthesized
+        // callables, and to preserve inferred returns for unannotated defs.
+        let is_identifier_name_range = !range.is_empty()
+            && self.identifier_at(handle, range.start()).is_some_and(
+                |IdentifierWithContext {
+                     identifier,
+                     context,
+                 }| {
+                    identifier.range == range
+                        && matches!(
+                            context,
+                            IdentifierContext::Expr(_)
+                                | IdentifierContext::FunctionDef { .. }
+                                | IdentifierContext::MethodDef { .. }
+                        )
+                },
+            );
+
+        if range.is_empty() || is_identifier_name_range {
             return self.get_type_at_preserving_declaration(handle, range.start());
         }
         self.get_type_trace(handle, range)
diff --git a/pyrefly/lib/test/tsp/tsp_interaction/get_type_queries.rs b/pyrefly/lib/test/tsp/tsp_interaction/get_type_queries.rs
index 0d7d268f38..044d7ce5b7 100644
--- a/pyrefly/lib/test/tsp/tsp_interaction/get_type_queries.rs
+++ b/pyrefly/lib/test/tsp/tsp_interaction/get_type_queries.rs
@@ -20,7 +20,26 @@ use crate::test::tsp::tsp_interaction::object_model::write_pyproject;
 fn setup_project(file_content: &str) -> (TspInteraction, String, i32) {
     let temp_dir = TempDir::new().unwrap();
     write_pyproject(temp_dir.path());
+    setup_project_in_dir(temp_dir, file_content)
+}
+
+fn setup_project_with_pyrefly_config(
+    file_content: &str,
+    pyrefly_config: &str,
+) -> (TspInteraction, String, i32) {
+    let temp_dir = TempDir::new().unwrap();
+    write_pyproject(temp_dir.path());
+
+    let pyproject = temp_dir.path().join("pyproject.toml");
+    let mut content = std::fs::read_to_string(&pyproject).unwrap();
+    content.push_str("\n[tool.pyrefly]\n");
+    content.push_str(pyrefly_config);
+    std::fs::write(pyproject, content).unwrap();
+
+    setup_project_in_dir(temp_dir, file_content)
+}
 
+fn setup_project_in_dir(temp_dir: TempDir, file_content: &str) -> (TspInteraction, String, i32) {
     let test_file = temp_dir.path().join("main.py");
     std::fs::write(&test_file, file_content).unwrap();
 
@@ -787,6 +806,53 @@ fn test_get_computed_type_function_has_return_type() {
     tsp.shutdown();
 }
 
+#[test]
+fn test_get_computed_type_definition_site_uses_inferred_return_type() {
+    let code = "\
+def plain():
+    return True
+
+class C:
+    def method(self):
+        return 1
+
+    @property
+    def ok(self):
+        return True
+
+plain
+";
+    let (mut tsp, file_uri, snapshot) = setup_project_with_pyrefly_config(
+        code,
+        "untyped-def-behavior = \"check-and-infer-return-type\"\n",
+    );
+
+    for (start_line, start_character, end_line, end_character, label) in [
+        (0, 4, 0, 9, "plain function definition"),
+        (4, 8, 4, 14, "method definition"),
+        (8, 8, 8, 10, "property getter definition"),
+        (11, 0, 11, 5, "plain function reference"),
+    ] {
+        let result = get_computed_type_range_ok(
+            &mut tsp,
+            &file_uri,
+            start_line,
+            start_character,
+            end_line,
+            end_character,
+            snapshot,
+        );
+        assert_kind(&result, TypeKind::Function);
+
+        let return_type = result
+            .get("returnType")
+            .unwrap_or_else(|| panic!("Expected returnType for {label}: {result}"));
+        assert_kind(return_type, TypeKind::Class);
+    }
+
+    tsp.shutdown();
+}
+
 #[test]
 fn test_get_computed_type_function_has_declaration() {
     // Verify that a def-function has a Regular declaration with name