From d0d8f6aae7185113ee00dd10c060fc40084cb217 Mon Sep 17 00:00:00 2001 From: Asuka Minato Date: Wed, 28 Oct 2026 17:34:51 +0900 Subject: [PATCH 1/2] test --- pyrefly/lib/test/lsp/code_actions.rs | 124 +++++++++++++++++++++++++++ 1 file changed, 124 insertions(+) diff --git a/pyrefly/lib/test/lsp/code_actions.rs b/pyrefly/lib/test/lsp/code_actions.rs index 2b83c8698a..3f00528f55 100644 --- a/pyrefly/lib/test/lsp/code_actions.rs +++ b/pyrefly/lib/test/lsp/code_actions.rs @@ -12,6 +12,7 @@ use pyrefly_build::handle::Handle; use pyrefly_python::ast::Ast; use pyrefly_python::module::Module; use ruff_python_ast::PySourceType; +use ruff_text_size::Ranged; use ruff_text_size::TextRange; use ruff_text_size::TextSize; @@ -29,6 +30,7 @@ use crate::test::util::extract_cursors_for_test; use crate::test::util::get_batched_lsp_operations_report_allow_error; use crate::test::util::mk_multi_file_state; use crate::test::util::mk_multi_file_state_assert_no_errors; +use crate::test::util::mk_multi_file_state_with_env; fn apply_patch(info: &ModuleInfo, range: TextRange, patch: String) -> (String, String) { let before = info.contents().as_str().to_owned(); @@ -1052,6 +1054,128 @@ my_module ); } +#[test] +fn insertion_test_implicit_import() { + for (code, after) in [ + ( + "import foo\nx = foo.bar.x\n", + Some("import foo.bar\nx = foo.bar.x\n"), + ), + ("import foo as f\nx = f.bar.x\n", None), + ("import foo.bar as foo\nx = foo.baz.x\n", None), + ( + "import foo.bar\nx = foo.bar.baz.x\n", + Some("import foo.bar.baz\nx = foo.bar.baz.x\n"), + ), + ("from foo import bar\nx = bar.baz.x\n", None), + ( + "import foo\r\nx = foo.bar.x\r\n", + Some("import foo.bar\r\nx = foo.bar.x\r\n"), + ), + ( + "def f():\n import foo\n return foo.bar.x\n", + Some("def f():\n import foo.bar\n return foo.bar.x\n"), + ), + ( + "import foo\ny = foo\ndef f():\n import foo\n return foo.bar.x\n", + Some("import foo\ny = foo\ndef f():\n import foo.bar\n return foo.bar.x\n"), + ), + ( + "def f(): import foo; return foo.bar.x\n", + Some("def f(): import foo.bar; return foo.bar.x\n"), + ), + ("def f():\n import foo as f\n return f.bar.x\n", None), + ("import str as s\ny = str(1)\nx = s.bar.x\n", None), + ( + "def f(flag: bool):\n if flag:\n import foo\n else:\n import foo\n return foo.bar.x\n", + None, + ), + ( + "def f(items: list[int]):\n import foo\n for _ in items:\n import foo\n return foo.bar.x\n", + None, + ), + ( + "def f(items: list[int]):\n import foo\n for _ in items:\n print(foo.bar.x)\n import foo\n", + None, + ), + ] { + let env = || { + let mut env = TestEnv::new(); + env.add_with_path("foo", "foo/__init__.py", ""); + env.add_with_path("foo.bar", "foo/bar/__init__.py", "x = 1\n"); + env.add_with_path("foo.bar.baz", "foo/bar/baz.py", "x = 2\n"); + env.add_with_path("str", "str/__init__.py", ""); + env.add_with_path("str.bar", "str/bar.py", "x = 1\n"); + env + }; + let (handles, state) = + mk_multi_file_state_with_env(env(), &[("main", code)], Require::Exports, false); + let handle = &handles["main"]; + let transaction = state.transaction(); + let errors = transaction + .get_errors(vec![handle]) + .collect_errors() + .ordinary; + assert_eq!(errors.len(), 1); + assert_eq!(errors[0].error_kind(), ErrorKind::ImplicitImport); + let position = errors[0].range().start(); + let actions = transaction + .local_quickfix_code_actions_sorted( + handle, + TextRange::empty(position), + ImportFormat::Absolute, + None, + ) + .unwrap_or_default(); + let Some(after) = after else { + assert!( + !actions + .iter() + .any(|(title, _)| title.starts_with("Insert import:")) + ); + assert!( + transaction + .get_answers(handle) + .unwrap() + .bindings() + .unused_imports() + .is_empty() + ); + continue; + }; + let (_, edits) = actions + .iter() + .find(|(title, _)| title.starts_with("Insert import:")) + .unwrap_or_else(|| panic!("expected import action for {code}")); + assert_eq!(apply_refactor_edits_for_module(&edits[0].0, edits), after); + let (fixed_handles, fixed_state) = + mk_multi_file_state_with_env(env(), &[("main", after)], Require::Exports, true); + assert!( + fixed_state + .transaction() + .get_answers(&fixed_handles["main"]) + .unwrap() + .bindings() + .unused_imports() + .is_empty() + ); + + let unrelated_actions = transaction + .local_quickfix_code_actions_sorted( + handle, + TextRange::empty(TextSize::new(0)), + ImportFormat::Absolute, + None, + ) + .unwrap_or_default(); + assert!( + !unrelated_actions + .iter() + .any(|(title, _)| title.starts_with("Insert import:")) + ); + } +} + #[test] fn insertion_test_common_alias_module_import() { let code = r#" From be600e7223cfad1792387dbe0284b605de7e9477 Mon Sep 17 00:00:00 2001 From: Asuka Minato Date: Tue, 22 Sep 2026 01:31:14 +0900 Subject: [PATCH 2/2] fix --- pyrefly/lib/alt/attr.rs | 9 +- pyrefly/lib/error/error.rs | 2 + pyrefly/lib/error/suppress.rs | 4 +- pyrefly/lib/state/lsp.rs | 25 ++++- pyrefly/lib/state/lsp/quick_fixes.rs | 1 + .../lib/state/lsp/quick_fixes/enum_member.rs | 4 +- .../state/lsp/quick_fixes/implicit_import.rs | 94 +++++++++++++++++++ .../state/lsp/quick_fixes/unused_ignore.rs | 4 +- 8 files changed, 134 insertions(+), 9 deletions(-) create mode 100644 pyrefly/lib/state/lsp/quick_fixes/implicit_import.rs diff --git a/pyrefly/lib/alt/attr.rs b/pyrefly/lib/alt/attr.rs index 4bce1c2f73..5a4c6782e2 100644 --- a/pyrefly/lib/alt/attr.rs +++ b/pyrefly/lib/alt/attr.rs @@ -1652,13 +1652,14 @@ impl<'ctx, 'answer, Ans: LookupAnswer> AnswersSolver<'ctx, 'answer, Ans> { } Attribute::Simple(ty) => Ok(ty), Attribute::ModuleFallback(_, name, ty) => { - self.error_with_context( - errors, + errors.error_builder( range, ErrorKind::ImplicitImport, format!("Module `{name}` exists, but was not imported explicitly. You are relying on other modules to load it."), - context, - ); + ) + .with_context(context) + .with_quick_fix(ErrorQuickFix::ImportModule { module: name }) + .emit(); Ok(ty) } Attribute::GetAttr(_, getattr_attr, name) => self diff --git a/pyrefly/lib/error/error.rs b/pyrefly/lib/error/error.rs index 886b496645..845f42df23 100644 --- a/pyrefly/lib/error/error.rs +++ b/pyrefly/lib/error/error.rs @@ -22,6 +22,7 @@ use pyrefly_python::ignore::SuppressionEffect; use pyrefly_python::ignore::Tool; use pyrefly_python::ignore::TypeIgnoreUnknownTagBehavior; use pyrefly_python::module::Module; +use pyrefly_python::module_name::ModuleName; use pyrefly_python::module_path::ModulePath; use pyrefly_util::display::number_thousands; use pyrefly_util::lined_buffer::DisplayRange; @@ -58,6 +59,7 @@ pub enum ErrorQuickFix { ReplaceWithEnumMember { replacement: String }, AssertNotNone, RemoveUnusedSuppression(SuppressionEdit), + ImportModule { module: ModuleName }, } impl From for ErrorQuickFix { diff --git a/pyrefly/lib/error/suppress.rs b/pyrefly/lib/error/suppress.rs index e1bec1426f..bd0ee28b0e 100644 --- a/pyrefly/lib/error/suppress.rs +++ b/pyrefly/lib/error/suppress.rs @@ -139,7 +139,9 @@ impl SerializedError { if let ModulePathDetails::FileSystem(path) = error.path().details() { let removal_edit = error.quick_fixes().iter().find_map(|fix| match fix { ErrorQuickFix::RemoveUnusedSuppression(edit) => Some(edit), - ErrorQuickFix::ReplaceWithEnumMember { .. } | ErrorQuickFix::AssertNotNone => None, + ErrorQuickFix::ReplaceWithEnumMember { .. } + | ErrorQuickFix::AssertNotNone + | ErrorQuickFix::ImportModule { .. } => None, }); let line = removal_edit .map(|edit| { diff --git a/pyrefly/lib/state/lsp.rs b/pyrefly/lib/state/lsp.rs index 47531097a6..75506ab2ff 100644 --- a/pyrefly/lib/state/lsp.rs +++ b/pyrefly/lib/state/lsp.rs @@ -82,6 +82,7 @@ use crate::binding::binding::Binding; use crate::binding::binding::Key; use crate::binding::binding::KeyAnnotation; use crate::config::error_kind::ErrorKind; +use crate::error::error::ErrorQuickFix; use crate::error::suppress::detect_line_ending; use crate::export::exports::Export; use crate::export::exports::ExportLocation; @@ -3621,6 +3622,26 @@ impl<'a> Transaction<'a> { } for error in errors { let error_range = error.range(); + if error_range.contains_range(range) { + for fix in error.quick_fixes() { + if let ErrorQuickFix::ImportModule { module } = fix + && let Some(action) = + quick_fixes::implicit_import::import_module_code_action( + self, + handle, + &module_info, + &ast, + error_range, + *module, + ) + { + let key = (action.0.clone(), action.2, action.3.clone()); + if other_action_keys.insert(key) { + other_actions.push(action); + } + } + } + } if error_range.contains_range(range) && let Some(action) = quick_fixes::enum_member::replace_with_enum_member_code_action( &module_info, @@ -3733,7 +3754,7 @@ impl<'a> Transaction<'a> { is_private_import, }); } - self.create_quickfix_action_for_fuzzy_match( + self.create_quickfix_action_for_module_import( handle, &module_info, &ast, @@ -3906,7 +3927,7 @@ impl<'a> Transaction<'a> { Some(module_name) } - fn create_quickfix_action_for_fuzzy_match( + fn create_quickfix_action_for_module_import( &self, handle: &Handle, module_info: &Module, diff --git a/pyrefly/lib/state/lsp/quick_fixes.rs b/pyrefly/lib/state/lsp/quick_fixes.rs index 2aeed6d263..f92818f4ec 100644 --- a/pyrefly/lib/state/lsp/quick_fixes.rs +++ b/pyrefly/lib/state/lsp/quick_fixes.rs @@ -17,6 +17,7 @@ mod extract_shared; pub(crate) mod extract_superclass; pub(crate) mod extract_variable; pub(crate) mod generate_code; +pub(crate) mod implicit_import; pub(crate) mod inline_method; pub(crate) mod inline_parameter; pub(crate) mod inline_variable; diff --git a/pyrefly/lib/state/lsp/quick_fixes/enum_member.rs b/pyrefly/lib/state/lsp/quick_fixes/enum_member.rs index 9da8f37d98..b17d7be526 100644 --- a/pyrefly/lib/state/lsp/quick_fixes/enum_member.rs +++ b/pyrefly/lib/state/lsp/quick_fixes/enum_member.rs @@ -35,7 +35,9 @@ pub(crate) fn replace_with_enum_member_code_action( fn enum_member_replacement(error: &Error) -> Option<&str> { error.quick_fixes().iter().find_map(|fix| match fix { ErrorQuickFix::ReplaceWithEnumMember { replacement } => Some(replacement.as_str()), - ErrorQuickFix::AssertNotNone | ErrorQuickFix::RemoveUnusedSuppression(_) => None, + ErrorQuickFix::AssertNotNone + | ErrorQuickFix::RemoveUnusedSuppression(_) + | ErrorQuickFix::ImportModule { .. } => None, }) } diff --git a/pyrefly/lib/state/lsp/quick_fixes/implicit_import.rs b/pyrefly/lib/state/lsp/quick_fixes/implicit_import.rs new file mode 100644 index 0000000000..8728c6bfdc --- /dev/null +++ b/pyrefly/lib/state/lsp/quick_fixes/implicit_import.rs @@ -0,0 +1,94 @@ +/* + * 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 dupe::Dupe; +use pyrefly_build::handle::Handle; +use pyrefly_python::ast::Ast; +use pyrefly_python::module::Module; +use pyrefly_python::module_name::ModuleName; +use pyrefly_python::short_identifier::ShortIdentifier; +use pyrefly_util::gas::Gas; +use ruff_python_ast::AnyNodeRef; +use ruff_python_ast::ModModule; +use ruff_text_size::Ranged; +use ruff_text_size::TextRange; +use starlark_map::Hashed; + +use crate::ModuleInfo; +use crate::binding::binding::Binding; +use crate::binding::binding::Key; +use crate::state::lsp::Transaction; + +/// Extend an unaliased import without changing its scope or binding a new name. +pub(crate) fn import_module_code_action( + transaction: &Transaction<'_>, + handle: &Handle, + module_info: &ModuleInfo, + ast: &ModModule, + error_range: TextRange, + module: ModuleName, +) -> Option<(String, Module, TextRange, String)> { + let receiver = Ast::locate_node(ast, error_range.start()) + .into_iter() + .find_map(|node| match node { + AnyNodeRef::ExprName(name) => Some(name), + _ => None, + })?; + if receiver.id != module.first_component() { + return None; + } + let answers = transaction.get_answers(handle)?; + let bindings = answers.bindings(); + let use_key = Key::BoundName(ShortIdentifier::expr_name(receiver)); + let mut idx = bindings.key_to_idx_hashed_opt(Hashed::new(&use_key))?; + let mut gas = Gas::new(100); + let import_range = loop { + if gas.stop() { + return None; + } + let key = bindings.idx_to_key(idx); + match bindings.get(idx) { + Binding::Module(_) if matches!(key, Key::Import(..) | Key::Definition(..)) => { + break key.range(); + } + Binding::Forward(next) + | Binding::PromoteForward(next) + | Binding::ForwardToFirstUse(next) + | Binding::Narrow(next, _, _) => idx = *next, + // Navigation selects one source at joins. An edit must not select one + // branch or ignore imports carried over from another loop iteration. + Binding::Phi(..) | Binding::LoopPhi(..) => return None, + _ => return None, + } + }; + let nodes = Ast::locate_node(ast, import_range.start()); + if !nodes + .iter() + .any(|node| matches!(node, AnyNodeRef::StmtImport(_))) + { + return None; + } + let alias = nodes.into_iter().find_map(|node| match node { + AnyNodeRef::Alias(alias) => Some(alias), + _ => None, + })?; + // Aliases need a reference rewrite or an extra binding, which can shadow names + // or produce an unused import that the editor would offer to remove. + if alias.asname.is_some() + || !module + .components() + .starts_with(&ModuleName::from_name(&alias.name.id).components()) + { + return None; + } + Some(( + format!("Insert import: `import {module}`"), + module_info.dupe(), + alias.name.range(), + module.as_str().to_owned(), + )) +} diff --git a/pyrefly/lib/state/lsp/quick_fixes/unused_ignore.rs b/pyrefly/lib/state/lsp/quick_fixes/unused_ignore.rs index 7635c08cab..d5833a0a21 100644 --- a/pyrefly/lib/state/lsp/quick_fixes/unused_ignore.rs +++ b/pyrefly/lib/state/lsp/quick_fixes/unused_ignore.rs @@ -25,7 +25,9 @@ pub(crate) fn remove_unused_ignore_code_action( let edit = error.quick_fixes().iter().find_map(|fix| match fix { ErrorQuickFix::RemoveUnusedSuppression(edit) => Some(edit), - ErrorQuickFix::ReplaceWithEnumMember { .. } | ErrorQuickFix::AssertNotNone => None, + ErrorQuickFix::ReplaceWithEnumMember { .. } + | ErrorQuickFix::AssertNotNone + | ErrorQuickFix::ImportModule { .. } => None, })?; let title = if edit.replacement.is_empty() { "Remove suppression comment"