mirror of https://github.com/astral-sh/ruff
98 lines
3.8 KiB
Rust
98 lines
3.8 KiB
Rust
use ruff_db::files::{File, FilePath, system_path_to_file};
|
|
use ruff_db::system::SystemPath;
|
|
use ty_python_semantic::{
|
|
ModuleName, resolve_module, resolve_module_confident, resolve_real_module,
|
|
resolve_real_module_confident,
|
|
};
|
|
|
|
use crate::ModuleDb;
|
|
use crate::collector::CollectedImport;
|
|
|
|
/// Collect all imports for a given Python file.
|
|
pub(crate) struct Resolver<'a> {
|
|
db: &'a ModuleDb,
|
|
file: Option<File>,
|
|
}
|
|
|
|
impl<'a> Resolver<'a> {
|
|
/// Initialize a [`Resolver`] with a given [`ModuleDb`].
|
|
pub(crate) fn new(db: &'a ModuleDb, path: &SystemPath) -> Self {
|
|
// If we know the importing file we can potentially resolve more imports
|
|
let file = system_path_to_file(db, path).ok();
|
|
Self { db, file }
|
|
}
|
|
|
|
/// Resolve the [`CollectedImport`] into a [`FilePath`].
|
|
pub(crate) fn resolve(&self, import: CollectedImport) -> impl Iterator<Item = &'a FilePath> {
|
|
match import {
|
|
CollectedImport::Import(import) => {
|
|
// Attempt to resolve the module (e.g., given `import foo`, look for `foo`).
|
|
let file = self.resolve_module(&import);
|
|
|
|
// If the file is a stub, look for the corresponding source file.
|
|
let source_file = file
|
|
.is_some_and(|file| file.extension() == Some("pyi"))
|
|
.then(|| self.resolve_real_module(&import))
|
|
.flatten();
|
|
|
|
std::iter::once(file)
|
|
.chain(std::iter::once(source_file))
|
|
.flatten()
|
|
}
|
|
CollectedImport::ImportFrom(import) => {
|
|
// Attempt to resolve the member (e.g., given `from foo import bar`, look for `foo.bar`).
|
|
if let Some(file) = self.resolve_module(&import) {
|
|
// If the file is a stub, look for the corresponding source file.
|
|
let source_file = (file.extension() == Some("pyi"))
|
|
.then(|| self.resolve_real_module(&import))
|
|
.flatten();
|
|
|
|
return std::iter::once(Some(file))
|
|
.chain(std::iter::once(source_file))
|
|
.flatten();
|
|
}
|
|
|
|
// Attempt to resolve the module (e.g., given `from foo import bar`, look for `foo`).
|
|
let parent = import.parent();
|
|
let file = parent
|
|
.as_ref()
|
|
.and_then(|parent| self.resolve_module(parent));
|
|
|
|
// If the file is a stub, look for the corresponding source file.
|
|
let source_file = file
|
|
.is_some_and(|file| file.extension() == Some("pyi"))
|
|
.then(|| {
|
|
parent
|
|
.as_ref()
|
|
.and_then(|parent| self.resolve_real_module(parent))
|
|
})
|
|
.flatten();
|
|
|
|
std::iter::once(file)
|
|
.chain(std::iter::once(source_file))
|
|
.flatten()
|
|
}
|
|
}
|
|
}
|
|
|
|
/// Resolves a module name to a module.
|
|
pub(crate) fn resolve_module(&self, module_name: &ModuleName) -> Option<&'a FilePath> {
|
|
let module = if let Some(file) = self.file {
|
|
resolve_module(self.db, file, module_name)?
|
|
} else {
|
|
resolve_module_confident(self.db, module_name)?
|
|
};
|
|
Some(module.file(self.db)?.path(self.db))
|
|
}
|
|
|
|
/// Resolves a module name to a module (stubs not allowed).
|
|
fn resolve_real_module(&self, module_name: &ModuleName) -> Option<&'a FilePath> {
|
|
let module = if let Some(file) = self.file {
|
|
resolve_real_module(self.db, file, module_name)?
|
|
} else {
|
|
resolve_real_module_confident(self.db, module_name)?
|
|
};
|
|
Some(module.file(self.db)?.path(self.db))
|
|
}
|
|
}
|