Move venv resolution into SearchPaths::from_settings

crates/red_knot/src/main.rs

@@ -7,14 +7,12 @@ use colored::Colorize;
 use crossbeam::channel as crossbeam_channel;
 use salsa::plumbing::ZalsaDatabase;
 
-use red_knot_python_semantic::PythonVersion;
+use red_knot_python_semantic::SitePackages;
 use red_knot_server::run_server;
 use red_knot_workspace::db::RootDatabase;
 use red_knot_workspace::watch;
 use red_knot_workspace::watch::WorkspaceWatcher;
-use red_knot_workspace::workspace::settings::{
-    SitePackages, WorkspaceConfiguration, WorkspaceConfigurationTransformer,
-};
+use red_knot_workspace::workspace::settings::Configuration;
 use red_knot_workspace::workspace::WorkspaceMetadata;
 use ruff_db::system::{OsSystem, System, SystemPath, SystemPathBuf};
 use target_version::TargetVersion;
@@ -87,6 +85,36 @@ to resolve type information for the project's third-party dependencies.",
     watch: bool,
 }
 
+impl Args {
+    fn to_configuration(&self, cli_cwd: &SystemPath) -> Configuration {
+        let mut configuration = Configuration::default();
+
+        if let Some(target_version) = self.target_version {
+            configuration.target_version = Some(target_version.into());
+        }
+
+        if let Some(venv_path) = &self.venv_path {
+            configuration.search_paths.site_packages = Some(SitePackages::Derived {
+                venv_path: SystemPath::absolute(venv_path, cli_cwd),
+            });
+        }
+
+        if let Some(custom_typeshed_dir) = &self.custom_typeshed_dir {
+            configuration.search_paths.custom_typeshed =
+                Some(SystemPath::absolute(custom_typeshed_dir, cli_cwd));
+        }
+
+        if let Some(extra_search_paths) = &self.extra_search_path {
+            configuration.search_paths.extra_paths = extra_search_paths
+                .iter()
+                .map(|path| Some(SystemPath::absolute(path, cli_cwd)))
+                .collect();
+        }
+
+        configuration
+    }
+}
+
 #[derive(Debug, clap::Subcommand)]
 pub enum Command {
     /// Start the language server
@@ -154,15 +182,18 @@ fn run() -> anyhow::Result<ExitStatus> {
         .unwrap_or_else(|| cli_base_path.clone());
 
     let system = OsSystem::new(cwd.clone());
-    let transformer = CliConfigurationTransformer::from_cli_arguments(&args, &cli_base_path);
-    let workspace_metadata =
-        WorkspaceMetadata::from_path(system.current_directory(), &system, &transformer)?;
+    let cli_configuration = args.to_configuration(&cwd);
+    let workspace_metadata = WorkspaceMetadata::from_path(
+        system.current_directory(),
+        &system,
+        Some(cli_configuration.clone()),
+    )?;
 
     // TODO: Use the `program_settings` to compute the key for the database's persistent
     //   cache and load the cache if it exists.
     let mut db = RootDatabase::new(workspace_metadata, system)?;
 
-    let (main_loop, main_loop_cancellation_token) = MainLoop::new(transformer);
+    let (main_loop, main_loop_cancellation_token) = MainLoop::new(cli_configuration);
 
     // Listen to Ctrl+C and abort the watch mode.
     let main_loop_cancellation_token = Mutex::new(Some(main_loop_cancellation_token));
@@ -215,21 +246,19 @@ struct MainLoop {
     /// The file system watcher, if running in watch mode.
     watcher: Option<WorkspaceWatcher>,
 
-    configuration_transformer: CliConfigurationTransformer,
+    cli_configuration: Configuration,
 }
 
 impl MainLoop {
-    fn new(
-        configuration_transformer: CliConfigurationTransformer,
-    ) -> (Self, MainLoopCancellationToken) {
+    fn new(cli_configuration: Configuration) -> (Self, MainLoopCancellationToken) {
         let (sender, receiver) = crossbeam_channel::bounded(10);
 
         (
             Self {
                 sender: sender.clone(),
                 receiver,
                 watcher: None,
-                configuration_transformer,
+                cli_configuration,
             },
             MainLoopCancellationToken { sender },
         )
@@ -312,7 +341,7 @@ impl MainLoop {
                 MainLoopMessage::ApplyChanges(changes) => {
                     revision += 1;
                     // Automatically cancels any pending queries and waits for them to complete.
-                    db.apply_changes(changes, &self.configuration_transformer);
+                    db.apply_changes(changes, Some(&self.cli_configuration));
                     if let Some(watcher) = self.watcher.as_mut() {
                         watcher.update(db);
                     }
@@ -353,69 +382,3 @@ enum MainLoopMessage {
     ApplyChanges(Vec<watch::ChangeEvent>),
     Exit,
 }
-
-#[derive(Debug, Default)]
-struct CliConfigurationTransformer {
-    venv_path: Option<SystemPathBuf>,
-    custom_typeshed_dir: Option<SystemPathBuf>,
-    extra_search_paths: Option<Vec<SystemPathBuf>>,
-    target_version: Option<PythonVersion>,
-}
-
-impl CliConfigurationTransformer {
-    fn from_cli_arguments(arguments: &Args, cli_cwd: &SystemPath) -> Self {
-        let Args {
-            venv_path,
-            custom_typeshed_dir,
-            extra_search_path,
-            target_version,
-            ..
-        } = arguments;
-
-        let venv_path = venv_path
-            .as_deref()
-            .map(|path| SystemPath::absolute(path, cli_cwd));
-
-        let custom_typeshed_dir = custom_typeshed_dir
-            .as_deref()
-            .map(|path| SystemPath::absolute(path, cli_cwd));
-
-        let extra_search_paths = extra_search_path.as_deref().map(|paths| {
-            paths
-                .iter()
-                .map(|path| SystemPath::absolute(path, cli_cwd))
-                .collect()
-        });
-
-        Self {
-            venv_path,
-            custom_typeshed_dir,
-            extra_search_paths,
-            target_version: target_version.map(PythonVersion::from),
-        }
-    }
-}
-
-impl WorkspaceConfigurationTransformer for CliConfigurationTransformer {
-    fn transform(&self, mut configuration: WorkspaceConfiguration) -> WorkspaceConfiguration {
-        if let Some(venv_path) = &self.venv_path {
-            configuration.search_paths.site_packages = Some(SitePackages::Derived {
-                venv_path: venv_path.clone(),
-            });
-        }
-
-        if let Some(custom_typeshed_dir) = &self.custom_typeshed_dir {
-            configuration.search_paths.custom_typeshed = Some(custom_typeshed_dir.clone());
-        }
-
-        if let Some(extra_search_paths) = &self.extra_search_paths {
-            configuration.search_paths.extra_paths = Some(extra_search_paths.clone());
-        }
-
-        if let Some(target_version) = self.target_version {
-            configuration.target_version = Some(target_version);
-        }
-
-        configuration
-    }
-}

crates/red_knot/tests/file_watching.rs

@@ -5,18 +5,16 @@ use std::time::Duration;
 
 use anyhow::{anyhow, Context};
 
-use red_knot_python_semantic::{resolve_module, ModuleName, Program, PythonVersion};
+use red_knot_python_semantic::{resolve_module, ModuleName, Program, PythonVersion, SitePackages};
 use red_knot_workspace::db::RootDatabase;
 use red_knot_workspace::watch;
 use red_knot_workspace::watch::{directory_watcher, WorkspaceWatcher};
-use red_knot_workspace::workspace::settings::{
-    SearchPathConfiguration, SitePackages, WorkspaceConfiguration,
-};
+use red_knot_workspace::workspace::settings::{Configuration, SearchPathConfiguration};
 use red_knot_workspace::workspace::WorkspaceMetadata;
 use ruff_db::files::{system_path_to_file, File, FileError};
 use ruff_db::source::source_text;
 use ruff_db::system::{OsSystem, SystemPath, SystemPathBuf};
-use ruff_db::{Db, Upcast};
+use ruff_db::Upcast;
 
 struct TestCase {
     db: RootDatabase,
@@ -26,7 +24,7 @@ struct TestCase {
     /// We need to hold on to it in the test case or the temp files get deleted.
     _temp_dir: tempfile::TempDir,
     root_dir: SystemPathBuf,
-    search_path_configuration: SearchPathConfiguration,
+    configuration: Configuration,
 }
 
 impl TestCase {
@@ -103,12 +101,7 @@ impl TestCase {
     }
 
     fn apply_changes(&mut self, changes: Vec<watch::ChangeEvent>) {
-        self.db
-            .apply_changes(changes, &|mut configuration: WorkspaceConfiguration| {
-                configuration.target_version = Some(PythonVersion::PY312);
-                configuration.search_paths = self.search_path_configuration.clone();
-                configuration
-            });
+        self.db.apply_changes(changes, Some(&self.configuration));
     }
 
     fn update_search_path_settings(
@@ -117,11 +110,10 @@ impl TestCase {
     ) -> anyhow::Result<()> {
         let program = Program::get(self.db());
 
-        let new_settings =
-            configuration.to_settings(self.db.workspace().root(&self.db), self.db.system())?;
+        let new_settings = configuration.to_settings(self.db.workspace().root(&self.db));
+        self.configuration.search_paths = configuration;
 
         program.update_search_paths(&mut self.db, new_settings)?;
-        self.search_path_configuration = configuration;
 
         if let Some(watcher) = &mut self.watcher {
             watcher.update(&self.db);
@@ -226,40 +218,33 @@ where
 
     let system = OsSystem::new(&workspace_path);
 
-    let search_path_configuration = create_search_paths(&root_path, &workspace_path);
+    let search_paths = create_search_paths(&root_path, &workspace_path);
 
-    let workspace = WorkspaceMetadata::from_path(
-        &workspace_path,
-        &system,
-        &|mut configuration: WorkspaceConfiguration| {
-            configuration.target_version = Some(PythonVersion::PY312);
-            configuration.search_paths = search_path_configuration.clone();
-            configuration
-        },
-    )?;
-
-    for path in search_path_configuration
+    for path in search_paths
         .extra_paths
         .iter()
         .flatten()
-        .chain(search_path_configuration.custom_typeshed.iter())
-        .chain(
-            search_path_configuration
-                .site_packages
-                .as_ref()
-                .and_then(|site_packages| {
-                    if let SitePackages::Known(path) = site_packages {
-                        Some(path)
-                    } else {
-                        None
-                    }
-                }),
-        )
+        .chain(search_paths.custom_typeshed.iter())
+        .chain(search_paths.site_packages.iter().flat_map(|site_packages| {
+            if let SitePackages::Known(path) = site_packages {
+                path.as_slice()
+            } else {
+                &[]
+            }
+        }))
     {
         std::fs::create_dir_all(path.as_std_path())
             .with_context(|| format!("Failed to create search path '{path}'"))?;
     }
 
+    let configuration = Configuration {
+        target_version: Some(PythonVersion::PY312),
+        search_paths,
+    };
+
+    let workspace =
+        WorkspaceMetadata::from_path(&workspace_path, &system, Some(configuration.clone()))?;
+
     let db = RootDatabase::new(workspace, system)?;
 
     let (sender, receiver) = crossbeam::channel::unbounded();
@@ -275,7 +260,7 @@ where
         watcher: Some(watcher),
         _temp_dir: temp_dir,
         root_dir: root_path,
-        search_path_configuration,
+        configuration,
     };
 
     // Sometimes the file watcher reports changes for events that happened before the watcher was started.
@@ -718,7 +703,7 @@ fn search_path() -> anyhow::Result<()> {
     let mut case = setup_with_search_paths(
         [("bar.py", "import sub.a")],
         |root_path, _workspace_path| SearchPathConfiguration {
-            site_packages: Some(SitePackages::Known(root_path.join("site_packages"))),
+            site_packages: Some(SitePackages::Known(vec![root_path.join("site_packages")])),
             ..SearchPathConfiguration::default()
         },
     )?;
@@ -756,7 +741,7 @@ fn add_search_path() -> anyhow::Result<()> {
 
     // Register site-packages as a search path.
     case.update_search_path_settings(SearchPathConfiguration {
-        site_packages: Some(SitePackages::Known(site_packages.clone())),
+        site_packages: Some(SitePackages::Known(vec![site_packages.clone()])),
         ..SearchPathConfiguration::default()
     })
     .expect("Search path settings to be valid");
@@ -777,7 +762,7 @@ fn remove_search_path() -> anyhow::Result<()> {
     let mut case = setup_with_search_paths(
         [("bar.py", "import sub.a")],
         |root_path, _workspace_path| SearchPathConfiguration {
-            site_packages: Some(SitePackages::Known(root_path.join("site_packages"))),
+            site_packages: Some(SitePackages::Known(vec![root_path.join("site_packages")])),
             ..SearchPathConfiguration::default()
         },
     )?;
@@ -1236,9 +1221,9 @@ mod unix {
                 Ok(())
             },
             |_root, workspace| SearchPathConfiguration {
-                site_packages: Some(SitePackages::Known(
-                    workspace.join(".venv/lib/python3.12/site-packages"),
-                )),
+                site_packages: Some(SitePackages::Known(vec![
+                    workspace.join(".venv/lib/python3.12/site-packages")
+                ])),
                 ..SearchPathConfiguration::default()
             },
         )?;

crates/red_knot_python_semantic/Cargo.toml

@@ -26,6 +26,7 @@ countme = { workspace = true }
 once_cell = { workspace = true }
 ordermap = { workspace = true }
 salsa = { workspace = true }
+thiserror = { workspace = true }
 tracing = { workspace = true }
 rustc-hash = { workspace = true }
 hashbrown = { workspace = true }

crates/red_knot_python_semantic/src/lib.rs

@@ -5,7 +5,7 @@ use rustc_hash::FxHasher;
 pub use db::Db;
 pub use module_name::ModuleName;
 pub use module_resolver::{resolve_module, system_module_search_paths, vendored_typeshed_stubs};
-pub use program::{Program, ProgramSettings, SearchPathSettings};
+pub use program::{Program, ProgramSettings, SearchPathSettings, SitePackages};
 pub use python_version::PythonVersion;
 pub use semantic_model::{HasTy, SemanticModel};
 
@@ -19,6 +19,7 @@ mod program;
 mod python_version;
 pub mod semantic_index;
 mod semantic_model;
+pub(crate) mod site_packages;
 pub mod types;
 
 type FxOrderSet<V> = ordermap::set::OrderSet<V, BuildHasherDefault<FxHasher>>;