fix

Author: asukaminato0721

pyrefly/lib/commands/infer.rs

@@ -25,6 +25,7 @@ use crate::commands::files::get_project_config_for_current_dir;
 use crate::commands::util::CommandExitStatus;
 use crate::config::error_kind::ErrorKind;
 use crate::state::ide::insert_import_edit_with_forced_import_format;
+use crate::state::ide::ImportEdit;
 use crate::state::lsp::AnnotationKind;
 use crate::state::lsp::ParameterAnnotation;
 use crate::state::require::Require;

pyrefly/lib/state/ide.rs

@@ -16,6 +16,8 @@ use ruff_python_ast::Expr;
 use ruff_python_ast::ModModule;
 use ruff_python_ast::helpers::is_docstring_stmt;
 use ruff_python_ast::name::Name;
+use ruff_python_ast::Stmt;
+use ruff_python_ast::StmtImportFrom;
 use ruff_text_size::Ranged;
 use ruff_text_size::TextRange;
 use ruff_text_size::TextSize;
@@ -32,6 +34,13 @@ use crate::state::lsp::ImportFormat;
 
 const KEY_TO_DEFINITION_INITIAL_GAS: Gas = Gas::new(100);
 
+#[derive(Clone, Debug, PartialEq, Eq)]
+pub struct ImportEdit {
+    pub position: TextSize,
+    pub insert_text: String,
+    pub display_text: String,
+}
+
 pub enum IntermediateDefinition {
     Local(Export),
     NamedImport(TextRange, ModuleName, Name, Option<TextRange>),
@@ -242,6 +251,25 @@ pub fn insert_import_edit_with_forced_import_format(
     } else {
         handle_to_import_from.module()
     };
+    let display_text = format!(
+        "from {} import {}",
+        module_name_to_import.as_str(),
+        export_name
+    );
+    if let Some((position, insert_text)) =
+        try_extend_existing_from_import(ast, module_name_to_import.as_str(), export_name)
+    {
+        return ImportEdit {
+            position,
+            insert_text,
+            display_text,
+        };
+    }
+    let position = if let Some(first_stmt) = ast.body.iter().find(|stmt| !is_docstring_stmt(stmt)) {
+        first_stmt.range().start()
+    } else {
+        ast.range.end()
+    };
     let insert_text = format!(
         "from {} import {}\n",
         module_name_to_import.as_str(),

pyrefly/lib/state/lsp.rs

@@ -1915,9 +1915,15 @@ impl<'a> Transaction<'a> {
                                 unknown_name,
                                 import_format,
                             );
-                            let range = TextRange::at(position, TextSize::new(0));
-                            let title = format!("Insert import: `{}`", insert_text.trim());
-                            code_actions.push((title, module_info.dupe(), range, insert_text));
+                            let range = TextRange::at(import_edit.position, TextSize::new(0));
+                            let title =
+                                format!("Insert import: `{}`", import_edit.display_text);
+                            code_actions.push((
+                                title,
+                                module_info.dupe(),
+                                range,
+                                import_edit.insert_text,
+                            ));
                         }
 
                         for module_name in self.search_modules_fuzzy(unknown_name) {
@@ -2445,8 +2451,11 @@ impl<'a> Transaction<'a> {
                         import_format,
                     );
                     let import_text_edit = TextEdit {
-                        range: module_info.to_lsp_range(TextRange::at(position, TextSize::new(0))),
-                        new_text: insert_text.clone(),
+                        range: module_info.to_lsp_range(TextRange::at(
+                            import_edit.position,
+                            TextSize::new(0),
+                        )),
+                        new_text: import_edit.insert_text.clone(),
                     };
                     (insert_text, Some(vec![import_text_edit]), module_name)
                 };

pyrefly/lib/test/lsp/code_actions.rs

@@ -210,8 +210,7 @@ fn insertion_test_duplicate_imports() {
         ],
         get_test_report,
     );
-    // The insertion won't attempt to merge imports from the same module.
-    // It's not illegal, but it would be nice if we do merge.
+    // When another import from the same module already exists, we should append to it.
     assert_eq!(
         r#"
 # a.py
@@ -227,8 +226,7 @@ from a import another_thing
 my_export
 # ^
 ## After:
-from a import my_export
-from a import another_thing
+from a import another_thing, my_export
 my_export
 # ^
 "#