From 2fa67eae6fe3200648bfb1d17913c2094e01f6a8 Mon Sep 17 00:00:00 2001 From: samypr100 <3933065+samypr100@users.noreply.github.com> Date: Thu, 22 Feb 2024 14:39:37 -0500 Subject: [PATCH] feat: allow passing extra config k,v pairs for pyvenv.cfg when creating a venv (#1852) ## Summary This modifies `gourgeist` to allow passing additional k,v pairs to add to the `pyvenv.cfg` file as proposed in #1697. I made it allow an arbitrary set of pairs (to decouple from `uv` since this is mainly a change to `gourgeist`) , but I can slim it down to just allow just a name and version strings if that's desired. The `pyvenv.cfg` will also have a `uv = ` when a venv is created via `uv venv` ~~and `uv-build = ` when it's created via `SourceBuild::setup`~~. Example below via `uv venv`: ```ini home = ... implementation = CPython version_info = 3.12 include-system-site-packages = false base-prefix = ... base-exec-prefix = ... base-executable = ... uv = 0.1.6 prompt = uv ``` Open to any suggestions, thanks! Closes #1697 ## Test Plan Added new test in `tests/venv.rs` called `verify_pyvenv_cfg` to verify that it contains the right uv version string. I didn't see tests configured in `gourgeist` itself, so I didn't add any there. --- crates/gourgeist/src/bare.rs | 75 ++++++++++++++++++++++------------ crates/gourgeist/src/lib.rs | 5 ++- crates/gourgeist/src/main.rs | 2 +- crates/uv-build/src/lib.rs | 1 + crates/uv/src/commands/venv.rs | 6 ++- crates/uv/tests/venv.rs | 21 +++++++++- 6 files changed, 79 insertions(+), 31 deletions(-) diff --git a/crates/gourgeist/src/bare.rs b/crates/gourgeist/src/bare.rs index aff617b70c27..da9a658377c4 100644 --- a/crates/gourgeist/src/bare.rs +++ b/crates/gourgeist/src/bare.rs @@ -13,7 +13,7 @@ use uv_fs::Normalized; use uv_interpreter::Interpreter; -use crate::Prompt; +use crate::{Error, Prompt}; /// The bash activate scripts with the venv dependent paths patches out const ACTIVATE_TEMPLATES: &[(&str, &str)] = &[ @@ -33,17 +33,10 @@ const ACTIVATE_TEMPLATES: &[(&str, &str)] = &[ const VIRTUALENV_PATCH: &str = include_str!("_virtualenv.py"); /// Very basic `.cfg` file format writer. -fn write_cfg( - f: &mut impl Write, - data: &[(&str, String); 8], - prompt: Option, -) -> io::Result<()> { +fn write_cfg(f: &mut impl Write, data: &[(String, String)]) -> io::Result<()> { for (key, value) in data { writeln!(f, "{key} = {value}")?; } - if let Some(prompt) = prompt { - writeln!(f, "prompt = {prompt}")?; - } Ok(()) } @@ -73,7 +66,8 @@ pub fn create_bare_venv( location: &Utf8Path, interpreter: &Interpreter, prompt: Prompt, -) -> io::Result { + extra_cfg: Vec<(String, String)>, +) -> Result { // We have to canonicalize the interpreter path, otherwise the home is set to the venv dir instead of the real root. // This would make python-build-standalone fail with the encodings module not being found because its home is wrong. let base_python: Utf8PathBuf = fs_err::canonicalize(interpreter.sys_executable())? @@ -84,10 +78,10 @@ pub fn create_bare_venv( match location.metadata() { Ok(metadata) => { if metadata.is_file() { - return Err(io::Error::new( + return Err(Error::IO(io::Error::new( io::ErrorKind::AlreadyExists, format!("File exists at `{location}`"), - )); + ))); } else if metadata.is_dir() { if location.join("pyvenv.cfg").is_file() { info!("Removing existing directory"); @@ -99,17 +93,17 @@ pub fn create_bare_venv( { info!("Ignoring empty directory"); } else { - return Err(io::Error::new( + return Err(Error::IO(io::Error::new( io::ErrorKind::AlreadyExists, format!("The directory `{location}` exists, but it's not a virtualenv"), - )); + ))); } } } Err(err) if err.kind() == io::ErrorKind::NotFound => { fs::create_dir_all(location)?; } - Err(err) => return Err(err), + Err(err) => return Err(Error::IO(err)), } // TODO(konstin): I bet on windows we'll have to strip the prefix again @@ -226,31 +220,58 @@ pub fn create_bare_venv( } else { unimplemented!("Only Windows and Unix are supported") }; - let pyvenv_cfg_data = &[ - ("home", python_home), + + // Validate extra_cfg + let reserved_keys = [ + "home", + "implementation", + "version_info", + "include-system-site-packages", + "base-prefix", + "base-exec-prefix", + "base-executable", + "prompt", + ]; + for (key, _) in &extra_cfg { + if reserved_keys.contains(&key.as_str()) { + return Err(Error::ReservedConfigKey(key.to_string())); + } + } + + let mut pyvenv_cfg_data: Vec<(String, String)> = vec![ + ("home".to_string(), python_home), ( - "implementation", + "implementation".to_string(), interpreter.markers().platform_python_implementation.clone(), ), ( - "version_info", + "version_info".to_string(), interpreter.markers().python_version.string.clone(), ), - ("gourgeist", env!("CARGO_PKG_VERSION").to_string()), - // I wouldn't allow this option anyway - ("include-system-site-packages", "false".to_string()), ( - "base-prefix", + "include-system-site-packages".to_string(), + "false".to_string(), + ), + ( + "base-prefix".to_string(), interpreter.base_prefix().to_string_lossy().to_string(), ), ( - "base-exec-prefix", + "base-exec-prefix".to_string(), interpreter.base_exec_prefix().to_string_lossy().to_string(), ), - ("base-executable", base_python.to_string()), - ]; + ("base-executable".to_string(), base_python.to_string()), + ] + .into_iter() + .chain(extra_cfg) + .collect(); + + if let Some(prompt) = prompt { + pyvenv_cfg_data.push(("prompt".to_string(), prompt)); + } + let mut pyvenv_cfg = BufWriter::new(File::create(location.join("pyvenv.cfg"))?); - write_cfg(&mut pyvenv_cfg, pyvenv_cfg_data, prompt)?; + write_cfg(&mut pyvenv_cfg, &pyvenv_cfg_data)?; drop(pyvenv_cfg); let site_packages = if cfg!(unix) { diff --git a/crates/gourgeist/src/lib.rs b/crates/gourgeist/src/lib.rs index 9545253023ff..c4c2ec83c0e1 100644 --- a/crates/gourgeist/src/lib.rs +++ b/crates/gourgeist/src/lib.rs @@ -21,6 +21,8 @@ pub enum Error { InvalidPythonInterpreter(#[source] Box), #[error(transparent)] Platform(#[from] PlatformError), + #[error("Reserved key used for pyvenv.cfg: {0}")] + ReservedConfigKey(String), } /// The value to use for the shell prompt when inside a virtual environment. @@ -51,11 +53,12 @@ pub fn create_venv( location: &Path, interpreter: Interpreter, prompt: Prompt, + extra_cfg: Vec<(String, String)>, ) -> Result { let location: &Utf8Path = location .try_into() .map_err(|err: FromPathError| err.into_io_error())?; - let paths = create_bare_venv(location, &interpreter, prompt)?; + let paths = create_bare_venv(location, &interpreter, prompt, extra_cfg)?; Ok(Virtualenv::from_interpreter( interpreter, paths.root.as_std_path(), diff --git a/crates/gourgeist/src/main.rs b/crates/gourgeist/src/main.rs index a6718fef24c2..13691a8985ec 100644 --- a/crates/gourgeist/src/main.rs +++ b/crates/gourgeist/src/main.rs @@ -36,7 +36,7 @@ fn run() -> Result<(), gourgeist::Error> { Cache::from_path(".gourgeist_cache")? }; let info = Interpreter::query(python.as_std_path(), &platform, &cache).unwrap(); - create_bare_venv(&location, &info, Prompt::from_args(cli.prompt))?; + create_bare_venv(&location, &info, Prompt::from_args(cli.prompt), Vec::new())?; Ok(()) } diff --git a/crates/uv-build/src/lib.rs b/crates/uv-build/src/lib.rs index 2708b194a5fb..d63ba89f6749 100644 --- a/crates/uv-build/src/lib.rs +++ b/crates/uv-build/src/lib.rs @@ -329,6 +329,7 @@ impl SourceBuild { &temp_dir.path().join(".venv"), interpreter.clone(), gourgeist::Prompt::None, + Vec::new(), )?; // Setup the build environment. diff --git a/crates/uv/src/commands/venv.rs b/crates/uv/src/commands/venv.rs index 72cbfc800463..ba38edd36382 100644 --- a/crates/uv/src/commands/venv.rs +++ b/crates/uv/src/commands/venv.rs @@ -119,8 +119,12 @@ async fn venv_impl( ) .into_diagnostic()?; + // Extra cfg for pyvenv.cfg to specify uv version + let extra_cfg = vec![("uv".to_string(), env!("CARGO_PKG_VERSION").to_string())]; + // Create the virtual environment. - let venv = gourgeist::create_venv(path, interpreter, prompt).map_err(VenvError::Creation)?; + let venv = gourgeist::create_venv(path, interpreter, prompt, extra_cfg) + .map_err(VenvError::Creation)?; // Install seed packages. if seed { diff --git a/crates/uv/tests/venv.rs b/crates/uv/tests/venv.rs index 81277e7df4af..587c1126769d 100644 --- a/crates/uv/tests/venv.rs +++ b/crates/uv/tests/venv.rs @@ -7,7 +7,9 @@ use assert_fs::prelude::*; use uv_fs::Normalized; -use crate::common::{create_bin_with_executables, get_bin, uv_snapshot, EXCLUDE_NEWER}; +use crate::common::{ + create_bin_with_executables, get_bin, uv_snapshot, TestContext, EXCLUDE_NEWER, +}; mod common; @@ -644,3 +646,20 @@ fn virtualenv_compatibility() -> Result<()> { Ok(()) } + +#[test] +fn verify_pyvenv_cfg() { + let context = TestContext::new("3.12"); + let venv = context.temp_dir.child(".venv"); + let pyvenv_cfg = venv.child("pyvenv.cfg"); + + venv.assert(predicates::path::is_dir()); + + // Check pyvenv.cfg exists + pyvenv_cfg.assert(predicates::path::is_file()); + + // Check if "uv = version" is present in the file + let version = env!("CARGO_PKG_VERSION").to_string(); + let search_string = format!("uv = {version}"); + pyvenv_cfg.assert(predicates::str::contains(search_string)); +}