go-to definition / hover on methods of .pyi

Summary:
D80145015 implemented .py preference on go-to definition, but it did not work for attributes. that's because when we solve a type `a`, `a.attribute` doesn't do another find, it simply goes to the definition of a from solving which is always the .pyi.

asuka implemented a fix for that [here](#1520) for hover docstrings. the intution is to do a second find on attribute definition lookup and look at the AST to find the corresponding definition. this works pretty well, and this diff stack generalizes it to definition.

Reviewed By: connernilsen

Differential Revision: D86814372

fbshipit-source-id: 8b7ed56d74071db55188d448f1eb76f91bcc612d

Author: kinto0

pyrefly/lib/state/lsp.rs

@@ -34,6 +34,7 @@ use pyrefly_python::module::TextRangeWithModule;
 use pyrefly_python::module_name::ModuleName;
 use pyrefly_python::module_path::ModulePath;
 use pyrefly_python::module_path::ModulePathDetails;
+use pyrefly_python::module_path::ModuleStyle;
 use pyrefly_python::short_identifier::ShortIdentifier;
 use pyrefly_python::symbol_kind::SymbolKind;
 use pyrefly_python::sys_info::SysInfo;
@@ -79,6 +80,7 @@ use crate::state::ide::IntermediateDefinition;
 use crate::state::ide::import_regular_import_edit;
 use crate::state::ide::insert_import_edit;
 use crate::state::ide::key_to_intermediate_definition;
+use crate::state::lsp_attributes::AttributeContext;
 use crate::state::require::Require;
 use crate::state::semantic_tokens::SemanticTokenBuilder;
 use crate::state::semantic_tokens::SemanticTokensLegends;
@@ -1154,6 +1156,19 @@ impl<'a> Transaction<'a> {
     ) -> Option<(TextRangeWithModule, Option<TextRange>)> {
         match definition {
             AttrDefinition::FullyResolved(text_range_with_module_info) => {
+                // If prefer_pyi is false and the current module is a .pyi file,
+                // try to find the corresponding .py file
+                if !preference.prefer_pyi
+                    && text_range_with_module_info.module.path().is_interface()
+                    && let Some((exec_module, exec_range)) = self
+                        .search_corresponding_py_module_for_attribute(
+                            handle,
+                            attr_name,
+                            &text_range_with_module_info,
+                        )
+                {
+                    return Some((TextRangeWithModule::new(exec_module, exec_range), None));
+                }
                 Some((text_range_with_module_info, docstring_range))
             }
             AttrDefinition::PartiallyResolvedImportedModuleAttribute { module_name } => {
@@ -1168,6 +1183,39 @@ impl<'a> Transaction<'a> {
         }
     }
 
+    /// Find the .py definition for a corresponding .pyi definition by importing
+    /// and parsing the AST, looking for classes/functions.
+    fn search_corresponding_py_module_for_attribute(
+        &self,
+        request_handle: &Handle,
+        attr_name: &Name,
+        pyi_definition: &TextRangeWithModule,
+    ) -> Option<(Module, TextRange)> {
+        let context = AttributeContext::from_module(&pyi_definition.module, pyi_definition.range)?;
+        let executable_handle = self
+            .import_handle_prefer_executable(request_handle, pyi_definition.module.name(), None)
+            .finding()?;
+        if executable_handle.path().style() != ModuleStyle::Executable {
+            return None;
+        }
+        let _ = self.get_exports(&executable_handle);
+        let executable_module = self.get_module_info(&executable_handle)?;
+        let ast = self.get_ast(&executable_handle).unwrap_or_else(|| {
+            Ast::parse(
+                executable_module.contents(),
+                executable_module.source_type(),
+            )
+            .0
+            .into()
+        });
+        let def_range = crate::state::lsp_attributes::definition_from_executable_ast(
+            ast.as_ref(),
+            &context,
+            attr_name,
+        )?;
+        Some((executable_module, def_range))
+    }
+
     fn key_to_export(
         &self,
         handle: &Handle,

pyrefly/lib/state/lsp_attributes.rs

@@ -0,0 +1,115 @@
+/*
+ * Copyright (c) Meta Platforms, Inc. and affiliates.
+ *
+ * This source code is licensed under the MIT license found in the
+ * LICENSE file in the root directory of this source tree.
+ */
+
+use pyrefly_python::ast::Ast;
+use pyrefly_python::module::Module;
+use ruff_python_ast::ModModule;
+use ruff_python_ast::Stmt;
+use ruff_text_size::Ranged;
+use ruff_text_size::TextRange;
+
+#[derive(Debug, Clone)]
+pub(crate) enum AttributeKind {
+    Function,
+    Class,
+}
+
+#[derive(Debug, Clone)]
+pub(crate) struct AttributeContext {
+    pub(crate) parent_classes: Vec<ruff_python_ast::name::Name>,
+    pub(crate) kind: AttributeKind,
+}
+
+impl AttributeContext {
+    pub(crate) fn from_module(
+        module: &Module,
+        target_range: TextRange,
+    ) -> Option<AttributeContext> {
+        let ast = Ast::parse(module.contents(), module.source_type()).0;
+        let mut parents = Vec::new();
+        Self::from_body(ast.body.as_slice(), &mut parents, target_range)
+    }
+
+    fn from_body(
+        body: &[Stmt],
+        parents: &mut Vec<ruff_python_ast::name::Name>,
+        target_range: TextRange,
+    ) -> Option<AttributeContext> {
+        for stmt in body {
+            match stmt {
+                Stmt::ClassDef(class_def) => {
+                    parents.push(class_def.name.id.clone());
+                    if let Some(ctx) = Self::from_body(&class_def.body, parents, target_range) {
+                        return Some(ctx);
+                    }
+                    parents.pop();
+                    if class_def.range().contains_range(target_range) {
+                        return Some(AttributeContext {
+                            parent_classes: parents.clone(),
+                            kind: AttributeKind::Class,
+                        });
+                    }
+                }
+                Stmt::FunctionDef(_) if stmt.range().contains_range(target_range) => {
+                    return Some(AttributeContext {
+                        parent_classes: parents.clone(),
+                        kind: AttributeKind::Function,
+                    });
+                }
+                _ => {}
+            }
+        }
+        None
+    }
+}
+
+pub(crate) fn definition_from_executable_ast(
+    ast: &ModModule,
+    context: &AttributeContext,
+    attr_name: &ruff_python_ast::name::Name,
+) -> Option<TextRange> {
+    let mut body = ast.body.as_slice();
+    for class_name in &context.parent_classes {
+        let class_def = body.iter().find_map(|stmt| match stmt {
+            Stmt::ClassDef(class_def) if &class_def.name.id == class_name => Some(class_def),
+            _ => None,
+        })?;
+        body = class_def.body.as_slice();
+    }
+    match context.kind {
+        AttributeKind::Function => definition_from_function(body, attr_name),
+        AttributeKind::Class => definition_from_class(body, attr_name),
+    }
+}
+
+fn definition_from_function(
+    body: &[Stmt],
+    attr_name: &ruff_python_ast::name::Name,
+) -> Option<TextRange> {
+    for stmt in body {
+        if let Stmt::FunctionDef(func_def) = stmt
+            && &func_def.name.id == attr_name
+        {
+            return Some(func_def.name.range);
+        }
+    }
+    None
+}
+
+fn definition_from_class(
+    body: &[Stmt],
+    attr_name: &ruff_python_ast::name::Name,
+) -> Option<TextRange> {
+    for stmt in body {
+        if let Stmt::ClassDef(class_def) = stmt
+            && &class_def.name.id == attr_name
+        {
+            return Some(class_def.name.range);
+        }
+    }
+    None
+}

pyrefly/lib/state/mod.rs

@@ -12,6 +12,7 @@ pub mod ide;
 pub mod load;
 pub mod loader;
 pub mod lsp;
+pub mod lsp_attributes;
 pub mod memory;
 pub mod notebook;
 pub mod require;

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

@@ -170,9 +170,8 @@ fn definition_in_builtins() {
     );
 }
 
-// todo(kylei): definition of an attribute of a pyi should still point to py
 #[test]
-fn definition_on_attr_of_pyi() {
+fn definition_on_attr_of_pyi_goes_to_py() {
     let root = get_test_files_root();
     let mut interaction = LspInteraction::new();
     interaction.set_root(root.path().to_path_buf());
@@ -184,7 +183,7 @@ fn definition_on_attr_of_pyi() {
     interaction.server.definition(file, 7, 4);
     interaction
         .client
-        .expect_definition_response_from_root("attributes_of_py/lib.pyi", 6, 8, 6, 9);
+        .expect_definition_response_from_root("attributes_of_py/lib.py", 7, 8, 7, 9);
     interaction.shutdown();
 }