diff --git a/crates/ruff/tests/analyze_graph.rs b/crates/ruff/tests/analyze_graph.rs index 437702a2a8..4d70bf6270 100644 --- a/crates/ruff/tests/analyze_graph.rs +++ b/crates/ruff/tests/analyze_graph.rs @@ -57,33 +57,40 @@ fn dependencies() -> Result<()> { .write_str(indoc::indoc! {r#" def f(): pass "#})?; + root.child("ruff") + .child("e.pyi") + .write_str(indoc::indoc! {r#" + def f() -> None: ... + "#})?; insta::with_settings!({ filters => INSTA_FILTERS.to_vec(), }, { - assert_cmd_snapshot!(command().current_dir(&root), @r###" - success: true - exit_code: 0 - ----- stdout ----- - { - "ruff/__init__.py": [], - "ruff/a.py": [ - "ruff/b.py" - ], - "ruff/b.py": [ - "ruff/c.py" - ], - "ruff/c.py": [ - "ruff/d.py" - ], - "ruff/d.py": [ - "ruff/e.py" - ], - "ruff/e.py": [] - } + assert_cmd_snapshot!(command().current_dir(&root), @r#" + success: true + exit_code: 0 + ----- stdout ----- + { + "ruff/__init__.py": [], + "ruff/a.py": [ + "ruff/b.py" + ], + "ruff/b.py": [ + "ruff/c.py" + ], + "ruff/c.py": [ + "ruff/d.py" + ], + "ruff/d.py": [ + "ruff/e.py", + "ruff/e.pyi" + ], + "ruff/e.py": [], + "ruff/e.pyi": [] + } - ----- stderr ----- - "###); + ----- stderr ----- + "#); }); Ok(()) diff --git a/crates/ruff_graph/src/lib.rs b/crates/ruff_graph/src/lib.rs index bbe4182f5b..eaf307018d 100644 --- a/crates/ruff_graph/src/lib.rs +++ b/crates/ruff_graph/src/lib.rs @@ -42,13 +42,11 @@ impl ModuleImports { // Resolve the imports. let mut resolved_imports = ModuleImports::default(); for import in imports { - let Some(resolved) = Resolver::new(db).resolve(import) else { - continue; - }; - let Some(path) = resolved.as_system_path() else { - continue; - }; - resolved_imports.insert(path.to_path_buf()); + for resolved in Resolver::new(db).resolve(import) { + if let Some(path) = resolved.as_system_path() { + resolved_imports.insert(path.to_path_buf()); + } + } } Ok(resolved_imports) diff --git a/crates/ruff_graph/src/resolver.rs b/crates/ruff_graph/src/resolver.rs index 70b1b083bc..546b83b3fc 100644 --- a/crates/ruff_graph/src/resolver.rs +++ b/crates/ruff_graph/src/resolver.rs @@ -1,5 +1,5 @@ use ruff_db::files::FilePath; -use ty_python_semantic::resolve_module; +use ty_python_semantic::{ModuleName, resolve_module, resolve_real_module}; use crate::ModuleDb; use crate::collector::CollectedImport; @@ -16,24 +16,67 @@ impl<'a> Resolver<'a> { } /// Resolve the [`CollectedImport`] into a [`FilePath`]. - pub(crate) fn resolve(&self, import: CollectedImport) -> Option<&'a FilePath> { + pub(crate) fn resolve(&self, import: CollectedImport) -> impl Iterator { match import { CollectedImport::Import(import) => { - let module = resolve_module(self.db, &import)?; - Some(module.file(self.db)?.path(self.db)) + // 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)); - let module = resolve_module(self.db, &import).or_else(|| { - // Attempt to resolve the module (e.g., given `from foo import bar`, look for `foo`). + // 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(); - resolve_module(self.db, &parent?) - })?; - - Some(module.file(self.db)?.path(self.db)) + std::iter::once(file) + .chain(std::iter::once(source_file)) + .flatten() } } } + + /// Resolves a module name to a module. + fn resolve_module(&self, module_name: &ModuleName) -> Option<&'a FilePath> { + let module = resolve_module(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 = resolve_real_module(self.db, module_name)?; + Some(module.file(self.db)?.path(self.db)) + } }