[ty] Add a new abstraction for adding imports to a module

This is somewhat inspired by a similar abstraction in
`ruff_linter`. The main idea is to create an importer once
for a module that you want to add imports to. And then call
`import` to generate an edit for each symbol you want to
add.

I haven't done any performance profiling here yet. I don't
know if it will be a bottleneck. In particular, I do expect
`Importer::import` (but not `Importer::new`) to get called
many times for a single completion request when auto-import
is enabled. Particularly in projects with a lot of unimported
symbols. Because I don't know the perf impact, I didn't do
any premature optimization here. But there are surely some
low hanging fruit if this does prove to be a problem.

New tests make up a big portion of the diff here. I tried to
think of a bunch of different cases, although I'm sure there
are more.
This commit is contained in:
Andrew Gallant 2025-09-16 13:29:41 -04:00 committed by Andrew Gallant
parent bcc8d6910b
commit b6a29592e7
3 changed files with 1812 additions and 32 deletions

File diff suppressed because it is too large Load diff

View file

@ -14,6 +14,7 @@ mod goto_definition;
mod goto_references;
mod goto_type_definition;
mod hover;
mod importer;
mod inlay_hints;
mod markup;
mod references;
@ -294,7 +295,10 @@ mod tests {
use ruff_db::Db;
use ruff_db::diagnostic::{Diagnostic, DiagnosticFormat, DisplayDiagnosticConfig};
use ruff_db::files::{File, FileRootKind, system_path_to_file};
use ruff_db::parsed::{ParsedModuleRef, parsed_module};
use ruff_db::source::{SourceText, source_text};
use ruff_db::system::{DbWithWritableSystem, SystemPath, SystemPathBuf};
use ruff_python_codegen::Stylist;
use ruff_python_trivia::textwrap::dedent;
use ruff_text_size::TextSize;
use ty_project::ProjectMetadata;
@ -350,11 +354,17 @@ mod tests {
}
}
/// The file and offset into that file containing
/// a `<CURSOR>` marker.
/// The file and offset into that file where a `<CURSOR>` marker
/// is located.
///
/// (Along with other information about that file, such as the
/// parsed AST.)
pub(super) struct Cursor {
pub(super) file: File,
pub(super) offset: TextSize,
pub(super) parsed: ParsedModuleRef,
pub(super) source: SourceText,
pub(super) stylist: Stylist<'static>,
}
#[derive(Default)]
@ -371,8 +381,20 @@ mod tests {
SystemPathBuf::from("/"),
));
let mut cursor: Option<Cursor> = None;
let search_paths = SearchPathSettings::new(vec![SystemPathBuf::from("/")])
.to_search_paths(db.system(), db.vendored())
.expect("Valid search path settings");
Program::from_settings(
&db,
ProgramSettings {
python_version: PythonVersionWithSource::default(),
python_platform: PythonPlatform::default(),
search_paths,
},
);
let mut cursor: Option<Cursor> = None;
for &Source {
ref path,
ref contents,
@ -405,25 +427,22 @@ mod tests {
cursor.is_none(),
"found more than one source that contains `<CURSOR>`"
);
cursor = Some(Cursor { file, offset });
let source = source_text(&db, file);
let parsed = parsed_module(&db, file).load(&db);
let stylist =
Stylist::from_tokens(parsed.tokens(), source.as_str()).into_owned();
cursor = Some(Cursor {
file,
offset,
parsed,
source,
stylist,
});
}
}
let search_paths = SearchPathSettings::new(vec![SystemPathBuf::from("/")])
.to_search_paths(db.system(), db.vendored())
.expect("Valid search path settings");
Program::from_settings(
&db,
ProgramSettings {
python_version: PythonVersionWithSource::default(),
python_platform: PythonPlatform::default(),
search_paths,
},
);
let mut insta_settings = insta::Settings::clone_current();
insta_settings.add_filter(r#"\\(\w\w|\s|\.|")"#, "/$1");
insta_settings.add_filter(r#"\\(\w\w|\.|")"#, "/$1");
// Filter out TODO types because they are different between debug and release builds.
insta_settings.add_filter(r"@Todo\(.+\)", "@Todo");

View file

@ -3,6 +3,7 @@ use ruff_db::source::line_index;
use ruff_python_ast as ast;
use ruff_python_ast::{Expr, ExprRef, HasNodeIndex, name::Name};
use ruff_source_file::LineIndex;
use rustc_hash::FxHashMap;
use crate::Db;
use crate::module_name::ModuleName;
@ -37,6 +38,37 @@ impl<'db> SemanticModel<'db> {
line_index(self.db, self.file)
}
/// Returns a map from symbol name to that symbol's
/// type and definition site (if available).
///
/// The symbols are the symbols in scope at the given
/// AST node.
pub fn members_in_scope_at(
&self,
node: ast::AnyNodeRef<'_>,
) -> FxHashMap<Name, MemberDefinition<'db>> {
let index = semantic_index(self.db, self.file);
let mut members = FxHashMap::default();
let Some(file_scope) = self.scope(node) else {
return members;
};
for (file_scope, _) in index.ancestor_scopes(file_scope) {
for memberdef in
all_declarations_and_bindings(self.db, file_scope.to_scope_id(self.db, self.file))
{
members.insert(
memberdef.member.name,
MemberDefinition {
ty: memberdef.member.ty,
definition: memberdef.definition,
},
);
}
}
members
}
pub fn resolve_module(&self, module_name: &ModuleName) -> Option<Module<'_>> {
resolve_module(self.db, module_name)
}
@ -212,20 +244,7 @@ impl<'db> SemanticModel<'db> {
pub fn scoped_completions(&self, node: ast::AnyNodeRef<'_>) -> Vec<Completion<'db>> {
let index = semantic_index(self.db, self.file);
// TODO: We currently use `try_expression_scope_id` here as a hotfix for [1].
// Revert this to use `expression_scope_id` once a proper fix is in place.
//
// [1] https://github.com/astral-sh/ty/issues/572
let Some(file_scope) = (match node {
ast::AnyNodeRef::Identifier(identifier) => index.try_expression_scope_id(identifier),
node => match node.as_expr_ref() {
// If we couldn't identify a specific
// expression that we're in, then just
// fall back to the global scope.
None => Some(FileScopeId::global()),
Some(expr) => index.try_expression_scope_id(&expr),
},
}) else {
let Some(file_scope) = self.scope(node) else {
return vec![];
};
let mut completions = vec![];
@ -244,6 +263,32 @@ impl<'db> SemanticModel<'db> {
completions.extend(self.module_completions(&builtins));
completions
}
fn scope(&self, node: ast::AnyNodeRef<'_>) -> Option<FileScopeId> {
let index = semantic_index(self.db, self.file);
// TODO: We currently use `try_expression_scope_id` here as a hotfix for [1].
// Revert this to use `expression_scope_id` once a proper fix is in place.
//
// [1] https://github.com/astral-sh/ty/issues/572
match node {
ast::AnyNodeRef::Identifier(identifier) => index.try_expression_scope_id(identifier),
node => match node.as_expr_ref() {
// If we couldn't identify a specific
// expression that we're in, then just
// fall back to the global scope.
None => Some(FileScopeId::global()),
Some(expr) => index.try_expression_scope_id(&expr),
},
}
}
}
/// The type and definition (if available) of a symbol.
#[derive(Clone, Debug)]
pub struct MemberDefinition<'db> {
pub ty: Type<'db>,
pub definition: Option<Definition<'db>>,
}
/// A classification of symbol names.