fix

Author: asukaminato0721

pyrefly/lib/lsp/non_wasm/server.rs

@@ -197,6 +197,8 @@ use pyrefly_build::handle::Handle;
 use pyrefly_config::config::ConfigSource;
 use pyrefly_config::error_kind::Severity;
 use pyrefly_python::PYTHON_EXTENSIONS;
+use pyrefly_python::ast::Ast;
+use pyrefly_python::module::Module;
 use pyrefly_python::module::TextRangeWithModule;
 use pyrefly_python::module_name::ModuleName;
 use pyrefly_python::module_name::ModuleNameWithKind;
@@ -314,9 +316,12 @@ use crate::lsp::wasm::provide_type::ProvideTypeResponse;
 use crate::lsp::wasm::provide_type::provide_type;
 use crate::module::bundled::BundledStub;
 use crate::state::load::LspFile;
+use crate::state::loader::LoaderFindCache;
 use crate::state::lsp::DisplayTypeErrors;
 use crate::state::lsp::FindDefinitionItemWithDocstring;
 use crate::state::lsp::FindPreference;
+use crate::state::lsp::IdentifierContext;
+use crate::state::lsp::IdentifierWithContext;
 use crate::state::lsp::ImportBehavior;
 use crate::state::lsp::LocalRefactorCodeAction;
 use crate::state::notebook::LspNotebook;
@@ -1825,6 +1830,28 @@ impl Server {
                     return Ok(ProcessEvent::Continue);
                 }
 
+                if let Some(params) = as_request::<GotoDefinition>(&x) {
+                    if let Some(params) = self
+                        .extract_request_params_or_send_err_response::<GotoDefinition>(
+                            params, &x.id,
+                        )
+                    {
+                        if let Some(response) = self.goto_definition_import_fast_path(&params) {
+                            let response = match response {
+                                Ok(response) => response,
+                                Err(reason) => {
+                                    telemetry_event.set_empty_response_reason(reason);
+                                    None
+                                }
+                            };
+                            self.send_response(new_response(x.id, Ok(response)));
+                            return Ok(ProcessEvent::Continue);
+                        }
+                    } else {
+                        return Ok(ProcessEvent::Continue);
+                    }
+                }
+
                 let mut transaction =
                     ide_transaction_manager.non_committable_transaction(&self.state);
 
@@ -3933,6 +3960,84 @@ impl Server {
             .map(|(handle, _)| handle)
     }
 
+    fn module_for_uri_fast_path(&self, uri: &Url, handle: &Handle) -> Option<Module> {
+        let path = self.path_for_uri_or_notebook_cell(uri)?;
+        if let Some(file) = self.open_files.read().get(&path).duped() {
+            return Some(match &*file {
+                LspFile::Source(contents) => {
+                    Module::new(handle.module(), handle.path().dupe(), Arc::clone(contents))
+                }
+                LspFile::Notebook(notebook) => Module::new_notebook(
+                    handle.module(),
+                    handle.path().dupe(),
+                    Arc::clone(notebook.ruff_notebook()),
+                ),
+            });
+        }
+        let contents = std::fs::read_to_string(&path).ok()?;
+        Some(Module::new(
+            handle.module(),
+            handle.path().dupe(),
+            Arc::new(contents),
+        ))
+    }
+
+    fn goto_definition_import_fast_path(
+        &self,
+        params: &GotoDefinitionParams,
+    ) -> Option<Result<Option<GotoDefinitionResponse>, EmptyResponseReason>> {
+        let uri = &params.text_document_position_params.text_document.uri;
+        let handle = match self.make_handle_if_enabled(uri, Some(GotoDefinition::METHOD)) {
+            Ok(handle) => handle,
+            Err(err) => return Some(Err(err.into())),
+        };
+        let module = match self.module_for_uri_fast_path(uri, &handle) {
+            Some(module) => module,
+            None => return Some(Err(EmptyResponseReason::ModuleInfoNotFound)),
+        };
+        let position = module.from_lsp_position(
+            params.text_document_position_params.position,
+            self.maybe_get_cell_index(uri),
+        );
+        let ast = Ast::parse(module.contents(), module.source_type()).0;
+        let covering_nodes = Ast::locate_node(&ast, position);
+        let Some(IdentifierWithContext {
+            identifier,
+            context: IdentifierContext::ImportedModule { name, dots },
+        }) = Transaction::identifier_from_covering_nodes(&covering_nodes)
+        else {
+            return None;
+        };
+        let target_module =
+            Transaction::target_imported_module_name(&handle, &identifier, name, dots, position);
+        let config = self
+            .state
+            .config_finder()
+            .python_file(handle.module_kind(), handle.path());
+        let loader = LoaderFindCache::new(config.dupe());
+        let Some(target_path) = loader
+            .find_import_prefer_executable(target_module, Some(handle.path()))
+            .finding()
+        else {
+            return Some(Err(EmptyResponseReason::ModuleNotFound));
+        };
+        let Some(path) = to_real_path(&target_path) else {
+            return Some(Err(EmptyResponseReason::ModuleInfoNotFound));
+        };
+        let path = match &self.path_remapper {
+            Some(remapper) => remapper(&path).into_owned(),
+            None => path,
+        };
+        let uri = match Url::from_file_path(path.absolutize()) {
+            Ok(uri) => uri,
+            Err(_) => return Some(Err(EmptyResponseReason::NoFilePath)),
+        };
+        Some(Ok(Some(GotoDefinitionResponse::Scalar(Location {
+            uri,
+            range: Range::default(),
+        }))))
+    }
+
     fn goto_definition(
         &self,
         transaction: &Transaction<'_>,

pyrefly/lib/state/lsp.rs

@@ -682,7 +682,7 @@ impl<'a> Transaction<'a> {
         Self::identifier_from_covering_nodes(&covering_nodes)
     }
 
-    fn identifier_from_covering_nodes(
+    pub(crate) fn identifier_from_covering_nodes(
         covering_nodes: &[AnyNodeRef],
     ) -> Option<IdentifierWithContext> {
         match (
@@ -825,6 +825,43 @@ impl<'a> Transaction<'a> {
         }
     }
 
+    pub(crate) fn target_imported_module_name(
+        handle: &Handle,
+        identifier: &Identifier,
+        module_name: ModuleName,
+        dots: u32,
+        position: TextSize,
+    ) -> ModuleName {
+        let resolved_module_name = if dots > 0 {
+            let is_init = handle.path().is_init();
+            let suffix = if module_name.as_str().is_empty() {
+                None
+            } else {
+                Some(&Name::new(module_name.as_str()))
+            };
+            handle
+                .module()
+                .new_maybe_relative(is_init, dots, suffix)
+                .unwrap_or(module_name)
+        } else {
+            module_name
+        };
+
+        let components = resolved_module_name.components();
+        if let Some(idx) = components.iter().position(|c| c == &identifier.id) {
+            ModuleName::from_parts(&components[..=idx])
+        } else if identifier.as_str() == resolved_module_name.as_str() {
+            let module_str = resolved_module_name.as_str();
+            let offset = (position - identifier.range.start())
+                .to_usize()
+                .min(module_str.len());
+            let idx = module_str[..offset].matches('.').count();
+            ModuleName::from_parts(&components[..=idx])
+        } else {
+            resolved_module_name
+        }
+    }
+
     fn callee_at(&self, handle: &Handle, position: TextSize) -> Option<ExprCall> {
         let mod_module = self.get_ast(handle)?;
         fn f(x: &Expr, find: TextSize, res: &mut Option<ExprCall>) {
@@ -1929,42 +1966,13 @@ impl<'a> Transaction<'a> {
                         dots,
                     },
             }) => {
-                // For relative imports (dots > 0), resolve the module name using
-                // the current file's module name as context.
-                let resolved_module_name = if dots > 0 {
-                    let is_init = handle.path().is_init();
-                    let suffix = if module_name.as_str().is_empty() {
-                        None
-                    } else {
-                        Some(&Name::new(module_name.as_str()))
-                    };
-                    handle
-                        .module()
-                        .new_maybe_relative(is_init, dots, suffix)
-                        .unwrap_or(module_name)
-                } else {
-                    module_name
-                };
-
-                // Build the module name for lookup based on identifier position.
-                let components = resolved_module_name.components();
-
-                let target_module_name =
-                    if let Some(idx) = components.iter().position(|c| c == &identifier.id) {
-                        // Identifier matches a module component.
-                        ModuleName::from_parts(&components[..=idx])
-                    } else if identifier.as_str() == resolved_module_name.as_str() {
-                        // Identifier matches full module name; decide which component based on position offset.
-                        let module_str = resolved_module_name.as_str();
-                        let offset = (position - identifier.range.start())
-                            .to_usize()
-                            .min(module_str.len());
-                        let idx = module_str[..offset].matches('.').count();
-                        ModuleName::from_parts(&components[..=idx])
-                    } else {
-                        // Fallback: use the whole module name.
-                        resolved_module_name
-                    };
+                let target_module_name = Self::target_imported_module_name(
+                    handle,
+                    &identifier,
+                    module_name,
+                    dots,
+                    position,
+                );
                 match self.find_definition_for_imported_module(
                     handle,
                     target_module_name,

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

@@ -5,6 +5,7 @@
  * LICENSE file in the root directory of this source tree.
  */
 
+use std::fs;
 use std::path::PathBuf;
 use std::sync::Arc;
 
@@ -199,6 +200,55 @@ fn test_go_to_def_relative_path_helper() {
     );
 }
 
+#[test]
+fn definition_on_import_module_while_recheck_is_blocked() {
+    let root = get_test_files_root();
+    let root_path = root.path().join("basic");
+    let scope_uri = Url::from_file_path(&root_path).unwrap();
+    let bar_path = root_path.join("bar.py");
+    let bar_contents = fs::read_to_string(&bar_path).unwrap();
+    let mut interaction = LspInteraction::new_with_indexing_mode(IndexingMode::LazyBlocking);
+    interaction.set_root(root_path.clone());
+    interaction
+        .initialize(InitializeSettings {
+            configuration: Some(Some(
+                json!([{"pyrefly": {"displayTypeErrors": "force-on"}}]),
+            )),
+            workspace_folders: Some(vec![("test".to_owned(), scope_uri)]),
+            file_watch: true,
+            ..Default::default()
+        })
+        .unwrap();
+    interaction.client.did_open("bar.py");
+    interaction
+        .client
+        .expect_publish_diagnostics_eventual_error_count(bar_path.clone(), 0)
+        .expect("Failed to receive initial diagnostics for bar");
+    interaction.client.did_open("foo.py");
+    interaction
+        .client
+        .expect_publish_diagnostics_eventual_error_count(root_path.join("foo.py"), 0)
+        .expect("Failed to receive initial diagnostics for foo");
+
+    interaction.do_not_commit_next_recheck();
+    interaction
+        .client
+        .edit_file("bar.py", &bar_contents.replace("foo = 3", "foo = 4"));
+    interaction
+        .client
+        .expect_publish_diagnostics_eventual_error_count(bar_path, 0)
+        .expect("Failed to receive diagnostics while the recheck was blocked");
+
+    interaction
+        .client
+        .definition("foo.py", 5, 7)
+        .expect_definition_response_from_root("bar.py", 0, 0, 0, 0)
+        .unwrap();
+
+    interaction.continue_recheck();
+    interaction.shutdown().unwrap();
+}
+
 #[test]
 fn definition_in_builtins() {
     let root = get_test_files_root();