Add `--seed` flag to `venv` to allow seed package environments (#865)

## 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

@@ -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<ExitStatus> {
-    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<ExitStatus> {
@@ -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)
 }

crates/puffin-cli/src/main.rs

@@ -408,9 +408,25 @@ struct VenvArgs {
     #[clap(short, long)]
     python: Option<PathBuf>,
 
+    /// 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<IndexUrl>,
+
+    /// 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<ExitStatus> {
         }
         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),
     }

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(())
+}

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" \