[red-knot] Allow multiple site-packages search paths (#12609)

AlexWaygood · web-flow · commit fbab04fbe1ea · 2024-08-02T13:33:19.000Z
diff --git a/crates/red_knot/src/main.rs b/crates/red_knot/src/main.rs
@@ -104,7 +104,7 @@ pub fn main() -> anyhow::Result<()> {
             extra_paths,
             workspace_root: workspace_metadata.root().to_path_buf(),
             custom_typeshed: custom_typeshed_dir,
-            site_packages: None,
+            site_packages: vec![],
         },
     };
 
diff --git a/crates/red_knot/tests/file_watching.rs b/crates/red_knot/tests/file_watching.rs
@@ -181,7 +181,7 @@ where
         extra_paths: vec![],
         workspace_root: workspace_path.to_path_buf(),
         custom_typeshed: None,
-        site_packages: None,
+        site_packages: vec![],
     })
 }
 
@@ -697,7 +697,7 @@ fn search_path() -> anyhow::Result<()> {
                 extra_paths: vec![],
                 workspace_root: workspace_path.to_path_buf(),
                 custom_typeshed: None,
-                site_packages: Some(root_path.join("site_packages")),
+                site_packages: vec![root_path.join("site_packages")],
             }
         })?;
 
@@ -734,7 +734,7 @@ fn add_search_path() -> anyhow::Result<()> {
 
     // Register site-packages as a search path.
     case.update_search_path_settings(|settings| SearchPathSettings {
-        site_packages: Some(site_packages.clone()),
+        site_packages: vec![site_packages.clone()],
         ..settings.clone()
     });
 
@@ -757,14 +757,14 @@ fn remove_search_path() -> anyhow::Result<()> {
                 extra_paths: vec![],
                 workspace_root: workspace_path.to_path_buf(),
                 custom_typeshed: None,
-                site_packages: Some(root_path.join("site_packages")),
+                site_packages: vec![root_path.join("site_packages")],
             }
         })?;
 
     // Remove site packages from the search path settings.
     let site_packages = case.root_path().join("site_packages");
     case.update_search_path_settings(|settings| SearchPathSettings {
-        site_packages: None,
+        site_packages: vec![],
         ..settings.clone()
     });
 
@@ -1175,7 +1175,7 @@ mod unix {
                 extra_paths: vec![],
                 workspace_root: workspace.to_path_buf(),
                 custom_typeshed: None,
-                site_packages: Some(workspace.join(".venv/lib/python3.12/site-packages")),
+                site_packages: vec![workspace.join(".venv/lib/python3.12/site-packages")],
             },
         )?;
 
diff --git a/crates/red_knot_module_resolver/src/path.rs b/crates/red_knot_module_resolver/src/path.rs
@@ -477,12 +477,6 @@ impl SearchPath {
         )
     }
 
-    /// Does this search path point to the `site-packages` directory?
-    #[must_use]
-    pub(crate) fn is_site_packages(&self) -> bool {
-        matches!(&*self.0, SearchPathInner::SitePackages(_))
-    }
-
     fn is_valid_extension(&self, extension: &str) -> bool {
         if self.is_standard_library() {
             extension == "pyi"
diff --git a/crates/red_knot_module_resolver/src/resolver.rs b/crates/red_knot_module_resolver/src/resolver.rs
@@ -160,12 +160,6 @@ fn try_resolve_module_resolution_settings(
         SearchPath::vendored_stdlib()
     });
 
-    if let Some(site_packages) = site_packages {
-        files.try_add_root(db.upcast(), site_packages, FileRootKind::LibrarySearchPath);
-
-        static_search_paths.push(SearchPath::site_packages(system, site_packages.clone())?);
-    };
-
     // TODO vendor typeshed's third-party stubs as well as the stdlib and fallback to them as a final step
 
     let target_version = program.target_version(db.upcast());
@@ -191,6 +185,7 @@ fn try_resolve_module_resolution_settings(
     Ok(ModuleResolutionSettings {
         target_version,
         static_search_paths,
+        site_packages_paths: site_packages.to_owned(),
     })
 }
 
@@ -200,67 +195,79 @@ pub(crate) fn module_resolution_settings(db: &dyn Db) -> ModuleResolutionSetting
     try_resolve_module_resolution_settings(db).unwrap()
 }
 
-/// Collect all dynamic search paths:
-/// search paths listed in `.pth` files in the `site-packages` directory
-/// due to editable installations of third-party packages.
+/// Collect all dynamic search paths. For each `site-packages` path:
+/// - Collect that `site-packages` path
+/// - Collect any search paths listed in `.pth` files in that `site-packages` directory
+///   due to editable installations of third-party packages.
+///
+/// The editable-install search paths for the first `site-packages` directory
+/// should come between the two `site-packages` directories when it comes to
+/// module-resolution priority.
 #[salsa::tracked(return_ref)]
-pub(crate) fn editable_install_resolution_paths(db: &dyn Db) -> Vec<SearchPath> {
-    let settings = module_resolution_settings(db);
-    let static_search_paths = &settings.static_search_paths;
-
-    let site_packages = static_search_paths
-        .iter()
-        .find(|path| path.is_site_packages());
-
-    let Some(site_packages) = site_packages else {
-        return Vec::new();
-    };
-
-    let site_packages = site_packages
-        .as_system_path()
-        .expect("Expected site-packages never to be a VendoredPath!");
-
-    let mut dynamic_paths = Vec::default();
+pub(crate) fn dynamic_resolution_paths(db: &dyn Db) -> Vec<SearchPath> {
+    let ModuleResolutionSettings {
+        target_version: _,
+        static_search_paths,
+        site_packages_paths,
+    } = module_resolution_settings(db);
 
-    // This query needs to be re-executed each time a `.pth` file
-    // is added, modified or removed from the `site-packages` directory.
-    // However, we don't use Salsa queries to read the source text of `.pth` files;
-    // we use the APIs on the `System` trait directly. As such, add a dependency on the
-    // site-package directory's revision.
-    if let Some(site_packages_root) = db.files().root(db.upcast(), site_packages) {
-        let _ = site_packages_root.revision(db.upcast());
-    }
+    let mut dynamic_paths = Vec::new();
 
-    // As well as modules installed directly into `site-packages`,
-    // the directory may also contain `.pth` files.
-    // Each `.pth` file in `site-packages` may contain one or more lines
-    // containing a (relative or absolute) path.
-    // Each of these paths may point to an editable install of a package,
-    // so should be considered an additional search path.
-    let Ok(pth_file_iterator) = PthFileIterator::new(db, site_packages) else {
+    if site_packages_paths.is_empty() {
         return dynamic_paths;
-    };
-
-    // The Python documentation specifies that `.pth` files in `site-packages`
-    // are processed in alphabetical order, so collecting and then sorting is necessary.
-    // https://docs.python.org/3/library/site.html#module-site
-    let mut all_pth_files: Vec<PthFile> = pth_file_iterator.collect();
-    all_pth_files.sort_by(|a, b| a.path.cmp(&b.path));
+    }
 
     let mut existing_paths: FxHashSet<_> = static_search_paths
         .iter()
         .filter_map(|path| path.as_system_path())
         .map(Cow::Borrowed)
         .collect();
 
-    dynamic_paths.reserve(all_pth_files.len());
+    let files = db.files();
+    let system = db.system();
 
-    for pth_file in &all_pth_files {
-        for installation in pth_file.editable_installations() {
-            if existing_paths.insert(Cow::Owned(
-                installation.as_system_path().unwrap().to_path_buf(),
-            )) {
-                dynamic_paths.push(installation);
+    for site_packages_dir in site_packages_paths {
+        if !existing_paths.insert(Cow::Borrowed(site_packages_dir)) {
+            continue;
+        }
+        let site_packages_root = files.try_add_root(
+            db.upcast(),
+            site_packages_dir,
+            FileRootKind::LibrarySearchPath,
+        );
+        // This query needs to be re-executed each time a `.pth` file
+        // is added, modified or removed from the `site-packages` directory.
+        // However, we don't use Salsa queries to read the source text of `.pth` files;
+        // we use the APIs on the `System` trait directly. As such, add a dependency on the
+        // site-package directory's revision.
+        site_packages_root.revision(db.upcast());
+
+        dynamic_paths
+            .push(SearchPath::site_packages(system, site_packages_dir.to_owned()).unwrap());
+
+        // As well as modules installed directly into `site-packages`,
+        // the directory may also contain `.pth` files.
+        // Each `.pth` file in `site-packages` may contain one or more lines
+        // containing a (relative or absolute) path.
+        // Each of these paths may point to an editable install of a package,
+        // so should be considered an additional search path.
+        let Ok(pth_file_iterator) = PthFileIterator::new(db, site_packages_dir) else {
+            continue;
+        };
+
+        // The Python documentation specifies that `.pth` files in `site-packages`
+        // are processed in alphabetical order, so collecting and then sorting is necessary.
+        // https://docs.python.org/3/library/site.html#module-site
+        let mut all_pth_files: Vec<PthFile> = pth_file_iterator.collect();
+        all_pth_files.sort_by(|a, b| a.path.cmp(&b.path));
+
+        for pth_file in &all_pth_files {
+            for installation in pth_file.editable_installations() {
+                if existing_paths.insert(Cow::Owned(
+                    installation.as_system_path().unwrap().to_path_buf(),
+                )) {
+                    dynamic_paths.push(installation);
+                }
             }
         }
     }
@@ -293,7 +300,7 @@ impl<'db> Iterator for SearchPathIterator<'db> {
 
         static_paths.next().or_else(|| {
             dynamic_paths
-                .get_or_insert_with(|| editable_install_resolution_paths(*db).iter())
+                .get_or_insert_with(|| dynamic_resolution_paths(*db).iter())
                 .next()
         })
     }
@@ -403,9 +410,18 @@ impl<'db> Iterator for PthFileIterator<'db> {
 #[derive(Clone, Debug, PartialEq, Eq)]
 pub(crate) struct ModuleResolutionSettings {
     target_version: TargetVersion,
+
     /// Search paths that have been statically determined purely from reading Ruff's configuration settings.
     /// These shouldn't ever change unless the config settings themselves change.
     static_search_paths: Vec<SearchPath>,
+
+    /// site-packages paths are not included in the above field:
+    /// if there are multiple site-packages paths, editable installations can appear
+    /// *between* the site-packages paths on `sys.path` at runtime.
+    /// That means we can't know where a second or third `site-packages` path should sit
+    /// in terms of module-resolution priority until we've discovered the editable installs
+    /// for the first `site-packages` path
+    site_packages_paths: Vec<SystemPathBuf>,
 }
 
 impl ModuleResolutionSettings {
@@ -630,6 +646,7 @@ mod tests {
     };
     use ruff_db::Db;
 
+    use crate::db::tests::TestDb;
     use crate::module::ModuleKind;
     use crate::module_name::ModuleName;
     use crate::testing::{FileSpec, MockedTypeshed, TestCase, TestCaseBuilder};
@@ -1180,7 +1197,7 @@ mod tests {
             extra_paths: vec![],
             workspace_root: src.clone(),
             custom_typeshed: Some(custom_typeshed.clone()),
-            site_packages: Some(site_packages.clone()),
+            site_packages: vec![site_packages],
         };
 
         Program::new(&db, TargetVersion::Py38, search_paths);
@@ -1578,7 +1595,7 @@ not_a_directory
             &FilePath::system("/y/src/bar.py")
         );
         let events = db.take_salsa_events();
-        assert_const_function_query_was_not_run(&db, editable_install_resolution_paths, &events);
+        assert_const_function_query_was_not_run(&db, dynamic_resolution_paths, &events);
     }
 
     #[test]
@@ -1656,4 +1673,53 @@ not_a_directory
         assert!(!search_paths
             .contains(&&SearchPath::editable(db.system(), SystemPathBuf::from("/src")).unwrap()));
     }
+
+    #[test]
+    fn multiple_site_packages_with_editables() {
+        let mut db = TestDb::new();
+
+        let venv_site_packages = SystemPathBuf::from("/venv-site-packages");
+        let site_packages_pth = venv_site_packages.join("foo.pth");
+        let system_site_packages = SystemPathBuf::from("/system-site-packages");
+        let editable_install_location = SystemPathBuf::from("/x/y/a.py");
+        let system_site_packages_location = system_site_packages.join("a.py");
+
+        db.memory_file_system()
+            .create_directory_all("/src")
+            .unwrap();
+        db.write_files([
+            (&site_packages_pth, "/x/y"),
+            (&editable_install_location, ""),
+            (&system_site_packages_location, ""),
+        ])
+        .unwrap();
+
+        Program::new(
+            &db,
+            TargetVersion::default(),
+            SearchPathSettings {
+                extra_paths: vec![],
+                workspace_root: SystemPathBuf::from("/src"),
+                custom_typeshed: None,
+                site_packages: vec![venv_site_packages, system_site_packages],
+            },
+        );
+
+        // The editable installs discovered from the `.pth` file in the first `site-packages` directory
+        // take precedence over the second `site-packages` directory...
+        let a_module_name = ModuleName::new_static("a").unwrap();
+        let a_module = resolve_module(&db, a_module_name.clone()).unwrap();
+        assert_eq!(a_module.file().path(&db), &editable_install_location);
+
+        db.memory_file_system()
+            .remove_file(&site_packages_pth)
+            .unwrap();
+        File::sync_path(&mut db, &site_packages_pth);
+
+        // ...But now that the `.pth` file in the first `site-packages` directory has been deleted,
+        // the editable install no longer exists, so the module now resolves to the file in the
+        // second `site-packages` directory
+        let a_module = resolve_module(&db, a_module_name).unwrap();
+        assert_eq!(a_module.file().path(&db), &system_site_packages_location);
+    }
 }
diff --git a/crates/red_knot_module_resolver/src/testing.rs b/crates/red_knot_module_resolver/src/testing.rs
@@ -12,6 +12,9 @@ pub(crate) struct TestCase<T> {
     pub(crate) db: TestDb,
     pub(crate) src: SystemPathBuf,
     pub(crate) stdlib: T,
+    // Most test cases only ever need a single `site-packages` directory,
+    // so this is a single directory instead of a `Vec` of directories,
+    // like it is in `ruff_db::Program`.
     pub(crate) site_packages: SystemPathBuf,
     pub(crate) target_version: TargetVersion,
 }
@@ -223,7 +226,7 @@ impl TestCaseBuilder<MockedTypeshed> {
                 extra_paths: vec![],
                 workspace_root: src.clone(),
                 custom_typeshed: Some(typeshed.clone()),
-                site_packages: Some(site_packages.clone()),
+                site_packages: vec![site_packages.clone()],
             },
         );
 
@@ -276,7 +279,7 @@ impl TestCaseBuilder<VendoredTypeshed> {
                 extra_paths: vec![],
                 workspace_root: src.clone(),
                 custom_typeshed: None,
-                site_packages: Some(site_packages.clone()),
+                site_packages: vec![site_packages.clone()],
             },
         );
 
diff --git a/crates/red_knot_python_semantic/src/semantic_model.rs b/crates/red_knot_python_semantic/src/semantic_model.rs
@@ -179,7 +179,7 @@ mod tests {
             SearchPathSettings {
                 extra_paths: vec![],
                 workspace_root: SystemPathBuf::from("/src"),
-                site_packages: None,
+                site_packages: vec![],
                 custom_typeshed: None,
             },
         );
diff --git a/crates/red_knot_python_semantic/src/types/infer.rs b/crates/red_knot_python_semantic/src/types/infer.rs
@@ -1515,7 +1515,7 @@ mod tests {
             SearchPathSettings {
                 extra_paths: Vec::new(),
                 workspace_root: SystemPathBuf::from("/src"),
-                site_packages: None,
+                site_packages: vec![],
                 custom_typeshed: None,
             },
         );
@@ -1532,7 +1532,7 @@ mod tests {
             SearchPathSettings {
                 extra_paths: Vec::new(),
                 workspace_root: SystemPathBuf::from("/src"),
-                site_packages: None,
+                site_packages: vec![],
                 custom_typeshed: Some(SystemPathBuf::from(typeshed)),
             },
         );
diff --git a/crates/red_knot_workspace/src/lint.rs b/crates/red_knot_workspace/src/lint.rs
@@ -326,7 +326,7 @@ mod tests {
             SearchPathSettings {
                 extra_paths: Vec::new(),
                 workspace_root,
-                site_packages: None,
+                site_packages: vec![],
                 custom_typeshed: None,
             },
         );
diff --git a/crates/red_knot_workspace/tests/check.rs b/crates/red_knot_workspace/tests/check.rs
diff --git a/crates/ruff_benchmark/benches/red_knot.rs b/crates/ruff_benchmark/benches/red_knot.rs
diff --git a/crates/ruff_db/src/program.rs b/crates/ruff_db/src/program.rs

Original file line number	Diff line number	Diff line change
`@@ -477,12 +477,6 @@ impl SearchPath {`
`477`	`477`	`)`
`478`	`478`	`}`
`479`	`479`
`480`		- /// Does this search path point to the `site-packages` directory?
`481`		`- #[must_use]`
`482`		`- pub(crate) fn is_site_packages(&self) -> bool {`
`483`		`- matches!(&*self.0, SearchPathInner::SitePackages(_))`
`484`		`- }`
`485`		`-`
`486`	`480`	`fn is_valid_extension(&self, extension: &str) -> bool {`
`487`	`481`	`if self.is_standard_library() {`
`488`	`482`	`extension == "pyi"`