Fix relative import resolution in site-packages packages when the site-packages search path is a subdirectory of the first-party search path (#17178)

## Summary

If a package in `site-packages` had this directory structure:

```py
# bar/__init__.py
from .a import A

# bar/a.py
class A: ...
```

then we would fail to resolve the `from .a import A` import _if_ (as is
usually the case!) the `site-packages` search path was located inside a
`.venv` directory that was a subdirectory of the project's first-party
search path. The reason for this is a bug in `file_to_module` in the
module resolver. In this loop, we would identify that
`/project_root/.venv/lib/python3.13/site-packages/foo/__init__.py` can
be turned into a path relative to the first-party search path
(`/project_root`):


https://github.com/astral-sh/ruff/blob/6e2b8f9696e87d786cfed6304dc3baa0f029d341/crates/red_knot_python_semantic/src/module_resolver/resolver.rs#L101-L110

but we'd then try to turn the relative path
(.venv/lib/python3.13/site-packages/foo/__init__.py`) into a module
path, realise that it wasn't a valid module path... and therefore
immediately `break` out of the loop before trying any other search paths
(such as the `site-packages` search path).

This bug was originally reported on Discord by @MatthewMckee4.

## Test Plan

I added a unit test for `file_to_module` in `resolver.rs`, and an
integration test that shows we can now resolve the import correctly in
`infer.rs`.

Author: AlexWaygood

crates/red_knot_python_semantic/src/db.rs

@@ -19,14 +19,14 @@ pub(crate) mod tests {
     use std::sync::Arc;
 
     use crate::program::{Program, SearchPathSettings};
-    use crate::{default_lint_registry, ProgramSettings, PythonPlatform};
+    use crate::{default_lint_registry, ProgramSettings, PythonPath, PythonPlatform};
 
     use super::Db;
     use crate::lint::{LintRegistry, RuleSelection};
     use anyhow::Context;
     use ruff_db::files::{File, Files};
     use ruff_db::system::{
-        DbWithTestSystem, DbWithWritableSystem as _, System, SystemPathBuf, TestSystem,
+        DbWithTestSystem, DbWithWritableSystem as _, System, SystemPath, SystemPathBuf, TestSystem,
     };
     use ruff_db::vendored::VendoredFileSystem;
     use ruff_db::{Db as SourceDb, Upcast};
@@ -139,8 +139,8 @@ pub(crate) mod tests {
         python_version: PythonVersion,
         /// Target Python platform
         python_platform: PythonPlatform,
-        /// Path to a custom typeshed directory
-        custom_typeshed: Option<SystemPathBuf>,
+        /// Paths to the directory to use for `site-packages`
+        site_packages: Vec<SystemPathBuf>,
         /// Path and content pairs for files that should be present
         files: Vec<(&'a str, &'a str)>,
     }
@@ -150,7 +150,7 @@ pub(crate) mod tests {
             Self {
                 python_version: PythonVersion::default(),
                 python_platform: PythonPlatform::default(),
-                custom_typeshed: None,
+                site_packages: vec![],
                 files: vec![],
             }
         }
@@ -160,8 +160,20 @@ pub(crate) mod tests {
             self
         }
 
-        pub(crate) fn with_file(mut self, path: &'a str, content: &'a str) -> Self {
-            self.files.push((path, content));
+        pub(crate) fn with_file(
+            mut self,
+            path: &'a (impl AsRef<SystemPath> + ?Sized),
+            content: &'a str,
+        ) -> Self {
+            self.files.push((path.as_ref().as_str(), content));
+            self
+        }
+
+        pub(crate) fn with_site_packages_search_path(
+            mut self,
+            path: &(impl AsRef<SystemPath> + ?Sized),
+        ) -> Self {
+            self.site_packages.push(path.as_ref().to_path_buf());
             self
         }
 
@@ -175,7 +187,7 @@ pub(crate) mod tests {
                 .context("Failed to write test files")?;
 
             let mut search_paths = SearchPathSettings::new(vec![src_root]);
-            search_paths.custom_typeshed = self.custom_typeshed;
+            search_paths.python_path = PythonPath::KnownSitePackages(self.site_packages);
 
             Program::from_settings(
                 &db,

crates/red_knot_python_semantic/src/module_resolver/resolver.rs

@@ -96,26 +96,21 @@ pub(crate) fn file_to_module(db: &dyn Db, file: File) -> Option<Module> {
         FilePath::SystemVirtual(_) => return None,
     };
 
-    let mut search_paths = search_paths(db);
-
-    let module_name = loop {
-        let candidate = search_paths.next()?;
+    let module_name = search_paths(db).find_map(|candidate| {
         let relative_path = match path {
             SystemOrVendoredPathRef::System(path) => candidate.relativize_system_path(path),
             SystemOrVendoredPathRef::Vendored(path) => candidate.relativize_vendored_path(path),
-        };
-        if let Some(relative_path) = relative_path {
-            break relative_path.to_module_name()?;
-        }
-    };
+        }?;
+        relative_path.to_module_name()
+    })?;
 
     // Resolve the module name to see if Python would resolve the name to the same path.
     // If it doesn't, then that means that multiple modules have the same name in different
     // root paths, but that the module corresponding to `path` is in a lower priority search path,
     // in which case we ignore it.
     let module = resolve_module(db, &module_name)?;
 
-    if file == module.file() {
+    if file.path(db) == module.file().path(db) {
         Some(module)
     } else {
         // This path is for a module with the same name but with a different precedence. For example:
@@ -1969,4 +1964,33 @@ not_a_directory
 
         Ok(())
     }
+
+    #[test]
+    fn file_to_module_where_one_search_path_is_subdirectory_of_other() {
+        let project_directory = SystemPathBuf::from("/project");
+        let site_packages = project_directory.join(".venv/lib/python3.13/site-packages");
+        let installed_foo_module = site_packages.join("foo/__init__.py");
+
+        let mut db = TestDb::new();
+        db.write_file(&installed_foo_module, "").unwrap();
+
+        Program::from_settings(
+            &db,
+            ProgramSettings {
+                python_version: PythonVersion::default(),
+                python_platform: PythonPlatform::default(),
+                search_paths: SearchPathSettings {
+                    extra_paths: vec![],
+                    src_roots: vec![project_directory],
+                    custom_typeshed: None,
+                    python_path: PythonPath::KnownSitePackages(vec![site_packages.clone()]),
+                },
+            },
+        )
+        .unwrap();
+
+        let foo_module_file = File::new(&db, FilePath::System(installed_foo_module));
+        let module = file_to_module(&db, foo_module_file).unwrap();
+        assert_eq!(module.search_path(), &site_packages);
+    }
 }

crates/red_knot_python_semantic/src/types/infer.rs

@@ -7371,15 +7371,15 @@ impl StringPartsCollector {
 
 #[cfg(test)]
 mod tests {
-    use crate::db::tests::{setup_db, TestDb};
+    use crate::db::tests::{setup_db, TestDb, TestDbBuilder};
     use crate::semantic_index::definition::Definition;
     use crate::semantic_index::symbol::FileScopeId;
     use crate::semantic_index::{global_scope, semantic_index, symbol_table, use_def_map};
     use crate::symbol::global_symbol;
     use crate::types::check_types;
     use ruff_db::diagnostic::Diagnostic;
     use ruff_db::files::{system_path_to_file, File};
-    use ruff_db::system::DbWithWritableSystem as _;
+    use ruff_db::system::{DbWithWritableSystem as _, SystemPath};
     use ruff_db::testing::{assert_function_query_was_not_run, assert_function_query_was_run};
 
     use super::*;
@@ -7535,6 +7535,26 @@ mod tests {
         Ok(())
     }
 
+    #[test]
+    fn relative_import_resolution_in_site_packages_when_site_packages_is_subdirectory_of_first_party_search_path(
+    ) {
+        let project_root = SystemPath::new("/src");
+        let foo_dot_py = project_root.join("foo.py");
+        let site_packages = project_root.join(".venv/lib/python3.13/site-packages");
+
+        let db = TestDbBuilder::new()
+            .with_site_packages_search_path(&site_packages)
+            .with_file(&foo_dot_py, "from bar import A")
+            .with_file(&site_packages.join("bar/__init__.py"), "from .a import *")
+            .with_file(&site_packages.join("bar/a.py"), "class A: ...")
+            .build()
+            .unwrap();
+
+        assert_file_diagnostics(&db, foo_dot_py.as_str(), &[]);
+        let a_symbol = get_symbol(&db, foo_dot_py.as_str(), &[], "A");
+        assert!(a_symbol.expect_type().is_class_literal());
+    }
+
     #[test]
     fn pep695_type_params() {
         let mut db = setup_db();