Add snapshot tests for resolver (#5404)

## Summary

This PR adds some snapshot tests for the resolver based on executing
resolutions within a "mock" of the Airflow repo (that is: a folder that
contains a subset of the repo's files, but all empty, and with an
only-partially-complete virtual environment). It's intended to act as a
lightweight integration test, to enable us to test resolutions on a
"real" project without adding a dependency on Airflow itself.

crates/ruff_python_resolver/Cargo.toml

@@ -19,3 +19,4 @@ log = { workspace = true }
 [dev-dependencies]
 env_logger = "0.10.0"
 tempfile = "3.6.0"
+insta = { workspace = true }

crates/ruff_python_resolver/resources/test/airflow/README.md

@@ -0,0 +1,3 @@
+# airflow
+
+This is a mock subset of the Airflow repository, used to test module resolution.

crates/ruff_python_resolver/resources/test/airflow/airflow/__init__.py

@@ -0,0 +1 @@
+

crates/ruff_python_resolver/resources/test/airflow/airflow/api/__init__.py

@@ -0,0 +1 @@
+

crates/ruff_python_resolver/resources/test/airflow/airflow/api/common/__init__.py

@@ -0,0 +1 @@
+

crates/ruff_python_resolver/resources/test/airflow/airflow/api/common/mark_tasks.py

@@ -0,0 +1,14 @@
+# Standard library.
+import os
+
+# First-party.
+from airflow.jobs.scheduler_job_runner import SchedulerJobRunner
+
+# Stub file.
+from airflow.compat.functools import cached_property
+
+# Namespace package.
+from airflow.providers.google.cloud.hooks.gcs import GCSHook
+
+# Third-party.
+from sqlalchemy.orm import Query

crates/ruff_python_resolver/resources/test/airflow/airflow/compat/__init__.py

@@ -0,0 +1,16 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements.  See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership.  The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License.  You may obtain a copy of the License at
+#
+#   http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied.  See the License for the
+# specific language governing permissions and limitations
+# under the License.

crates/ruff_python_resolver/resources/test/airflow/airflow/compat/functools.py

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

crates/ruff_python_resolver/resources/test/airflow/airflow/compat/functools.pyi

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

crates/ruff_python_resolver/resources/test/airflow/airflow/jobs/scheduler_job_runner.py

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

crates/ruff_python_resolver/resources/test/airflow/airflow/providers/google/cloud/hooks/gcs.py

@@ -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/sqlalchemy/__init__.py

@@ -0,0 +1 @@
+

crates/ruff_python_resolver/resources/test/airflow/venv/lib/python3.11/site-packages/sqlalchemy/orm/base.py

@@ -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/sqlalchemy/orm/dependency.py

@@ -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/sqlalchemy/orm/query.py

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

crates/ruff_python_resolver/src/implicit_imports.rs

@@ -1,4 +1,4 @@
-use std::collections::HashMap;
+use std::collections::BTreeMap;
 use std::ffi::OsStr;
 use std::fs;
 use std::path::{Path, PathBuf};
@@ -24,8 +24,8 @@ pub(crate) struct ImplicitImport {
 }
 
 /// Find the "implicit" imports within the namespace package at the given path.
-pub(crate) fn find(dir_path: &Path, exclusions: &[&Path]) -> HashMap<String, ImplicitImport> {
-    let mut implicit_imports = HashMap::new();
+pub(crate) fn find(dir_path: &Path, exclusions: &[&Path]) -> BTreeMap<String, ImplicitImport> {
+    let mut implicit_imports = BTreeMap::new();
 
     // Enumerate all files and directories in the path, expanding links.
     let Ok(entries) = fs::read_dir(dir_path) else {
@@ -128,14 +128,14 @@ pub(crate) fn find(dir_path: &Path, exclusions: &[&Path]) -> HashMap<String, Imp
 
 /// Filter a map of implicit imports to only include those that were actually imported.
 pub(crate) fn filter(
-    implicit_imports: &HashMap<String, ImplicitImport>,
+    implicit_imports: &BTreeMap<String, ImplicitImport>,
     imported_symbols: &[String],
-) -> Option<HashMap<String, ImplicitImport>> {
+) -> Option<BTreeMap<String, ImplicitImport>> {
     if implicit_imports.is_empty() || imported_symbols.is_empty() {
         return None;
     }
 
-    let mut filtered_imports = HashMap::new();
+    let mut filtered_imports = BTreeMap::new();
     for implicit_import in implicit_imports.values() {
         if imported_symbols.contains(&implicit_import.name) {
             filtered_imports.insert(implicit_import.name.clone(), implicit_import.clone());

crates/ruff_python_resolver/src/import_result.rs

@@ -1,9 +1,9 @@
 //! Interface that describes the output of the import resolver.
 
-use crate::implicit_imports::ImplicitImport;
-use std::collections::HashMap;
+use std::collections::BTreeMap;
 use std::path::PathBuf;
 
+use crate::implicit_imports::ImplicitImport;
 use crate::py_typed::PyTypedInfo;
 
 #[derive(Debug, Clone, PartialEq, Eq)]
@@ -69,11 +69,11 @@ pub(crate) struct ImportResult {
 
     /// A map from file to resolved path, for all implicitly imported
     /// modules that are part of a namespace package.
-    pub(crate) implicit_imports: HashMap<String, ImplicitImport>,
+    pub(crate) implicit_imports: BTreeMap<String, ImplicitImport>,
 
     /// Any implicit imports whose symbols were explicitly imported (i.e., via
     /// a `from x import y` statement).
-    pub(crate) filtered_implicit_imports: HashMap<String, ImplicitImport>,
+    pub(crate) filtered_implicit_imports: BTreeMap<String, ImplicitImport>,
 
     /// If the import resolved to a type hint (i.e., a `.pyi` file), then
     /// a non-type-hint resolution will be stored here.
@@ -105,8 +105,8 @@ impl ImportResult {
             is_stdlib_typeshed_file: false,
             is_third_party_typeshed_file: false,
             is_local_typings_file: false,
-            implicit_imports: HashMap::default(),
-            filtered_implicit_imports: HashMap::default(),
+            implicit_imports: BTreeMap::default(),
+            filtered_implicit_imports: BTreeMap::default(),
             non_stub_import_result: None,
             py_typed_info: None,
             package_directory: None,

crates/ruff_python_resolver/src/lib.rs

@@ -12,5 +12,3 @@ mod python_platform;
 mod python_version;
 mod resolver;
 mod search;
-
-pub(crate) const SITE_PACKAGES: &str = "site-packages";

crates/ruff_python_resolver/src/resolver.rs

@@ -1,6 +1,6 @@
 //! Resolves Python imports to their corresponding files on disk.
 
-use std::collections::HashMap;
+use std::collections::BTreeMap;
 use std::path::{Path, PathBuf};
 
 use log::debug;
@@ -36,7 +36,7 @@ fn resolve_module_descriptor(
     let mut is_stub_package = false;
     let mut is_stub_file = false;
     let mut is_native_lib = false;
-    let mut implicit_imports = HashMap::new();
+    let mut implicit_imports = BTreeMap::new();
     let mut package_directory = None;
     let mut py_typed_info = None;
 
@@ -194,7 +194,7 @@ fn resolve_module_descriptor(
         is_third_party_typeshed_file: false,
         is_local_typings_file: false,
         implicit_imports,
-        filtered_implicit_imports: HashMap::default(),
+        filtered_implicit_imports: BTreeMap::default(),
         non_stub_import_result: None,
         py_typed_info,
         package_directory,
@@ -424,7 +424,7 @@ fn resolve_best_absolute_import<Host: host::Host>(
 /// are all satisfied by submodules (as listed in the implicit imports).
 fn is_namespace_package_resolved(
     module_descriptor: &ImportModuleDescriptor,
-    implicit_imports: &HashMap<String, ImplicitImport>,
+    implicit_imports: &BTreeMap<String, ImplicitImport>,
 ) -> bool {
     if !module_descriptor.imported_symbols.is_empty() {
         // Pyright uses `!Array.from(moduleDescriptor.importedSymbols.keys()).some((symbol) => implicitImports.has(symbol))`.
@@ -774,6 +774,7 @@ fn resolve_import<Host: host::Host>(
 
 #[cfg(test)]
 mod tests {
+    use insta::assert_debug_snapshot;
     use std::fs::{create_dir_all, File};
     use std::io::{self, Write};
     use std::path::{Path, PathBuf};
@@ -823,6 +824,8 @@ mod tests {
         library: Option<PathBuf>,
         stub_path: Option<PathBuf>,
         typeshed_path: Option<PathBuf>,
+        venv_path: Option<PathBuf>,
+        venv: Option<PathBuf>,
     }
 
     fn resolve_options(
@@ -836,6 +839,8 @@ mod tests {
             library,
             stub_path,
             typeshed_path,
+            venv_path,
+            venv,
         } = options;
 
         let execution_environment = ExecutionEnvironment {
@@ -860,8 +865,8 @@ mod tests {
         let config = Config {
             typeshed_path,
             stub_path,
-            venv_path: None,
-            venv: None,
+            venv_path,
+            venv,
         };
 
         let host = host::StaticHost::new(if let Some(library) = library {
@@ -1545,4 +1550,109 @@ mod tests {
 
         Ok(())
     }
+
+    #[test]
+    fn airflow_standard_library() {
+        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,
+            "os",
+            root.clone(),
+            ResolverOptions {
+                venv_path: Some(root),
+                venv: Some(PathBuf::from("venv")),
+                ..Default::default()
+            },
+        );
+
+        assert_debug_snapshot!(result);
+    }
+
+    #[test]
+    fn airflow_first_party() {
+        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,
+            "airflow.jobs.scheduler_job_runner",
+            root.clone(),
+            ResolverOptions {
+                venv_path: Some(root),
+                venv: Some(PathBuf::from("venv")),
+                ..Default::default()
+            },
+        );
+
+        assert_debug_snapshot!(result);
+    }
+
+    #[test]
+    fn airflow_stub_file() {
+        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,
+            "airflow.compat.functools",
+            root.clone(),
+            ResolverOptions {
+                venv_path: Some(root),
+                venv: Some(PathBuf::from("venv")),
+                ..Default::default()
+            },
+        );
+
+        assert_debug_snapshot!(result);
+    }
+
+    #[test]
+    fn airflow_namespace_package() {
+        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,
+            "airflow.providers.google.cloud.hooks.gcs",
+            root.clone(),
+            ResolverOptions {
+                venv_path: Some(root),
+                venv: Some(PathBuf::from("venv")),
+                ..Default::default()
+            },
+        );
+
+        assert_debug_snapshot!(result);
+    }
+
+    #[test]
+    fn airflow_third_party() {
+        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,
+            "sqlalchemy.orm",
+            root.clone(),
+            ResolverOptions {
+                venv_path: Some(root),
+                venv: Some(PathBuf::from("venv")),
+                ..Default::default()
+            },
+        );
+
+        assert_debug_snapshot!(result);
+    }
 }

crates/ruff_python_resolver/src/search.rs

@@ -8,9 +8,11 @@ use std::path::{Path, PathBuf};
 use log::debug;
 
 use crate::config::Config;
+use crate::host;
 use crate::module_descriptor::ImportModuleDescriptor;
 use crate::python_version::PythonVersion;
-use crate::{host, SITE_PACKAGES};
+
+const SITE_PACKAGES: &str = "site-packages";
 
 /// Find the `site-packages` directory for the specified Python version.
 fn find_site_packages_path(

crates/ruff_python_resolver/src/snapshots/ruff_python_resolver__resolver__tests__airflow_first_party.snap

@@ -0,0 +1,33 @@
+---
+source: crates/ruff_python_resolver/src/resolver.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: Local,
+    resolved_paths: [
+        "./resources/test/airflow/airflow/__init__.py",
+        "./resources/test/airflow/airflow/jobs/__init__.py",
+        "./resources/test/airflow/airflow/jobs/scheduler_job_runner.py",
+    ],
+    search_path: Some(
+        "./resources/test/airflow",
+    ),
+    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: {},
+    filtered_implicit_imports: {},
+    non_stub_import_result: None,
+    py_typed_info: None,
+    package_directory: Some(
+        "./resources/test/airflow/airflow",
+    ),
+}

crates/ruff_python_resolver/src/snapshots/ruff_python_resolver__resolver__tests__airflow_namespace_package.snap

@@ -0,0 +1,36 @@
+---
+source: crates/ruff_python_resolver/src/resolver.rs
+expression: result
+---
+ImportResult {
+    is_relative: false,
+    is_import_found: true,
+    is_partly_resolved: false,
+    is_namespace_package: true,
+    is_init_file_present: true,
+    is_stub_package: false,
+    import_type: Local,
+    resolved_paths: [
+        "./resources/test/airflow/airflow/__init__.py",
+        "",
+        "./resources/test/airflow/airflow/providers/google/__init__.py",
+        "./resources/test/airflow/airflow/providers/google/cloud/__init__.py",
+        "./resources/test/airflow/airflow/providers/google/cloud/hooks/__init__.py",
+        "./resources/test/airflow/airflow/providers/google/cloud/hooks/gcs.py",
+    ],
+    search_path: Some(
+        "./resources/test/airflow",
+    ),
+    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: {},
+    filtered_implicit_imports: {},
+    non_stub_import_result: None,
+    py_typed_info: None,
+    package_directory: Some(
+        "./resources/test/airflow/airflow",
+    ),
+}

crates/ruff_python_resolver/src/snapshots/ruff_python_resolver__resolver__tests__airflow_standard_library.snap

@@ -0,0 +1,25 @@
+---
+source: crates/ruff_python_resolver/src/resolver.rs
+expression: result
+---
+ImportResult {
+    is_relative: false,
+    is_import_found: false,
+    is_partly_resolved: false,
+    is_namespace_package: false,
+    is_init_file_present: false,
+    is_stub_package: false,
+    import_type: Local,
+    resolved_paths: [],
+    search_path: None,
+    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: {},
+    filtered_implicit_imports: {},
+    non_stub_import_result: None,
+    py_typed_info: None,
+    package_directory: None,
+}

crates/ruff_python_resolver/src/snapshots/ruff_python_resolver__resolver__tests__airflow_stub_file.snap

@@ -0,0 +1,63 @@
+---
+source: crates/ruff_python_resolver/src/resolver.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: Local,
+    resolved_paths: [
+        "./resources/test/airflow/airflow/__init__.py",
+        "./resources/test/airflow/airflow/compat/__init__.py",
+        "./resources/test/airflow/airflow/compat/functools.pyi",
+    ],
+    search_path: Some(
+        "./resources/test/airflow",
+    ),
+    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: {},
+    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: Local,
+            resolved_paths: [
+                "./resources/test/airflow/airflow/__init__.py",
+                "./resources/test/airflow/airflow/compat/__init__.py",
+                "./resources/test/airflow/airflow/compat/functools.py",
+            ],
+            search_path: Some(
+                "./resources/test/airflow",
+            ),
+            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: {},
+            filtered_implicit_imports: {},
+            non_stub_import_result: None,
+            py_typed_info: None,
+            package_directory: Some(
+                "./resources/test/airflow/airflow",
+            ),
+        },
+    ),
+    py_typed_info: None,
+    package_directory: Some(
+        "./resources/test/airflow/airflow",
+    ),
+}

crates/ruff_python_resolver/src/snapshots/ruff_python_resolver__resolver__tests__airflow_third_party.snap

@@ -0,0 +1,54 @@
+---
+source: crates/ruff_python_resolver/src/resolver.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/sqlalchemy/__init__.py",
+        "./resources/test/airflow/venv/Lib/python3.11/site-packages/sqlalchemy/orm/__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: {
+        "base": ImplicitImport {
+            is_stub_file: false,
+            is_native_lib: false,
+            name: "base",
+            path: "./resources/test/airflow/venv/Lib/python3.11/site-packages/sqlalchemy/orm/base.py",
+            py_typed: None,
+        },
+        "dependency": ImplicitImport {
+            is_stub_file: false,
+            is_native_lib: false,
+            name: "dependency",
+            path: "./resources/test/airflow/venv/Lib/python3.11/site-packages/sqlalchemy/orm/dependency.py",
+            py_typed: None,
+        },
+        "query": ImplicitImport {
+            is_stub_file: false,
+            is_native_lib: false,
+            name: "query",
+            path: "./resources/test/airflow/venv/Lib/python3.11/site-packages/sqlalchemy/orm/query.py",
+            py_typed: None,
+        },
+    },
+    filtered_implicit_imports: {},
+    non_stub_import_result: None,
+    py_typed_info: None,
+    package_directory: Some(
+        "./resources/test/airflow/venv/Lib/python3.11/site-packages/sqlalchemy",
+    ),
+}