fix Inlay hint insertion should consider imports/use qualified paths to auto import correctly

crates/pyrefly_types/src/display.rs

@@ -7,9 +7,11 @@
 
 //! Display a type. The complexity comes from if we have two classes with the same name,
 //! we want to display disambiguating information (e.g. module name or location).
+use std::cell::RefCell;
 use std::fmt;
 use std::fmt::Display;
 
+use dupe::Dupe;
 use pyrefly_python::module_name::ModuleName;
 use pyrefly_python::qname::QName;
 use pyrefly_util::display::Fmt;
@@ -89,6 +91,8 @@ pub struct TypeDisplayContext<'a> {
     /// Should we display for IDE Hover? This makes type names more readable but less precise.
     hover: bool,
     always_display_module_name: bool,
+    /// Modules encountered while formatting, used downstream (e.g. to decide which imports are required).
+    modules: RefCell<SmallSet<ModuleName>>,
 }
 
 impl<'a> TypeDisplayContext<'a> {
@@ -234,6 +238,9 @@ impl<'a> TypeDisplayContext<'a> {
         name: &str,
         output: &mut impl TypeOutput,
     ) -> fmt::Result {
+        self.modules
+            .borrow_mut()
+            .insert(ModuleName::from_str(module));
         if self.always_display_module_name {
             // write!(f, "{module}.{name}")
             output.write_str(&format!("{}.{}", module, name))
@@ -243,6 +250,16 @@ impl<'a> TypeDisplayContext<'a> {
         }
     }
 
+    pub fn referenced_modules(&self) -> SmallSet<ModuleName> {
+        let mut modules = self.modules.borrow().clone();
+        for info in self.qnames.values() {
+            for module in info.info.keys() {
+                modules.insert(module.dupe());
+            }
+        }
+        modules
+    }
+
     fn fmt_helper_generic(
         &self,
         t: &Type,

pyrefly/lib/lsp/non_wasm/server.rs

@@ -2621,22 +2621,31 @@ impl Server {
         )?;
         let res = t
             .into_iter()
-            .filter_map(|(text_size, label_text, _locations)| {
+            .filter_map(|hint| {
                 // If the url is a notebook cell, filter out inlay hints for other cells
-                if info.to_cell_for_lsp(text_size) != maybe_cell_idx {
+                if info.to_cell_for_lsp(hint.position) != maybe_cell_idx {
                     return None;
                 }
-                let position = info.to_lsp_position(text_size);
+                let position = info.to_lsp_position(hint.position);
                 // The range is half-open, so the end position is exclusive according to the spec.
                 if position >= range.start && position < range.end {
+                    let mut text_edits = Vec::with_capacity(1 + hint.import_edits.len());
+                    text_edits.push(TextEdit {
+                        range: Range::new(position, position),
+                        new_text: hint.label.clone(),
+                    });
+                    for (offset, import_text) in hint.import_edits {
+                        let insert_position = info.to_lsp_position(offset);
+                        text_edits.push(TextEdit {
+                            range: Range::new(insert_position, insert_position),
+                            new_text: import_text,
+                        });
+                    }
                     Some(InlayHint {
                         position,
-                        label: InlayHintLabel::String(label_text.clone()),
+                        label: InlayHintLabel::String(hint.label),
                         kind: None,
-                        text_edits: Some(vec![TextEdit {
-                            range: Range::new(position, position),
-                            new_text: label_text,
-                        }]),
+                        text_edits: Some(text_edits),
                         tooltip: None,
                         padding_left: None,
                         padding_right: None,

pyrefly/lib/playground.rs

@@ -521,9 +521,12 @@ impl Playground {
             .get_module_info(handle)
             .zip(transaction.inlay_hints(handle, Default::default()))
             .map(|(info, hints)| {
-                hints.into_map(|(position, label, _locations)| {
-                    let position = Position::from_display_pos(info.display_pos(position));
-                    InlayHint { label, position }
+                hints.into_map(|hint| {
+                    let position = Position::from_display_pos(info.display_pos(hint.position));
+                    InlayHint {
+                        label: hint.label,
+                        position,
+                    }
                 })
             })
             .unwrap_or_default()

pyrefly/lib/state/import_tracker.rs

@@ -0,0 +1,212 @@
+/*
+ * 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.
+ */
+
+//! Helpers for harvesting imports and formatting type strings for inlay hints.
+
+use std::cmp::Reverse;
+
+use dupe::Dupe;
+use pyrefly_python::module_name::ModuleName;
+use ruff_python_ast::ModModule;
+use ruff_python_ast::Stmt;
+use ruff_python_ast::StmtImport;
+use starlark_map::small_set::SmallSet;
+
+use crate::types::display::TypeDisplayContext;
+use crate::types::types::Type;
+
+/// Tracks imports already present in a module and can determine which modules are still missing
+/// for a given set of referenced modules. Also supports alias-aware replacement when displaying
+/// type strings.
+#[derive(Default)]
+pub struct ImportTracker {
+    canonical_modules: SmallSet<ModuleName>,
+    alias_modules: Vec<(ModuleName, String)>,
+}
+
+impl ImportTracker {
+    /// Build an import tracker from the top-level `import ...` statements in a module.
+    pub fn from_ast(ast: &ModModule) -> Self {
+        let mut tracker = Self::default();
+        for stmt in &ast.body {
+            if let Stmt::Import(stmt_import) = stmt {
+                tracker.record_import(stmt_import);
+            }
+        }
+        tracker
+            .alias_modules
+            .sort_by_key(|(module, _)| Reverse(module.as_str().len()));
+        tracker
+    }
+
+    /// Record an `import ...` statement into the tracker.
+    pub fn record_import(&mut self, stmt_import: &StmtImport) {
+        for alias in &stmt_import.names {
+            let module_name = ModuleName::from_str(alias.name.as_str());
+            if let Some(asname) = &alias.asname {
+                self.alias_modules
+                    .push((module_name, asname.id.to_string()));
+            } else {
+                self.canonical_modules.insert(module_name);
+            }
+        }
+    }
+
+    /// Replace any module prefixes that have been imported under an alias (e.g. `import typing as t`).
+    pub fn apply_aliases(&self, text: &str) -> String {
+        if self.alias_modules.is_empty() {
+            return text.to_owned();
+        }
+        let bytes = text.as_bytes();
+        let mut result = String::with_capacity(text.len());
+        let mut i = 0;
+        while i < bytes.len() {
+            let mut replaced = false;
+            for (module, alias) in &self.alias_modules {
+                let module_str = module.as_str();
+                if module_str.is_empty() {
+                    continue;
+                }
+                let module_bytes = module_str.as_bytes();
+                if i + module_bytes.len() <= bytes.len()
+                    && &bytes[i..i + module_bytes.len()] == module_bytes
+                    && Self::is_boundary(bytes, i, i + module_bytes.len())
+                {
+                    result.push_str(alias);
+                    i += module_bytes.len();
+                    replaced = true;
+                    break;
+                }
+            }
+            if !replaced {
+                result.push(bytes[i] as char);
+                i += 1;
+            }
+        }
+        result
+    }
+
+    /// Modules that are referenced in the type string but not yet imported (excluding builtins/current).
+    pub fn missing_modules(
+        &self,
+        modules: &SmallSet<ModuleName>,
+        current_module: ModuleName,
+    ) -> SmallSet<ModuleName> {
+        let mut missing = SmallSet::new();
+        for module in modules.iter() {
+            let module = module.dupe();
+            if module.as_str().is_empty()
+                || module == current_module
+                || module == ModuleName::builtins()
+                || module == ModuleName::extra_builtins()
+            {
+                continue;
+            }
+            if self.module_is_imported(module) {
+                continue;
+            }
+            missing.insert(module);
+        }
+        missing
+    }
+
+    fn module_is_imported(&self, module: ModuleName) -> bool {
+        self.alias_for(module).is_some() || self.has_canonical(module)
+    }
+
+    fn alias_for(&self, module: ModuleName) -> Option<String> {
+        let target = module.as_str();
+        for (alias_module, alias_name) in &self.alias_modules {
+            let alias_module_str = alias_module.as_str();
+            if alias_module_str.is_empty() {
+                continue;
+            }
+            if target == alias_module_str {
+                return Some(alias_name.clone());
+            }
+            if target.len() > alias_module_str.len()
+                && target.starts_with(alias_module_str)
+                && target.as_bytes()[alias_module_str.len()] == b'.'
+            {
+                let remainder = &target[alias_module_str.len()..];
+                return Some(format!("{alias_name}{remainder}"));
+            }
+        }
+        None
+    }
+
+    fn has_canonical(&self, module: ModuleName) -> bool {
+        let target = module.as_str();
+        self.canonical_modules.iter().any(|imported| {
+            let imported_str = imported.as_str();
+            imported_str == target
+                || (target.len() > imported_str.len()
+                    && target.starts_with(imported_str)
+                    && target.as_bytes()[imported_str.len()] == b'.')
+        })
+    }
+
+    fn is_boundary(bytes: &[u8], start: usize, end: usize) -> bool {
+        (start == 0 || !Self::is_ident(bytes[start - 1]))
+            && (end == bytes.len() || !Self::is_ident(bytes[end]))
+    }
+
+    fn is_ident(byte: u8) -> bool {
+        matches!(byte, b'a'..=b'z' | b'A'..=b'Z' | b'0'..=b'9' | b'_')
+    }
+}
+
+/// Produce a user-facing type string (without module qualifiers) together with all referenced modules
+/// (captured with module qualification) so callers can insert the necessary imports.
+pub fn format_type_for_annotation(ty: &Type) -> (String, SmallSet<ModuleName>) {
+    // First pass: force module names so referenced_modules collects everything, but ignore the text.
+    let mut module_ctx = TypeDisplayContext::new(&[ty]);
+    module_ctx.always_display_module_name_except_builtins();
+    let _ = module_ctx.display(ty).to_string();
+    let modules = module_ctx.referenced_modules();
+
+    // Second pass: produce a concise label without module qualifiers.
+    let display_ctx = TypeDisplayContext::new(&[ty]);
+    let text = display_ctx.display(ty).to_string();
+    (text, modules)
+}
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+
+    #[test]
+    fn aliases_are_applied_at_boundaries_only() {
+        let module = ModuleName::from_str("typing");
+        let mut tracker = ImportTracker::default();
+        tracker.alias_modules.push((module, "t".to_owned()));
+        assert_eq!(tracker.apply_aliases("typing.Literal"), "t.Literal");
+        // Do not replace inside longer identifiers
+        assert_eq!(tracker.apply_aliases("mytyping"), "mytyping");
+    }
+
+    #[test]
+    fn missing_modules_skips_builtin_and_current() {
+        let tracker = ImportTracker::default();
+        let mut modules = SmallSet::new();
+        let current = ModuleName::from_str("pkg.mod");
+        modules.insert(current.dupe());
+        modules.insert(ModuleName::builtins());
+        modules.insert(ModuleName::from_str("typing"));
+        let missing = tracker.missing_modules(&modules, current);
+        assert!(missing.contains(&ModuleName::from_str("typing")));
+        assert_eq!(missing.len(), 1);
+    }
+
+    #[test]
+    fn format_type_collects_modules_but_returns_short_label() {
+        let ty = Type::LiteralString;
+        let (text, modules) = format_type_for_annotation(&ty);
+        assert_eq!(text, "LiteralString");
+        assert!(modules.contains(&ModuleName::from_str("typing")));
+    }
+}