Add a snapshot test for native module resolution (#5423)

crates/ruff_python_resolver/resources/test/airflow/venv/lib/python3.11/site-packages/_watchdog_fsevents.cpython-311-darwin.so

@@ -0,0 +1 @@
+# Empty file included to support filesystem-based resolver tests.

crates/ruff_python_resolver/resources/test/airflow/venv/lib/python3.11/site-packages/orjson/orjson.cpython-311-darwin.so

@@ -0,0 +1 @@
+# Empty file included to support filesystem-based resolver tests.

crates/ruff_python_resolver/src/lib.rs

@@ -896,4 +896,46 @@ mod tests {
 
         assert_debug_snapshot!(result);
     }
+
+    #[test]
+    fn airflow_explicit_native_module() {
+        setup();
+
+        let root = PathBuf::from("./resources/test/airflow");
+        let source_file = root.join("airflow/api/common/mark_tasks.py");
+
+        let result = resolve_options(
+            source_file,
+            "_watchdog_fsevents",
+            root.clone(),
+            ResolverOptions {
+                venv_path: Some(root),
+                venv: Some(PathBuf::from("venv")),
+                ..Default::default()
+            },
+        );
+
+        assert_debug_snapshot!(result);
+    }
+
+    #[test]
+    fn airflow_implicit_native_module() {
+        setup();
+
+        let root = PathBuf::from("./resources/test/airflow");
+        let source_file = root.join("airflow/api/common/mark_tasks.py");
+
+        let result = resolve_options(
+            source_file,
+            "orjson",
+            root.clone(),
+            ResolverOptions {
+                venv_path: Some(root),
+                venv: Some(PathBuf::from("venv")),
+                ..Default::default()
+            },
+        );
+
+        assert_debug_snapshot!(result);
+    }
 }

crates/ruff_python_resolver/src/native_module.rs

@@ -1,6 +1,7 @@
 //! Support for native Python extension modules.
 
 use std::ffi::OsStr;
+use std::io;
 use std::path::{Path, PathBuf};
 
 /// Returns `true` if the given file extension is that of a native module.
@@ -37,15 +38,16 @@ pub(crate) fn is_native_module_file_name(module_name: &str, file_name: &Path) ->
 }
 
 /// Find the native module within the namespace package at the given path.
-pub(crate) fn find_native_module(dir_path: &Path) -> Option<PathBuf> {
-    let module_name = dir_path.file_name()?.to_str()?;
-    dir_path
-        .read_dir()
-        .ok()?
+pub(crate) fn find_native_module(
+    module_name: &str,
+    dir_path: &Path,
+) -> io::Result<Option<PathBuf>> {
+    Ok(dir_path
+        .read_dir()?
         .flatten()
         .filter(|entry| entry.file_type().map_or(false, |ft| ft.is_file()))
         .map(|entry| entry.path())
-        .find(|path| is_native_module_file_name(module_name, path))
+        .find(|path| is_native_module_file_name(module_name, path)))
 }
 
 #[cfg(test)]

crates/ruff_python_resolver/src/resolver.rs

@@ -1,6 +1,7 @@
 //! Resolves Python imports to their corresponding files on disk.
 
 use std::collections::BTreeMap;
+use std::ffi::OsStr;
 use std::path::{Path, PathBuf};
 
 use log::debug;
@@ -66,22 +67,22 @@ fn resolve_module_descriptor(
             let is_last_part = i == module_descriptor.name_parts.len() - 1;
 
             // Extend the directory path with the next segment.
-            if use_stub_package && is_first_part {
-                dir_path = dir_path.join(format!("{part}-stubs"));
+            let module_dir_path = if use_stub_package && is_first_part {
                 is_stub_package = true;
+                dir_path.join(format!("{part}-stubs"))
             } else {
-                dir_path = dir_path.join(part);
-            }
+                dir_path.join(part)
+            };
 
-            let found_directory = dir_path.is_dir();
+            let found_directory = module_dir_path.is_dir();
             if found_directory {
                 if is_first_part {
-                    package_directory = Some(dir_path.clone());
+                    package_directory = Some(module_dir_path.clone());
                 }
 
                 // Look for an `__init__.py[i]` in the directory.
-                let py_file_path = dir_path.join("__init__.py");
-                let pyi_file_path = dir_path.join("__init__.pyi");
+                let py_file_path = module_dir_path.join("__init__.py");
+                let pyi_file_path = module_dir_path.join("__init__.pyi");
                 is_init_file_present = false;
 
                 if allow_pyi && pyi_file_path.is_file() {
@@ -99,7 +100,7 @@ fn resolve_module_descriptor(
 
                 if look_for_py_typed {
                     py_typed_info =
-                        py_typed_info.or_else(|| py_typed::get_py_typed_info(&dir_path));
+                        py_typed_info.or_else(|| py_typed::get_py_typed_info(&module_dir_path));
                 }
 
                 // We haven't reached the end of the import, and we found a matching directory.
@@ -110,21 +111,23 @@ fn resolve_module_descriptor(
                         is_namespace_package = true;
                         py_typed_info = None;
                     }
+
+                    dir_path = module_dir_path;
                     continue;
                 }
 
                 if is_init_file_present {
                     implicit_imports =
-                        implicit_imports::find(&dir_path, &[&py_file_path, &pyi_file_path]);
+                        implicit_imports::find(&module_dir_path, &[&py_file_path, &pyi_file_path]);
                     break;
                 }
             }
 
             // We couldn't find a matching directory, or the directory didn't contain an
             // `__init__.py[i]` file. Look for an `.py[i]` file with the same name as the
             // segment, in lieu of a directory.
-            let py_file_path = dir_path.with_extension("py");
-            let pyi_file_path = dir_path.with_extension("pyi");
+            let py_file_path = module_dir_path.with_extension("py");
+            let pyi_file_path = module_dir_path.with_extension("pyi");
 
             if allow_pyi && pyi_file_path.is_file() {
                 debug!("Resolved import with file: {pyi_file_path:?}");
@@ -138,10 +141,14 @@ fn resolve_module_descriptor(
             } else {
                 if allow_native_lib && dir_path.is_dir() {
                     // We couldn't find a `.py[i]` file; search for a native library.
-                    if let Some(native_lib_path) = native_module::find_native_module(&dir_path) {
-                        debug!("Resolved import with file: {native_lib_path:?}");
-                        is_native_lib = true;
-                        resolved_paths.push(native_lib_path);
+                    if let Some(module_name) = module_dir_path.file_name().and_then(OsStr::to_str) {
+                        if let Ok(Some(native_lib_path)) =
+                            native_module::find_native_module(module_name, &dir_path)
+                        {
+                            debug!("Resolved import with file: {native_lib_path:?}");
+                            is_native_lib = true;
+                            resolved_paths.push(native_lib_path);
+                        }
                     }
                 }
 
@@ -153,10 +160,9 @@ fn resolve_module_descriptor(
                             implicit_imports::find(&dir_path, &[&py_file_path, &pyi_file_path]);
                         is_namespace_package = true;
                     }
-                } else if is_native_lib {
-                    debug!("Did not find file {py_file_path:?} or {pyi_file_path:?}");
                 }
             }
+
             break;
         }
     }

crates/ruff_python_resolver/src/snapshots/ruff_python_resolver__tests__airflow_explicit_native_module.snap

@@ -0,0 +1,29 @@
+---
+source: crates/ruff_python_resolver/src/lib.rs
+expression: result
+---
+ImportResult {
+    is_relative: false,
+    is_import_found: true,
+    is_partly_resolved: false,
+    is_namespace_package: false,
+    is_init_file_present: false,
+    is_stub_package: false,
+    import_type: ThirdParty,
+    resolved_paths: [
+        "./resources/test/airflow/venv/lib/python3.11/site-packages/_watchdog_fsevents.cpython-311-darwin.so",
+    ],
+    search_path: Some(
+        "./resources/test/airflow/venv/lib/python3.11/site-packages",
+    ),
+    is_stub_file: false,
+    is_native_lib: true,
+    is_stdlib_typeshed_file: false,
+    is_third_party_typeshed_file: false,
+    is_local_typings_file: false,
+    implicit_imports: {},
+    filtered_implicit_imports: {},
+    non_stub_import_result: None,
+    py_typed_info: None,
+    package_directory: None,
+}

crates/ruff_python_resolver/src/snapshots/ruff_python_resolver__tests__airflow_implicit_native_module.snap

@@ -0,0 +1,85 @@
+---
+source: crates/ruff_python_resolver/src/lib.rs
+expression: result
+---
+ImportResult {
+    is_relative: false,
+    is_import_found: true,
+    is_partly_resolved: false,
+    is_namespace_package: false,
+    is_init_file_present: true,
+    is_stub_package: false,
+    import_type: ThirdParty,
+    resolved_paths: [
+        "./resources/test/airflow/venv/lib/python3.11/site-packages/orjson/__init__.pyi",
+    ],
+    search_path: Some(
+        "./resources/test/airflow/venv/lib/python3.11/site-packages",
+    ),
+    is_stub_file: true,
+    is_native_lib: false,
+    is_stdlib_typeshed_file: false,
+    is_third_party_typeshed_file: false,
+    is_local_typings_file: false,
+    implicit_imports: {
+        "orjson": ImplicitImport {
+            is_stub_file: false,
+            is_native_lib: true,
+            name: "orjson",
+            path: "./resources/test/airflow/venv/lib/python3.11/site-packages/orjson/orjson.cpython-311-darwin.so",
+            py_typed: None,
+        },
+    },
+    filtered_implicit_imports: {},
+    non_stub_import_result: Some(
+        ImportResult {
+            is_relative: false,
+            is_import_found: true,
+            is_partly_resolved: false,
+            is_namespace_package: false,
+            is_init_file_present: true,
+            is_stub_package: false,
+            import_type: ThirdParty,
+            resolved_paths: [
+                "./resources/test/airflow/venv/lib/python3.11/site-packages/orjson/__init__.py",
+            ],
+            search_path: Some(
+                "./resources/test/airflow/venv/lib/python3.11/site-packages",
+            ),
+            is_stub_file: false,
+            is_native_lib: false,
+            is_stdlib_typeshed_file: false,
+            is_third_party_typeshed_file: false,
+            is_local_typings_file: false,
+            implicit_imports: {
+                "orjson": ImplicitImport {
+                    is_stub_file: false,
+                    is_native_lib: true,
+                    name: "orjson",
+                    path: "./resources/test/airflow/venv/lib/python3.11/site-packages/orjson/orjson.cpython-311-darwin.so",
+                    py_typed: None,
+                },
+            },
+            filtered_implicit_imports: {},
+            non_stub_import_result: None,
+            py_typed_info: Some(
+                PyTypedInfo {
+                    py_typed_path: "./resources/test/airflow/venv/lib/python3.11/site-packages/orjson/py.typed",
+                    is_partially_typed: false,
+                },
+            ),
+            package_directory: Some(
+                "./resources/test/airflow/venv/lib/python3.11/site-packages/orjson",
+            ),
+        },
+    ),
+    py_typed_info: Some(
+        PyTypedInfo {
+            py_typed_path: "./resources/test/airflow/venv/lib/python3.11/site-packages/orjson/py.typed",
+            is_partially_typed: false,
+        },
+    ),
+    package_directory: Some(
+        "./resources/test/airflow/venv/lib/python3.11/site-packages/orjson",
+    ),
+}