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;
@@ -289,7 +290,7 @@ impl InferArgs {
                             if let Some(ast) = transaction.get_ast(&handle) {
                                 let error_range = error.range();
                                 let unknown_name = module_info.code_at(error_range);
-                                let imports: Vec<(TextSize, String)> = transaction
+                                let imports: Vec<ImportEdit> = transaction
                                     .search_exports_exact(unknown_name)
                                     .into_iter()
                                     .map(|handle_to_import_from| {
@@ -336,14 +337,17 @@ impl InferArgs {
 
     fn add_imports_to_file(
         file_path: &Path,
-        imports: Vec<(TextSize, String)>,
+        imports: Vec<ImportEdit>,
     ) -> anyhow::Result<()> {
         let file_content = fs_anyhow::read_to_string(file_path)?;
         let mut result = file_content;
-        for (position, import) in imports {
-            let offset = (position).into();
-            if !result.contains(&import) {
-                result.insert_str(offset, &import);
+        for import_edit in imports {
+            if import_edit.insert_text.is_empty() {
+                continue;
+            }
+            let offset = (import_edit.position).into();
+            if offset <= result.len() && !result.contains(&import_edit.insert_text) {
+                result.insert_str(offset, &import_edit.insert_text);
             }
         }
         fs_anyhow::write(file_path, result)

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>),
@@ -189,7 +198,7 @@ pub fn insert_import_edit(
     handle_to_import_from: Handle,
     export_name: &str,
     import_format: ImportFormat,
-) -> (TextSize, String) {
+) -> ImportEdit {
     let use_absolute_import = match import_format {
         ImportFormat::Absolute => true,
         ImportFormat::Relative => {
@@ -226,12 +235,7 @@ pub fn insert_import_edit_with_forced_import_format(
     handle_to_import_from: Handle,
     export_name: &str,
     use_absolute_import: bool,
-) -> (TextSize, String) {
-    let position = if let Some(first_stmt) = ast.body.iter().find(|stmt| !is_docstring_stmt(stmt)) {
-        first_stmt.range().start()
-    } else {
-        ast.range.end()
-    };
+) -> ImportEdit {
     let module_name_to_import = if use_absolute_import {
         handle_to_import_from.module()
     } else if let Some(relative_module) = ModuleName::relative_module_name_between(
@@ -242,12 +246,65 @@ 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(),
         export_name
     );
-    (position, insert_text)
+    ImportEdit {
+        position,
+        insert_text,
+        display_text,
+    }
+}
+
+fn try_extend_existing_from_import(
+    ast: &ModModule,
+    target_module_name: &str,
+    export_name: &str,
+) -> Option<(TextSize, String)> {
+    for stmt in &ast.body {
+        if let Stmt::ImportFrom(import_from) = stmt {
+            if import_from_module_name(import_from) == target_module_name {
+                if let Some(last_alias) = import_from.names.last() {
+                    let position = last_alias.range.end();
+                    let insert_text = format!(", {}", export_name);
+                    return Some((position, insert_text));
+                }
+            }
+        }
+    }
+    None
+}
+
+fn import_from_module_name(import_from: &StmtImportFrom) -> String {
+    let mut module_name = String::new();
+    if import_from.level > 0 {
+        module_name.push_str(&".".repeat(import_from.level as usize));
+    }
+    if let Some(module) = &import_from.module {
+        module_name.push_str(module.as_str());
+    }
+    module_name
 }
 
 /// Some handles must be imported in absolute style,

pyrefly/lib/state/lsp.rs

@@ -1859,17 +1859,23 @@ impl<'a> Transaction<'a> {
                     if error_range.contains_range(range) {
                         let unknown_name = module_info.code_at(error_range);
                         for handle_to_import_from in self.search_exports_exact(unknown_name) {
-                            let (position, insert_text) = insert_import_edit(
+                            let import_edit = insert_import_edit(
                                 &ast,
                                 self.config_finder(),
                                 handle.dupe(),
                                 handle_to_import_from,
                                 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) {
@@ -2391,8 +2397,8 @@ impl<'a> Transaction<'a> {
                 {
                     continue;
                 }
-                let (insert_text, additional_text_edits) = {
-                    let (position, insert_text) = insert_import_edit(
+                let (detail_text, additional_text_edits) = {
+                    let import_edit = insert_import_edit(
                         &ast,
                         self.config_finder(),
                         handle.dupe(),
@@ -2401,14 +2407,17 @@ 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(),
                     };
-                    (Some(insert_text), Some(vec![import_text_edit]))
+                    (Some(import_edit.display_text), Some(vec![import_text_edit]))
                 };
                 completions.push(CompletionItem {
                     label: name,
-                    detail: insert_text,
+                    detail: detail_text,
                     kind: export
                         .symbol_kind
                         .map_or(Some(CompletionItemKind::VARIABLE), |k| {

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