From fbb57b24dd949d466ef5815d678505aba9599620 Mon Sep 17 00:00:00 2001 From: Charlie Marsh Date: Tue, 9 Jan 2024 20:45:56 -0500 Subject: [PATCH] Add `--seed` flag to `venv` to allow seed package environments (#865) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit ## Summary Installs the seed packages you get with `virtualenv`, but opt-in rather than opt-out. Closes https://github.com/astral-sh/puffin/issues/852. ## Test Plan ``` ❯ ./scripts/benchmarks/venv.sh + hyperfine --runs 20 --warmup 3 --prepare 'rm -rf .venv' './target/release/puffin venv' --prepare 'rm -rf .venv' 'virtualenv --without-pip .venv' --prepare 'rm -rf .venv' 'python -m venv --without-pip .venv' Benchmark 1: ./target/release/puffin venv Time (mean ± σ): 4.6 ms ± 0.2 ms [User: 2.4 ms, System: 3.6 ms] Range (min … max): 4.3 ms … 4.9 ms 20 runs Warning: Command took less than 5 ms to complete. Note that the results might be inaccurate because hyperfine can not calibrate the shell startup time much more precise than this limit. You can try to use the `-N`/`--shell=none` option to disable the shell completely. Benchmark 2: virtualenv --without-pip .venv Time (mean ± σ): 73.3 ms ± 0.3 ms [User: 57.4 ms, System: 14.2 ms] Range (min … max): 72.8 ms … 74.0 ms 20 runs Benchmark 3: python -m venv --without-pip .venv Time (mean ± σ): 22.5 ms ± 0.3 ms [User: 17.0 ms, System: 4.9 ms] Range (min … max): 22.0 ms … 23.2 ms 20 runs Summary './target/release/puffin venv' ran 4.92 ± 0.20 times faster than 'python -m venv --without-pip .venv' 16.00 ± 0.63 times faster than 'virtualenv --without-pip .venv' + hyperfine --runs 20 --warmup 3 --prepare 'rm -rf .venv' './target/release/puffin venv --seed' --prepare 'rm -rf .venv' 'virtualenv .venv' --prepare 'rm -rf .venv' 'python -m venv .venv' Benchmark 1: ./target/release/puffin venv --seed Time (mean ± σ): 20.2 ms ± 0.4 ms [User: 8.6 ms, System: 15.7 ms] Range (min … max): 19.7 ms … 21.2 ms 20 runs Benchmark 2: virtualenv .venv Time (mean ± σ): 135.1 ms ± 2.4 ms [User: 66.7 ms, System: 65.7 ms] Range (min … max): 133.2 ms … 142.8 ms 20 runs Benchmark 3: python -m venv .venv Time (mean ± σ): 1.656 s ± 0.014 s [User: 1.447 s, System: 0.186 s] Range (min … max): 1.641 s … 1.697 s 20 runs Summary './target/release/puffin venv --seed' ran 6.67 ± 0.17 times faster than 'virtualenv .venv' 81.79 ± 1.70 times faster than 'python -m venv .venv' ``` --- crates/puffin-cli/src/commands/venv.rs | 66 ++++++++++++++++++++++++-- crates/puffin-cli/src/main.rs | 30 +++++++++++- crates/puffin-cli/tests/venv.rs | 36 ++++++++++++++ scripts/benchmarks/venv.sh | 8 ++-- 4 files changed, 131 insertions(+), 9 deletions(-) diff --git a/crates/puffin-cli/src/commands/venv.rs b/crates/puffin-cli/src/commands/venv.rs index e3d18ecde..957d78cc4 100644 --- a/crates/puffin-cli/src/commands/venv.rs +++ b/crates/puffin-cli/src/commands/venv.rs @@ -1,5 +1,6 @@ use std::fmt::Write; use std::path::{Path, PathBuf}; +use std::str::FromStr; use anyhow::Result; use fs_err as fs; @@ -7,22 +8,29 @@ use miette::{Diagnostic, IntoDiagnostic}; use owo_colors::OwoColorize; use thiserror::Error; +use distribution_types::{DistributionMetadata, IndexUrls, Name}; +use pep508_rs::Requirement; use platform_host::Platform; use puffin_cache::Cache; +use puffin_client::RegistryClientBuilder; +use puffin_dispatch::BuildDispatch; use puffin_interpreter::Interpreter; +use puffin_traits::{BuildContext, SetupPyStrategy}; use crate::commands::ExitStatus; use crate::printer::Printer; /// Create a virtual environment. #[allow(clippy::unnecessary_wraps)] -pub(crate) fn venv( +pub(crate) async fn venv( path: &Path, base_python: Option<&Path>, + index_urls: &IndexUrls, + seed: bool, cache: &Cache, printer: Printer, ) -> Result { - match venv_impl(path, base_python, cache, printer) { + match venv_impl(path, base_python, index_urls, seed, cache, printer).await { Ok(status) => Ok(status), Err(err) => { #[allow(clippy::print_stderr)] @@ -51,12 +59,18 @@ enum VenvError { #[error("Failed to create virtual environment")] #[diagnostic(code(puffin::venv::creation))] CreationError(#[source] gourgeist::Error), + + #[error("Failed to install seed packages")] + #[diagnostic(code(puffin::venv::seed))] + SeedError(#[source] anyhow::Error), } /// Create a virtual environment. -fn venv_impl( +async fn venv_impl( path: &Path, base_python: Option<&Path>, + index_urls: &IndexUrls, + seed: bool, cache: &Cache, mut printer: Printer, ) -> miette::Result { @@ -96,7 +110,51 @@ fn venv_impl( .into_diagnostic()?; // Create the virtual environment. - gourgeist::create_venv(path, interpreter).map_err(VenvError::CreationError)?; + let venv = gourgeist::create_venv(path, interpreter).map_err(VenvError::CreationError)?; + + // Install seed packages. + if seed { + // Instantiate a client. + let client = RegistryClientBuilder::new(cache.clone()).build(); + + // Prep the build context. + let build_dispatch = BuildDispatch::new( + &client, + cache, + venv.interpreter(), + index_urls, + venv.python_executable(), + SetupPyStrategy::default(), + true, + ); + + // Resolve the seed packages. + let resolution = build_dispatch + .resolve(&[ + Requirement::from_str("wheel").unwrap(), + Requirement::from_str("pip").unwrap(), + Requirement::from_str("setuptools").unwrap(), + ]) + .await + .map_err(VenvError::SeedError)?; + + // Install into the environment. + build_dispatch + .install(&resolution, &venv) + .await + .map_err(VenvError::SeedError)?; + + for distribution in resolution.distributions() { + writeln!( + printer, + " {} {}{}", + "+".green(), + distribution.name().as_ref().white().bold(), + distribution.version_or_url().dimmed() + ) + .into_diagnostic()?; + } + } Ok(ExitStatus::Success) } diff --git a/crates/puffin-cli/src/main.rs b/crates/puffin-cli/src/main.rs index aeb9a1771..3abc587ee 100644 --- a/crates/puffin-cli/src/main.rs +++ b/crates/puffin-cli/src/main.rs @@ -408,9 +408,25 @@ struct VenvArgs { #[clap(short, long)] python: Option, + /// Install seed packages (`pip`, `setuptools`, and `wheel`) into the virtual environment. + #[clap(long)] + seed: bool, + /// The path to the virtual environment to create. #[clap(default_value = ".venv")] name: PathBuf, + + /// The URL of the Python Package Index. + #[clap(long, short, default_value = IndexUrl::Pypi.as_str(), env = "PUFFIN_INDEX_URL")] + index_url: IndexUrl, + + /// Extra URLs of package indexes to use, in addition to `--index-url`. + #[clap(long)] + extra_index_url: Vec, + + /// Ignore the package index, instead relying on local archives and caches. + #[clap(long, conflicts_with = "index_url", conflicts_with = "extra_index_url")] + no_index: bool, } #[derive(Args)] @@ -598,7 +614,19 @@ async fn inner() -> Result { } Commands::Clean(args) => commands::clean(&cache, &args.package, printer), Commands::PipFreeze(args) => commands::freeze(&cache, args.strict, printer), - Commands::Venv(args) => commands::venv(&args.name, args.python.as_deref(), &cache, printer), + Commands::Venv(args) => { + let index_urls = + IndexUrls::from_args(args.index_url, args.extra_index_url, args.no_index); + commands::venv( + &args.name, + args.python.as_deref(), + &index_urls, + args.seed, + &cache, + printer, + ) + .await + } Commands::Add(args) => commands::add(&args.name, printer), Commands::Remove(args) => commands::remove(&args.name, printer), } diff --git a/crates/puffin-cli/tests/venv.rs b/crates/puffin-cli/tests/venv.rs index 5da45bf6a..03d88833c 100644 --- a/crates/puffin-cli/tests/venv.rs +++ b/crates/puffin-cli/tests/venv.rs @@ -73,3 +73,39 @@ fn create_venv_defaults_to_cwd() -> Result<()> { Ok(()) } + +#[test] +fn seed() -> Result<()> { + let temp_dir = assert_fs::TempDir::new()?; + let venv = temp_dir.child(".venv"); + + insta::with_settings!({ + filters => vec![ + (r"Using Python 3\.\d+\.\d+ at .+", "Using Python [VERSION] at [PATH]"), + (temp_dir.to_str().unwrap(), "/home/ferris/project"), + ] + }, { + assert_cmd_snapshot!(Command::new(get_cargo_bin(BIN_NAME)) + .arg("venv") + .arg(venv.as_os_str()) + .arg("--seed") + .arg("--python") + .arg("python3.12") + .current_dir(&temp_dir), @r###" + success: true + exit_code: 0 + ----- stdout ----- + + ----- stderr ----- + Using Python [VERSION] at [PATH] + Creating virtual environment at: /home/ferris/project/.venv + + setuptools==69.0.3 + + pip==23.3.2 + + wheel==0.42.0 + "###); + }); + + venv.assert(predicates::path::is_dir()); + + Ok(()) +} diff --git a/scripts/benchmarks/venv.sh b/scripts/benchmarks/venv.sh index 8aa2c1c49..8ef0cb36a 100755 --- a/scripts/benchmarks/venv.sh +++ b/scripts/benchmarks/venv.sh @@ -5,7 +5,7 @@ # # Example usage: # -# ./scripts/benchmarks/venv.sh ./scripts/benchmarks/requirements.txt +# ./scripts/benchmarks/venv.sh ### set -euxo pipefail @@ -15,7 +15,7 @@ set -euxo pipefail ### hyperfine --runs 20 --warmup 3 \ --prepare "rm -rf .venv" \ - "./target/release/puffin venv --no-cache" \ + "./target/release/puffin venv" \ --prepare "rm -rf .venv" \ "virtualenv --without-pip .venv" \ --prepare "rm -rf .venv" \ @@ -23,10 +23,10 @@ hyperfine --runs 20 --warmup 3 \ ### # Create a virtual environment with seed packages. -# -# TODO(charlie): Support seed packages in `puffin venv`. ### hyperfine --runs 20 --warmup 3 \ + --prepare "rm -rf .venv" \ + "./target/release/puffin venv --seed" \ --prepare "rm -rf .venv" \ "virtualenv .venv" \ --prepare "rm -rf .venv" \