fix add helper for Module exists, but was not imported explicitly. You are relying on other modules to load it. #4713

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

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<SuppressionEdit> for ErrorQuickFix {

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| {

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,

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;

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,
     })
 }
 

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(),
+    ))
+}

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"

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#"