Support .pyi files in ruff analyze graph

Author: charliermarsh

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(())

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)

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<Item = &'a FilePath> {
         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`).
-                let parent = import.parent();
+                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();
 
-                let module = resolve_module(self.db, &import).or_else(|| {
-                    // Attempt to resolve the module (e.g., given `from foo import bar`, look for `foo`).
+                    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));
 
-                    resolve_module(self.db, &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();
 
-                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))
+    }
 }