fix

Author: asukaminato0721

pyrefly/lib/commands/infer.rs

@@ -10,6 +10,7 @@ use std::collections::HashSet;
 use std::path::Path;
 
 use clap::Parser;
+use dupe::Dupe;
 use pyrefly_config::args::ConfigOverrideArgs;
 use pyrefly_config::base::InferReturnTypes;
 use pyrefly_config::finder::ConfigFinder;
@@ -19,9 +20,8 @@ use pyrefly_util::forgetter::Forgetter;
 use pyrefly_util::fs_anyhow;
 use pyrefly_util::includes::Includes;
 use pyrefly_util::thread_pool::ThreadCount;
-use ruff_python_ast::Stmt;
-use ruff_python_ast::helpers::is_docstring_stmt;
 use ruff_text_size::Ranged;
+use ruff_text_size::TextRange;
 use ruff_text_size::TextSize;
 
 use crate::annotation::format_annotation;
@@ -30,6 +30,8 @@ use crate::commands::config_finder::ConfigConfigurerWrapper;
 use crate::commands::files::FilesArgs;
 use crate::commands::util::CommandExitStatus;
 use crate::lsp::wasm::inlay_hints::ParameterAnnotation;
+use crate::state::ide::ImportEdit;
+use crate::state::ide::insert_import_edit_with_forced_import_format;
 use crate::state::lsp::AnnotationKind;
 use crate::state::require::Require;
 use crate::state::state::State;
@@ -39,15 +41,6 @@ use crate::types::simplify::unions_with_literals;
 use crate::types::stdlib::Stdlib;
 use crate::types::types::Type;
 
-/// Check if a statement is a `from __future__ import ...` statement.
-/// New imports must be inserted after `__future__` imports to produce valid Python.
-fn is_future_import_stmt(stmt: &Stmt) -> bool {
-    matches!(
-        stmt,
-        Stmt::ImportFrom(import_from) if import_from.module.as_ref().is_some_and(|m| m.id == "__future__")
-    )
-}
-
 #[deny(clippy::missing_docs_in_private_items)]
 /// Flags for controlling the behavior of the autotype command
 #[derive(Debug, Clone, Parser)]
@@ -278,15 +271,6 @@ fn format_hints(
     (qualified_hints, needed_imports)
 }
 
-// Sort the hints by reverse order so we don't have to recalculate positions
-fn sort_inlay_hints(
-    inlay_hints: Vec<(ruff_text_size::TextSize, String)>,
-) -> Vec<(ruff_text_size::TextSize, String)> {
-    let mut sorted_inlay_hints = inlay_hints;
-    sorted_inlay_hints.sort_by(|(a, _), (b, _)| b.cmp(a));
-    sorted_inlay_hints
-}
-
 fn hint_to_string(
     hint: Type,
     stdlib: &Stdlib,
@@ -371,7 +355,7 @@ impl InferArgs {
             if let Some(inferred_types) = inferred_types {
                 parameter_types.extend(inferred_types);
                 let heap = TypeHeap::new();
-                let (formatted, needed_imports) = format_hints(
+                let (annotations, needed_imports) = format_hints(
                     parameter_types,
                     &stdlib,
                     &|cls| {
@@ -389,46 +373,44 @@ impl InferArgs {
                     &heap,
                     handle.module(),
                 );
-                let sorted = sort_inlay_hints(formatted);
                 let file_path = handle.path().as_path();
                 // Build import edits once so dry-run and write share them.
-                let imports: Vec<(TextSize, String, String)> = if flags.imports()
+                let imports: Vec<ImportEdit> = if flags.imports()
                     && !needed_imports.is_empty()
                     && let Some(ast) = transaction.get_ast(&handle)
                 {
-                    let position = ast
-                        .body
-                        .iter()
-                        .find(|stmt| !is_docstring_stmt(stmt) && !is_future_import_stmt(stmt))
-                        .map_or(ast.range.end(), |stmt| stmt.range().start());
-                    let mut imports: Vec<(TextSize, String, String)> = needed_imports
+                    needed_imports
                         .into_iter()
                         .map(|(module_name, name)| {
-                            let import_text =
-                                format!("from {} import {}\n", module_name.as_str(), name);
-                            (position, import_text, module_name.as_str().to_owned())
+                            let handle_to_import_from = transaction
+                                .import_handle(&handle, module_name, None)
+                                .finding()
+                                .expect("infer import source should always resolve");
+                            insert_import_edit_with_forced_import_format(
+                                &ast,
+                                handle.dupe(),
+                                handle_to_import_from.dupe(),
+                                &name,
+                                true,
+                            )
                         })
-                        .collect();
-                    imports.sort_by(|(_, a, _), (_, b, _)| a.cmp(b));
-                    imports
+                        .filter(|edit| !edit.insert_text.is_empty())
+                        .collect()
                 } else {
                     Vec::new()
                 };
                 if flags.dry_run() {
-                    if !sorted.is_empty() || !imports.is_empty() {
+                    if !annotations.is_empty() || !imports.is_empty() {
                         any_changes = true;
                         println!(
                             "{}: would add {} annotation(s) and {} import(s)",
                             file_path.display(),
-                            sorted.len(),
+                            annotations.len(),
                             imports.len()
                         );
                     }
                 } else {
-                    Self::add_annotations_to_file(file_path, sorted)?;
-                    if !imports.is_empty() {
-                        Self::add_imports_to_file(file_path, imports)?;
-                    }
+                    Self::apply_edits_to_file(file_path, annotations, imports)?;
                 }
             }
         }
@@ -439,34 +421,30 @@ impl InferArgs {
         })
     }
 
-    fn add_annotations_to_file(
-        file_path: &Path,
-        sorted: Vec<(TextSize, String)>,
-    ) -> anyhow::Result<()> {
-        let file_content = fs_anyhow::read_to_string(file_path)?;
-        let mut result = file_content;
-        for inlay_hint in sorted {
-            let (position, hint) = inlay_hint;
-            // Convert the TextSize to a byte offset
-            let offset = (position).into();
-            if offset <= result.len() {
-                result.insert_str(offset, &hint);
-            }
-        }
-        fs_anyhow::write(file_path, result)
-    }
-
-    fn add_imports_to_file(
+    /// Apply all edits from one AST snapshot in descending offset order.
+    fn apply_edits_to_file(
         file_path: &Path,
-        imports: Vec<(TextSize, String, String)>,
+        annotations: 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);
-            }
+        let mut result = fs_anyhow::read_to_string(file_path)?;
+        let mut edits: Vec<(TextRange, String)> = imports
+            .into_iter()
+            .map(|edit| (edit.range, edit.insert_text))
+            .chain(
+                annotations
+                    .into_iter()
+                    .map(|(position, text)| (TextRange::empty(position), text)),
+            )
+            .collect();
+        edits.sort_by(|(range1, text1), (range2, text2)| {
+            range2
+                .start()
+                .cmp(&range1.start())
+                .then_with(|| text1.cmp(text2))
+        });
+        for (range, edit_text) in edits {
+            result.replace_range(range.start().to_usize()..range.end().to_usize(), &edit_text);
         }
         fs_anyhow::write(file_path, result)
     }

pyrefly/lib/lsp/wasm/inlay_hints.rs

@@ -34,6 +34,7 @@ use crate::binding::binding::Key;
 use crate::binding::binding::KeyClassField;
 use crate::binding::binding::UnpackedPosition;
 use crate::binding::bindings::Bindings;
+use crate::state::ide::fallback_import_insertion_point;
 use crate::state::ide::import_regular_import_edit;
 use crate::state::ide::insert_import_edit;
 use crate::state::import_tracker::ImportTracker;
@@ -140,17 +141,19 @@ impl TypeHintRenderer<'_, '_> {
 
         let mut missing = missing.into_iter().collect::<Vec<_>>();
         missing.sort_by(|left, right| left.as_str().cmp(right.as_str()));
-        let mut import_edit = None;
+        // Inlay hints apply all imports as one edit, so every import must be a
+        // complete line at the same legal insertion point rather than a merge suffix.
+        let import_position = fallback_import_insertion_point(self.ast.as_ref());
+        let mut import_text = String::new();
         for module in missing {
             if let Some(handle_to_import) = self
                 .transaction
                 .import_handle(self.handle, module, None)
                 .finding()
             {
-                let (position, insert_text, _) =
+                let (_, insert_text, _) =
                     import_regular_import_edit(self.ast.as_ref(), handle_to_import, None);
-                let (_, combined) = import_edit.get_or_insert((position, String::new()));
-                combined.push_str(&insert_text);
+                import_text.push_str(&insert_text);
             }
         }
         for import in direct {
@@ -162,12 +165,13 @@ impl TypeHintRenderer<'_, '_> {
                 &import.heads.join(", "),
                 self.import_format,
             );
-            let (_, combined) = import_edit.get_or_insert((edit.range.start(), String::new()));
-            combined.push_str(&edit.insert_text);
+            import_text.push_str(&edit.display_text);
+            import_text.push('\n');
         }
-        let imports = match import_edit {
-            Some(import_edit) => vec![import_edit],
-            None => Vec::new(),
+        let imports = if import_text.is_empty() {
+            Vec::new()
+        } else {
+            vec![(import_position, import_text)]
         };
         InlayHintEdits {
             annotation: format!("{prefix}{text}"),