Support more formats in puffin venv, incl. windows support (#1039)

Mirroring `virtualenv -p` and driven by the lack of `pythonx.y` in
`PATH` on windows, this PR adds `-p x.y` support to `puffin venv` (first
commit).

Supported formats:
* NEW: `-p 3.10` searches for an installed Python 3.10 (Looking for
`python3.10` on linux/mac).
  Specifying a patch version is not supported
* `-p python3.10` or `-p python.exe` looks for a binary in `PATH`
* `-p /home/ferris/.local/bin/python3.10` uses this exact Python

In the second commit, we add python interpreter search on windows using
`py --list-paths`. On windows, all python are called `python.exe` so the
unix trick of looking for `python{}.{}` in `PATH` doesn't work. Instead,
we ask the python launcher for windows to tell us about all installed
packages. We should eventually migrate this to [PEP
514](https://peps.python.org/pep-0514/) by reading the registry entries
ourselves.

crates/puffin-interpreter/Cargo.toml

@@ -23,6 +23,7 @@ puffin-fs = { path = "../puffin-fs" }
 
 fs-err = { workspace = true, features = ["tokio"] }
 once_cell = { workspace = true }
+regex = { workspace = true }
 rmp-serde = { workspace = true }
 serde = { workspace = true, features = ["derive"] }
 serde_json = { workspace = true }
@@ -32,5 +33,8 @@ tracing = { workspace = true }
 which = { workspace = true}
 
 [dev-dependencies]
+anyhow = { version = "1.0.79" }
 indoc = { version = "2.0.4" }
+insta = { version = "1.34.0" }
+itertools = { version = "0.12.0" }
 tempfile = { version = "3.9.0" }

crates/puffin-interpreter/src/interpreter.rs

@@ -15,6 +15,7 @@ use puffin_cache::{Cache, CacheBucket, CachedByTimestamp};
 use puffin_fs::write_atomic_sync;
 
 use crate::python_platform::PythonPlatform;
+use crate::python_query::find_python_windows;
 use crate::virtual_env::detect_virtual_env;
 use crate::{Error, PythonVersion};
 
@@ -89,7 +90,7 @@ impl Interpreter {
     /// We check, in order:
     /// * `VIRTUAL_ENV` and `CONDA_PREFIX`
     /// * A `.venv` folder
-    /// * If a python version is given: `pythonx.y` (TODO(konstin): `py -x.y` on windows),
+    /// * If a python version is given: `pythonx.y`
     /// * `python3` (unix) or `python.exe` (windows)
     pub fn find(
         python_version: Option<&PythonVersion>,
@@ -103,8 +104,7 @@ impl Interpreter {
             return Ok(interpreter);
         };
 
-        #[cfg(unix)]
-        {
+        if cfg!(unix) {
             if let Some(python_version) = python_version {
                 let requested = format!(
                     "python{}.{}",
@@ -123,23 +123,21 @@ impl Interpreter {
             debug!("Resolved python3 to {}", executable.display());
             let interpreter = Interpreter::query(&executable, platform.0, cache)?;
             Ok(interpreter)
-        }
-
-        #[cfg(windows)]
-        {
-            if let Some(_python_version) = python_version {
-                unimplemented!("Implement me")
+        } else if cfg!(windows) {
+            if let Some(python_version) = python_version {
+                if let Some(path) =
+                    find_python_windows(python_version.major(), python_version.minor())?
+                {
+                    return Interpreter::query(&path, platform.0, cache);
+                }
             }
 
             let executable = which::which("python.exe")
                 .map_err(|err| Error::WhichNotFound("python.exe".to_string(), err))?;
             let interpreter = Interpreter::query(&executable, platform.0, cache)?;
             Ok(interpreter)
-        }
-
-        #[cfg(not(any(unix, windows)))]
-        {
-            compile_error!("only unix (like mac and linux) and windows are supported")
+        } else {
+            unimplemented!("Only Windows and Unix are supported");
         }
     }
 

crates/puffin-interpreter/src/lib.rs

@@ -6,12 +6,14 @@ use thiserror::Error;
 
 pub use crate::cfg::Configuration;
 pub use crate::interpreter::Interpreter;
+pub use crate::python_query::find_requested_python;
 pub use crate::python_version::PythonVersion;
 pub use crate::virtual_env::Virtualenv;
 
 mod cfg;
 mod interpreter;
 mod python_platform;
+mod python_query;
 mod python_version;
 mod virtual_env;
 
@@ -37,6 +39,10 @@ pub enum Error {
         #[source]
         err: io::Error,
     },
+    #[error("Failed to run `py --list-paths` to find Python installations")]
+    PyList(#[source] io::Error),
+    #[error("No Python {major}.{minor} found through `py --list-paths`")]
+    NoSuchPython { major: u8, minor: u8 },
     #[error("{message}:\n--- stdout:\n{stdout}\n--- stderr:\n{stderr}\n---")]
     PythonSubcommandOutput {
         message: String,
@@ -51,4 +57,6 @@ pub enum Error {
     Encode(#[from] rmp_serde::encode::Error),
     #[error("Failed to parse pyvenv.cfg")]
     Cfg(#[from] cfg::Error),
+    #[error("Couldn't find `{0}` in PATH")]
+    Which(PathBuf, #[source] which::Error),
 }

crates/puffin-interpreter/src/python_query.rs

@@ -0,0 +1,149 @@
+//! Find a user requested python version/interpreter.
+
+use std::path::PathBuf;
+use std::process::Command;
+
+use once_cell::sync::Lazy;
+use regex::Regex;
+use tracing::info_span;
+
+use crate::Error;
+
+/// ```text
+/// -V:3.12          C:\Users\Ferris\AppData\Local\Programs\Python\Python312\python.exe
+/// -V:3.8           C:\Users\Ferris\AppData\Local\Programs\Python\Python38\python.exe
+/// ```
+static PY_LIST_PATHS: Lazy<Regex> = Lazy::new(|| {
+    // Without the `R` flag, paths have trailing \r
+    Regex::new(r"(?mR)^ -(?:V:)?(\d).(\d+)-?(?:arm)?(?:\d*)\s*\*?\s*(.*)$").unwrap()
+});
+
+/// Find a user requested python version/interpreter.
+///
+/// Supported formats:
+/// * `-p 3.10` searches for an installed Python 3.10 (`py --list-paths` on Windows, `python3.10` on Linux/Mac).
+///   Specifying a patch version is not supported.
+/// * `-p python3.10` or `-p python.exe` looks for a binary in `PATH`.
+/// * `-p /home/ferris/.local/bin/python3.10` uses this exact Python.
+pub fn find_requested_python(request: &str) -> Result<PathBuf, Error> {
+    let major_minor = request
+        .split_once('.')
+        .and_then(|(major, minor)| Some((major.parse::<u8>().ok()?, minor.parse::<u8>().ok()?)));
+    if let Some((major, minor)) = major_minor {
+        // `-p 3.10`
+        if cfg!(unix) {
+            let formatted = PathBuf::from(format!("python{major}.{minor}"));
+            which::which_global(&formatted).map_err(|err| Error::Which(formatted, err))
+        } else if cfg!(windows) {
+            find_python_windows(major, minor)?.ok_or(Error::NoSuchPython { major, minor })
+        } else {
+            unimplemented!("Only Windows and Unix are supported")
+        }
+    } else if !request.contains(std::path::MAIN_SEPARATOR) {
+        // `-p python3.10`; Generally not used on windows because all Python are `python.exe`.
+        let request = PathBuf::from(request);
+        which::which_global(&request).map_err(|err| Error::Which(request, err))
+    } else {
+        // `-p /home/ferris/.local/bin/python3.10`
+        Ok(fs_err::canonicalize(request)?)
+    }
+}
+
+/// Run `py --list-paths` to find the installed pythons.
+///
+/// The command takes 8ms on my machine. TODO(konstin): Implement <https://peps.python.org/pep-0514/> to read python
+/// installations from the registry instead.
+fn installed_pythons_windows() -> Result<Vec<(u8, u8, PathBuf)>, Error> {
+    let output = info_span!("py_list_paths")
+        .in_scope(|| Command::new("py").arg("--list-paths").output())
+        .map_err(Error::PyList)?;
+
+    // There shouldn't be any output on stderr.
+    if !output.status.success() || !output.stderr.is_empty() {
+        return Err(Error::PythonSubcommandOutput {
+            message: format!(
+                "Running `py --list-paths` failed with status {}",
+                output.status
+            ),
+            stdout: String::from_utf8_lossy(&output.stdout).trim().to_string(),
+            stderr: String::from_utf8_lossy(&output.stderr).trim().to_string(),
+        });
+    }
+
+    // Find the first python of the version we want in the list
+    let stdout =
+        String::from_utf8(output.stdout.clone()).map_err(|err| Error::PythonSubcommandOutput {
+            message: format!("Running `py --list-paths` stdout isn't UTF-8 encoded: {err}"),
+            stdout: String::from_utf8_lossy(&output.stdout).trim().to_string(),
+            stderr: String::from_utf8_lossy(&output.stderr).trim().to_string(),
+        })?;
+    Ok(PY_LIST_PATHS
+        .captures_iter(&stdout)
+        .filter_map(|captures| {
+            let (_, [major, minor, path]) = captures.extract();
+            Some((
+                major.parse::<u8>().ok()?,
+                minor.parse::<u8>().ok()?,
+                PathBuf::from(path),
+            ))
+        })
+        .collect())
+}
+
+pub(crate) fn find_python_windows(major: u8, minor: u8) -> Result<Option<PathBuf>, Error> {
+    Ok(installed_pythons_windows()?
+        .into_iter()
+        .find(|(major_, minor_, _path)| *major_ == major && *minor_ == minor)
+        .map(|(_, _, path)| path))
+}
+
+#[cfg(test)]
+mod tests {
+    use std::fmt::Debug;
+
+    use insta::{assert_display_snapshot, assert_snapshot};
+    use itertools::Itertools;
+
+    use crate::python_query::find_requested_python;
+    use crate::Error;
+
+    fn format_err<T: Debug>(err: Result<T, Error>) -> String {
+        anyhow::Error::new(err.unwrap_err())
+            .chain()
+            .join("\n  Caused by: ")
+    }
+
+    #[cfg(unix)]
+    #[test]
+    fn python312() {
+        assert_eq!(
+            find_requested_python("3.12").unwrap(),
+            find_requested_python("python3.12").unwrap()
+        );
+    }
+
+    #[test]
+    fn no_such_python_version() {
+        assert_snapshot!(format_err(find_requested_python("3.1000")), @r###"
+        Couldn't find `3.1000` in PATH
+          Caused by: cannot find binary path
+        "###);
+    }
+
+    #[test]
+    fn no_such_python_binary() {
+        assert_display_snapshot!(format_err(find_requested_python("python3.1000")), @r###"
+        Couldn't find `python3.1000` in PATH
+          Caused by: cannot find binary path
+        "###);
+    }
+
+    #[test]
+    fn no_such_python_path() {
+        assert_display_snapshot!(
+            format_err(find_requested_python("/does/not/exists/python3.12")), @r###"
+        failed to canonicalize path `/does/not/exists/python3.12`
+          Caused by: No such file or directory (os error 2)
+        "###);
+    }
+}

crates/puffin/src/commands/venv.rs

@@ -1,5 +1,5 @@
 use std::fmt::Write;
-use std::path::{Path, PathBuf};
+use std::path::Path;
 use std::str::FromStr;
 
 use anyhow::Result;
@@ -15,7 +15,7 @@ use platform_host::Platform;
 use puffin_cache::Cache;
 use puffin_client::{FlatIndex, FlatIndexClient, RegistryClientBuilder};
 use puffin_dispatch::BuildDispatch;
-use puffin_interpreter::Interpreter;
+use puffin_interpreter::{find_requested_python, Interpreter};
 use puffin_resolver::InMemoryIndex;
 use puffin_traits::{BuildContext, InFlight, SetupPyStrategy};
 
@@ -26,13 +26,13 @@ use crate::printer::Printer;
 #[allow(clippy::unnecessary_wraps)]
 pub(crate) async fn venv(
     path: &Path,
-    base_python: Option<&Path>,
+    python_request: Option<&str>,
     index_locations: &IndexLocations,
     seed: bool,
     cache: &Cache,
     printer: Printer,
 ) -> Result<ExitStatus> {
-    match venv_impl(path, base_python, index_locations, seed, cache, printer).await {
+    match venv_impl(path, python_request, index_locations, seed, cache, printer).await {
         Ok(status) => Ok(status),
         Err(err) => {
             #[allow(clippy::print_stderr)]
@@ -50,10 +50,6 @@ enum VenvError {
     #[diagnostic(code(puffin::venv::python_not_found))]
     PythonNotFound,
 
-    #[error("Unable to find a Python interpreter {0}")]
-    #[diagnostic(code(puffin::venv::python_not_found))]
-    UserPythonNotFound(PathBuf),
-
     #[error("Failed to extract Python interpreter info")]
     #[diagnostic(code(puffin::venv::interpreter))]
     InterpreterError(#[source] puffin_interpreter::Error),
@@ -78,19 +74,15 @@ enum VenvError {
 /// Create a virtual environment.
 async fn venv_impl(
     path: &Path,
-    base_python: Option<&Path>,
+    python_request: Option<&str>,
     index_locations: &IndexLocations,
     seed: bool,
     cache: &Cache,
     mut printer: Printer,
 ) -> miette::Result<ExitStatus> {
     // Locate the Python interpreter.
-    let base_python = if let Some(base_python) = base_python {
-        fs::canonicalize(
-            which::which_global(base_python)
-                .map_err(|_| VenvError::UserPythonNotFound(base_python.to_path_buf()))?,
-        )
-        .into_diagnostic()?
+    let base_python = if let Some(python_request) = python_request {
+        find_requested_python(python_request).into_diagnostic()?
     } else {
         fs::canonicalize(
             which::which_global("python3")

crates/puffin/src/main.rs

@@ -517,10 +517,18 @@ struct CleanArgs {
 #[allow(clippy::struct_excessive_bools)]
 struct VenvArgs {
     /// The Python interpreter to use for the virtual environment.
+    ///
+    /// Supported formats:
+    /// * `-p 3.10` searches for an installed Python 3.10 (`py --list-paths` on Windows, `python3.10` on Linux/Mac).
+    ///   Specifying a patch version is not supported.
+    /// * `-p python3.10` or `-p python.exe` looks for a binary in `PATH`.
+    /// * `-p /home/ferris/.local/bin/python3.10` uses this exact Python.
+    ///
+    /// Note that this is different from `--python-version` in `pip compile`, which takes `3.10` or `3.10.13` and
+    /// doesn't look for a Python interpreter on disk.
     // Short `-p` to match `virtualenv`
-    // TODO(konstin): Support e.g. `-p 3.10`
     #[clap(short, long)]
-    python: Option<PathBuf>,
+    python: Option<String>,
 
     /// Install seed packages (`pip`, `setuptools`, and `wheel`) into the virtual environment.
     #[clap(long)]