Make site_packages field on Program an enum so we can continue to…

… mock out `site-packages` in tests
astral-sh · Aug 1, 2024 · 7a4763c · 7a4763c
1 parent ce7e9f7
commit 7a4763c
Show file tree

Hide file tree

Showing 10 changed files with 62 additions and 28 deletions.
diff --git a/crates/red_knot/src/lint.rs b/crates/red_knot/src/lint.rs
@@ -307,7 +307,7 @@ enum AnyImportRef<'a> {
 #[cfg(test)]
 mod tests {
     use ruff_db::files::system_path_to_file;
-    use ruff_db::program::{Program, SearchPathSettings, TargetVersion};
+    use ruff_db::program::{Program, SearchPathSettings, SitePackages, TargetVersion};
     use ruff_db::system::{DbWithTestSystem, SystemPathBuf};
 
     use super::{lint_semantic, Diagnostics};
@@ -326,7 +326,7 @@ mod tests {
             SearchPathSettings {
                 extra_paths: Vec::new(),
                 workspace_root,
-                site_packages: vec![],
+                site_packages: SitePackages::none(),
                 custom_typeshed: None,
             },
         );

diff --git a/crates/red_knot/src/main.rs b/crates/red_knot/src/main.rs
@@ -13,7 +13,7 @@ use red_knot::db::RootDatabase;
 use red_knot::watch;
 use red_knot::watch::WorkspaceWatcher;
 use red_knot::workspace::WorkspaceMetadata;
-use ruff_db::program::{ProgramSettings, SearchPathSettings};
+use ruff_db::program::{ProgramSettings, SearchPathSettings, SitePackages};
 use ruff_db::system::{OsSystem, System, SystemPathBuf};
 
 use cli::target_version::TargetVersion;
@@ -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: vec![],
+            site_packages: SitePackages::none(),
         },
     };
 

diff --git a/crates/red_knot/tests/check.rs b/crates/red_knot/tests/check.rs
@@ -2,7 +2,7 @@ use red_knot::db::RootDatabase;
 use red_knot::lint::lint_semantic;
 use red_knot::workspace::WorkspaceMetadata;
 use ruff_db::files::system_path_to_file;
-use ruff_db::program::{ProgramSettings, SearchPathSettings, TargetVersion};
+use ruff_db::program::{ProgramSettings, SearchPathSettings, SitePackages, TargetVersion};
 use ruff_db::system::{OsSystem, SystemPathBuf};
 use std::fs;
 use std::path::PathBuf;
@@ -14,7 +14,7 @@ fn setup_db(workspace_root: SystemPathBuf) -> anyhow::Result<RootDatabase> {
         extra_paths: vec![],
         workspace_root,
         custom_typeshed: None,
-        site_packages: vec![],
+        site_packages: SitePackages::none(),
     };
     let settings = ProgramSettings {
         target_version: TargetVersion::default(),

diff --git a/crates/red_knot/tests/file_watching.rs b/crates/red_knot/tests/file_watching.rs
@@ -12,7 +12,7 @@ use red_knot::watch::{directory_watcher, WorkspaceWatcher};
 use red_knot::workspace::WorkspaceMetadata;
 use red_knot_module_resolver::{resolve_module, ModuleName};
 use ruff_db::files::{system_path_to_file, File, FileError};
-use ruff_db::program::{Program, ProgramSettings, SearchPathSettings, TargetVersion};
+use ruff_db::program::{Program, ProgramSettings, SearchPathSettings, SitePackages, TargetVersion};
 use ruff_db::source::source_text;
 use ruff_db::system::{OsSystem, SystemPath, SystemPathBuf};
 use ruff_db::Upcast;
@@ -98,7 +98,7 @@ where
             extra_paths: vec![],
             workspace_root: workspace_path.to_path_buf(),
             custom_typeshed: None,
-            site_packages: vec![],
+            site_packages: SitePackages::none(),
         }
     })
 }
@@ -149,11 +149,14 @@ where
 
     let workspace = WorkspaceMetadata::from_path(&workspace_path, &system)?;
     let search_paths = create_search_paths(&root_path, workspace.root());
+    let SitePackages::Known(site_packages_dirs) = &search_paths.site_packages else {
+        panic!("Expected site-packages directories to be mocked out in file-watching tests")
+    };
 
     for path in search_paths
         .extra_paths
         .iter()
-        .chain(search_paths.site_packages.iter())
+        .chain(site_packages_dirs)
         .chain(search_paths.custom_typeshed.iter())
     {
         std::fs::create_dir_all(path.as_std_path())
@@ -656,7 +659,7 @@ fn search_path() -> anyhow::Result<()> {
                 extra_paths: vec![],
                 workspace_root: workspace_path.to_path_buf(),
                 custom_typeshed: None,
-                site_packages: vec![root_path.join("site_packages")],
+                site_packages: SitePackages::Known(vec![root_path.join("site_packages")]),
             }
         })?;
 
@@ -693,7 +696,7 @@ fn add_search_path() -> anyhow::Result<()> {
 
     // Register site-packages as a search path.
     case.update_search_path_settings(|settings| SearchPathSettings {
-        site_packages: vec![site_packages.clone()],
+        site_packages: SitePackages::Known(vec![site_packages.clone()]),
         ..settings.clone()
     });
 
@@ -716,14 +719,14 @@ fn remove_search_path() -> anyhow::Result<()> {
                 extra_paths: vec![],
                 workspace_root: workspace_path.to_path_buf(),
                 custom_typeshed: None,
-                site_packages: vec![root_path.join("site_packages")],
+                site_packages: SitePackages::Known(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: vec![],
+        site_packages: SitePackages::none(),
         ..settings.clone()
     });
 

diff --git a/crates/red_knot_module_resolver/src/resolver.rs b/crates/red_knot_module_resolver/src/resolver.rs
@@ -5,7 +5,7 @@ use once_cell::sync::Lazy;
 use rustc_hash::{FxBuildHasher, FxHashSet};
 
 use ruff_db::files::{File, FilePath, FileRootKind};
-use ruff_db::program::{Program, SearchPathSettings, TargetVersion};
+use ruff_db::program::{Program, SearchPathSettings, SitePackages, TargetVersion};
 use ruff_db::system::{DirectoryEntry, System, SystemPath, SystemPathBuf};
 use ruff_db::vendored::VendoredPath;
 
@@ -161,7 +161,14 @@ fn try_resolve_module_resolution_settings(
         SearchPath::vendored_stdlib()
     });
 
-    for site_packages_dir in site_packages {
+    let site_packages_dirs = match site_packages {
+        SitePackages::Derived { venv_path: _ } => {
+            todo!("Derive site-packages from a Python executable")
+        }
+        SitePackages::Known(site_packages_dirs) => site_packages_dirs,
+    };
+
+    for site_packages_dir in site_packages_dirs {
         files.try_add_root(
             db.upcast(),
             site_packages_dir,
@@ -1187,7 +1194,7 @@ mod tests {
             extra_paths: vec![],
             workspace_root: src.clone(),
             custom_typeshed: Some(custom_typeshed.clone()),
-            site_packages: vec![site_packages],
+            site_packages: SitePackages::Known(vec![site_packages]),
         };
 
         Program::new(&db, TargetVersion::Py38, search_paths);

diff --git a/crates/red_knot_module_resolver/src/testing.rs b/crates/red_knot_module_resolver/src/testing.rs
@@ -1,4 +1,4 @@
-use ruff_db::program::{Program, SearchPathSettings, TargetVersion};
+use ruff_db::program::{Program, SearchPathSettings, SitePackages, TargetVersion};
 use ruff_db::system::{DbWithTestSystem, SystemPath, SystemPathBuf};
 use ruff_db::vendored::VendoredPathBuf;
 
@@ -226,7 +226,7 @@ impl TestCaseBuilder<MockedTypeshed> {
                 extra_paths: vec![],
                 workspace_root: src.clone(),
                 custom_typeshed: Some(typeshed.clone()),
-                site_packages: vec![site_packages.clone()],
+                site_packages: SitePackages::Known(vec![site_packages.clone()]),
             },
         );
 
@@ -279,7 +279,7 @@ impl TestCaseBuilder<VendoredTypeshed> {
                 extra_paths: vec![],
                 workspace_root: src.clone(),
                 custom_typeshed: None,
-                site_packages: vec![site_packages.clone()],
+                site_packages: SitePackages::Known(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
@@ -164,7 +164,7 @@ impl HasTy for ast::Alias {
 mod tests {
     use ruff_db::files::system_path_to_file;
     use ruff_db::parsed::parsed_module;
-    use ruff_db::program::{Program, SearchPathSettings, TargetVersion};
+    use ruff_db::program::{Program, SearchPathSettings, SitePackages, TargetVersion};
     use ruff_db::system::{DbWithTestSystem, SystemPathBuf};
 
     use crate::db::tests::TestDb;
@@ -179,7 +179,7 @@ mod tests {
             SearchPathSettings {
                 extra_paths: vec![],
                 workspace_root: SystemPathBuf::from("/src"),
-                site_packages: vec![],
+                site_packages: SitePackages::none(),
                 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
@@ -1493,7 +1493,7 @@ impl<'db> TypeInferenceBuilder<'db> {
 mod tests {
     use ruff_db::files::{system_path_to_file, File};
     use ruff_db::parsed::parsed_module;
-    use ruff_db::program::{Program, SearchPathSettings, TargetVersion};
+    use ruff_db::program::{Program, SearchPathSettings, SitePackages, TargetVersion};
     use ruff_db::system::{DbWithTestSystem, SystemPathBuf};
     use ruff_db::testing::assert_function_query_was_not_run;
     use ruff_python_ast::name::Name;
@@ -1515,7 +1515,7 @@ mod tests {
             SearchPathSettings {
                 extra_paths: Vec::new(),
                 workspace_root: SystemPathBuf::from("/src"),
-                site_packages: vec![],
+                site_packages: SitePackages::none(),
                 custom_typeshed: None,
             },
         );
@@ -1532,7 +1532,7 @@ mod tests {
             SearchPathSettings {
                 extra_paths: Vec::new(),
                 workspace_root: SystemPathBuf::from("/src"),
-                site_packages: vec![],
+                site_packages: SitePackages::none(),
                 custom_typeshed: Some(SystemPathBuf::from(typeshed)),
             },
         );

diff --git a/crates/ruff_benchmark/benches/red_knot.rs b/crates/ruff_benchmark/benches/red_knot.rs
@@ -6,7 +6,7 @@ use red_knot::db::RootDatabase;
 use red_knot::workspace::WorkspaceMetadata;
 use ruff_db::files::{system_path_to_file, vendored_path_to_file, File};
 use ruff_db::parsed::parsed_module;
-use ruff_db::program::{ProgramSettings, SearchPathSettings, TargetVersion};
+use ruff_db::program::{ProgramSettings, SearchPathSettings, SitePackages, TargetVersion};
 use ruff_db::system::{MemoryFileSystem, SystemPath, TestSystem};
 use ruff_db::vendored::VendoredPath;
 use ruff_db::Upcast;
@@ -74,7 +74,7 @@ fn setup_case() -> Case {
         search_paths: SearchPathSettings {
             extra_paths: vec![],
             workspace_root: workspace_root.to_path_buf(),
-            site_packages: vec![],
+            site_packages: SitePackages::none(),
             custom_typeshed: None,
         },
     };

diff --git a/crates/ruff_db/src/program.rs b/crates/ruff_db/src/program.rs
@@ -64,6 +64,24 @@ impl std::fmt::Debug for TargetVersion {
     }
 }
 
+/// Configuration options for `site-packages` search paths.
+#[derive(Debug, Clone, PartialEq, Eq)]
+pub enum SitePackages {
+    /// Derive the `site-packages` paths from a path to a Python executable
+    Derived { venv_path: SystemPathBuf },
+
+    /// Pretend that all `site-packages` paths are known before the red-knot process starts.
+    ///
+    /// This isn't particularly realistic, but it's useful for mocking out site-packages in tests.
+    Known(Vec<SystemPathBuf>),
+}
+
+impl SitePackages {
+    pub fn none() -> Self {
+        Self::Known(vec![])
+    }
+}
+
 /// Configures the search paths for module resolution.
 #[derive(Eq, PartialEq, Debug, Clone)]
 pub struct SearchPathSettings {
@@ -80,6 +98,12 @@ pub struct SearchPathSettings {
     /// bundled as a zip file in the binary
     pub custom_typeshed: Option<SystemPathBuf>,
 
-    /// The path to the user's `site-packages` directory, where third-party packages from ``PyPI`` are installed.
-    pub site_packages: Vec<SystemPathBuf>,
+    /// The paths to the user's `site-packages` directories,
+    /// where third-party packages from ``PyPI`` are installed.
+    ///
+    /// Usually there will be exactly 1 `site-packages` directory.
+    /// However, there may be 0 (if the user runs red-knot with `--isolated`),
+    /// or many (if the user has activate a virtual environment with
+    /// `include-system-site-packages = true` in its `pyvenv.cfg` file)
+    pub site_packages: SitePackages,
 }