Add Python interpreter detection (#11)

Closes https://github.com/astral-sh/puffin/issues/2.

crates/puffin-cli/Cargo.toml

@@ -3,10 +3,9 @@ name = "puffin-cli"
 version = "0.1.0"
 edition = "2021"
 
-# See more keys and their definitions at https://doc.rust-lang.org/cargo/reference/manifest.html
-
 [dependencies]
 puffin-client = { path = "../puffin-client" }
+puffin-interpreter = { path = "../puffin-interpreter" }
 puffin-requirements = { path = "../puffin-requirements" }
 
 anyhow = { version = "1.0.75" }

crates/puffin-cli/src/commands/install.rs

@@ -6,10 +6,11 @@ use anyhow::Result;
 use futures::future::Either;
 use futures::{StreamExt, TryFutureExt};
 use pep440_rs::Version;
-use pep508_rs::{MarkerEnvironment, Requirement, StringVersion, VersionOrUrl};
+use pep508_rs::{Requirement, VersionOrUrl};
 use tracing::debug;
 
 use puffin_client::{File, PypiClientBuilder, SimpleJson};
+use puffin_interpreter::PythonExecutable;
 use puffin_requirements::metadata::Metadata21;
 use puffin_requirements::package_name::PackageName;
 use puffin_requirements::wheel::WheelName;
@@ -29,27 +30,19 @@ enum Response {
 }
 
 pub(crate) async fn install(src: &Path) -> Result<ExitStatus> {
-    // TODO(charlie): Fetch from the environment.
-    let env = MarkerEnvironment {
-        implementation_name: String::new(),
-        implementation_version: StringVersion::from_str("3.10.0").unwrap(),
-        os_name: String::new(),
-        platform_machine: String::new(),
-        platform_python_implementation: String::new(),
-        platform_release: String::new(),
-        platform_system: String::new(),
-        platform_version: String::new(),
-        python_full_version: StringVersion::from_str("3.10.0").unwrap(),
-        python_version: StringVersion::from_str("3.10.0").unwrap(),
-        sys_platform: String::new(),
-    };
-
     // Read the `requirements.txt` from disk.
     let requirements_txt = std::fs::read_to_string(src)?;
 
     // Parse the `requirements.txt` into a list of requirements.
     let requirements = puffin_requirements::Requirements::from_str(&requirements_txt)?;
 
+    // Detect the current Python interpreter.
+    let python = PythonExecutable::from_env()?;
+    debug!(
+        "Using Python interpreter: {}",
+        python.executable().display()
+    );
+
     // Instantiate a client.
     let pypi_client = PypiClientBuilder::default().build();
     let proxy_client = PypiClientBuilder::default().build();
@@ -130,9 +123,10 @@ pub(crate) async fn install(src: &Path) -> Result<ExitStatus> {
 
                     // Enqueue its dependencies.
                     for dependency in metadata.requires_dist {
-                        if !dependency
-                            .evaluate_markers(&env, requirement.extras.clone().unwrap_or_default())
-                        {
+                        if !dependency.evaluate_markers(
+                            python.markers(),
+                            requirement.extras.clone().unwrap_or_default(),
+                        ) {
                             debug!("--> ignoring {dependency} due to environment mismatch");
                             continue;
                         }

crates/puffin-client/Cargo.toml

@@ -3,8 +3,6 @@ name = "puffin-client"
 version = "0.1.0"
 edition = "2021"
 
-# See more keys and their definitions at https://doc.rust-lang.org/cargo/reference/manifest.html
-
 [dependencies]
 puffin-requirements = { path = "../puffin-requirements" }
 

crates/puffin-interpreter/Cargo.toml

@@ -0,0 +1,16 @@
+[package]
+name = "puffin-interpreter"
+version = "0.1.0"
+edition.workspace = true
+rust-version.workspace = true
+homepage.workspace = true
+documentation.workspace = true
+repository.workspace = true
+authors.workspace = true
+license.workspace = true
+
+[dependencies]
+anyhow = { version = "1.0.75" }
+pep508_rs = { version = "0.2.3", features = ["serde"] }
+serde_json = { version = "1.0.107" }
+tracing = { version = "0.1.37" }

crates/puffin-interpreter/src/lib.rs

@@ -0,0 +1,40 @@
+use std::path::{Path, PathBuf};
+
+use anyhow::Result;
+use pep508_rs::MarkerEnvironment;
+
+use crate::platform::Platform;
+
+mod markers;
+mod platform;
+mod virtual_env;
+
+/// A Python executable and its associated platform markers.
+#[derive(Debug)]
+pub struct PythonExecutable {
+    executable: PathBuf,
+    markers: MarkerEnvironment,
+}
+
+impl PythonExecutable {
+    /// Detect the current Python executable from the host environment.
+    pub fn from_env() -> Result<Self> {
+        let target = Platform::from_host();
+        let venv = virtual_env::detect_virtual_env(&target)?;
+        let executable = target.get_venv_python(venv);
+        let markers = markers::detect_markers(&executable)?;
+
+        Ok(Self {
+            executable,
+            markers,
+        })
+    }
+
+    pub fn executable(&self) -> &Path {
+        self.executable.as_path()
+    }
+
+    pub fn markers(&self) -> &MarkerEnvironment {
+        &self.markers
+    }
+}

crates/puffin-interpreter/src/markers.rs

@@ -0,0 +1,59 @@
+use std::ffi::OsStr;
+use std::path::Path;
+use std::process::{Command, Output};
+
+use anyhow::{Context, Result};
+use pep508_rs::MarkerEnvironment;
+
+/// Return the resolved [`MarkerEnvironment`] for the given Python executable.
+pub(crate) fn detect_markers(python: impl AsRef<Path>) -> Result<MarkerEnvironment> {
+    let output = call_python(python.as_ref(), ["-c", CAPTURE_MARKERS_SCRIPT])?;
+    Ok(serde_json::from_slice::<MarkerEnvironment>(&output.stdout)?)
+}
+
+const CAPTURE_MARKERS_SCRIPT: &str = "
+import os
+import sys
+import platform
+import json
+def format_full_version(info):
+    version = '{0.major}.{0.minor}.{0.micro}'.format(info)
+    kind = info.releaselevel
+    if kind != 'final':
+        version += kind[0] + str(info.serial)
+    return version
+
+if hasattr(sys, 'implementation'):
+    implementation_version = format_full_version(sys.implementation.version)
+    implementation_name = sys.implementation.name
+else:
+    implementation_version = '0'
+    implementation_name = ''
+bindings = {
+    'implementation_name': implementation_name,
+    'implementation_version': implementation_version,
+    'os_name': os.name,
+    'platform_machine': platform.machine(),
+    'platform_python_implementation': platform.python_implementation(),
+    'platform_release': platform.release(),
+    'platform_system': platform.system(),
+    'platform_version': platform.version(),
+    'python_full_version': platform.python_version(),
+    'python_version': '.'.join(platform.python_version_tuple()[:2]),
+    'sys_platform': sys.platform,
+}
+json.dump(bindings, sys.stdout)
+sys.stdout.flush()
+";
+
+/// Run a Python script and return its output.
+fn call_python<I, S>(python: &Path, args: I) -> Result<Output>
+where
+    I: IntoIterator<Item = S>,
+    S: AsRef<OsStr>,
+{
+    Command::new(python)
+        .args(args)
+        .output()
+        .context(format!("Failed to run `python` at: {:?}", &python))
+}

crates/puffin-interpreter/src/platform.rs

@@ -0,0 +1,105 @@
+use std::env;
+use std::fmt;
+use std::path::Path;
+use std::path::PathBuf;
+
+#[derive(Debug, Clone, Eq, PartialEq)]
+pub(crate) struct Platform {
+    os: Option<Os>,
+}
+
+impl Platform {
+    /// Infer the target based on the current version used for compilation.
+    pub(crate) fn from_host() -> Self {
+        Self {
+            os: if cfg!(windows) {
+                Some(Os::Windows)
+            } else if cfg!(unix) {
+                Some(Os::Linux)
+            } else if cfg!(macos) {
+                Some(Os::Macos)
+            } else {
+                None
+            },
+        }
+    }
+
+    /// Returns `true` if the current platform is Linux.
+    #[allow(unused)]
+    #[inline]
+    pub(crate) fn is_linux(&self) -> bool {
+        self.os == Some(Os::Linux)
+    }
+
+    /// Returns `true` if the current platform is macOS.
+    #[allow(unused)]
+    #[inline]
+    pub(crate) fn is_macos(&self) -> bool {
+        self.os == Some(Os::Macos)
+    }
+
+    /// Returns `true` if the current platform is Windows.
+    #[allow(unused)]
+    #[inline]
+    pub(crate) fn is_windows(&self) -> bool {
+        self.os == Some(Os::Windows)
+    }
+
+    /// Returns the path to the `python` executable inside a virtual environment.
+    pub(crate) fn get_venv_python(&self, venv_base: impl AsRef<Path>) -> PathBuf {
+        self.get_venv_bin_dir(venv_base).join(self.get_python())
+    }
+
+    /// Returns the directory in which the binaries are stored inside a virtual environment.
+    pub(crate) fn get_venv_bin_dir(&self, venv_base: impl AsRef<Path>) -> PathBuf {
+        let venv = venv_base.as_ref();
+        if self.is_windows() {
+            let bin_dir = venv.join("Scripts");
+            if bin_dir.join("python.exe").exists() {
+                return bin_dir;
+            }
+            // Python installed via msys2 on Windows might produce a POSIX-like venv
+            // See https://github.com/PyO3/maturin/issues/1108
+            let bin_dir = venv.join("bin");
+            if bin_dir.join("python.exe").exists() {
+                return bin_dir;
+            }
+            // for conda environment
+            venv.to_path_buf()
+        } else {
+            venv.join("bin")
+        }
+    }
+
+    /// Returns the path to the `python` executable.
+    ///
+    /// For Windows, it's always `python.exe`. For UNIX, it's the `python` in the virtual
+    /// environment; or, if there is no virtual environment, `python3`.
+    pub(crate) fn get_python(&self) -> PathBuf {
+        if self.is_windows() {
+            PathBuf::from("python.exe")
+        } else if env::var_os("VIRTUAL_ENV").is_some() {
+            PathBuf::from("python")
+        } else {
+            PathBuf::from("python3")
+        }
+    }
+}
+
+/// All supported operating systems.
+#[derive(Debug, Clone, Copy, Eq, PartialEq)]
+enum Os {
+    Linux,
+    Windows,
+    Macos,
+}
+
+impl fmt::Display for Os {
+    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
+        match *self {
+            Os::Linux => write!(f, "Linux"),
+            Os::Windows => write!(f, "Windows"),
+            Os::Macos => write!(f, "macOS"),
+        }
+    }
+}

crates/puffin-interpreter/src/virtual_env.rs

@@ -0,0 +1,47 @@
+use std::env;
+use std::path::PathBuf;
+
+use anyhow::{bail, Result};
+use tracing::debug;
+
+use crate::platform::Platform;
+
+/// Locate the current virtual environment.
+pub(crate) fn detect_virtual_env(target: &Platform) -> Result<PathBuf> {
+    match (env::var_os("VIRTUAL_ENV"), env::var_os("CONDA_PREFIX")) {
+        (Some(dir), None) => return Ok(PathBuf::from(dir)),
+        (None, Some(dir)) => return Ok(PathBuf::from(dir)),
+        (Some(_), Some(_)) => {
+            bail!("Both VIRTUAL_ENV and CONDA_PREFIX are set. Please unset one of them.")
+        }
+        (None, None) => {
+            // No environment variables set. Try to find a virtualenv in the current directory.
+        }
+    };
+
+    // Search for a `.venv` directory in the current or any parent directory.
+    let current_dir = env::current_dir().expect("Failed to detect current directory");
+    for dir in current_dir.ancestors() {
+        let dot_venv = dir.join(".venv");
+        if dot_venv.is_dir() {
+            if !dot_venv.join("pyvenv.cfg").is_file() {
+                bail!(
+                    "Expected {} to be a virtual environment, but pyvenv.cfg is missing",
+                    dot_venv.display()
+                );
+            }
+            let python = target.get_venv_python(&dot_venv);
+            if !python.is_file() {
+                bail!(
+                    "Your virtualenv at {} is broken. It contains a pyvenv.cfg but no python at {}",
+                    dot_venv.display(),
+                    python.display()
+                );
+            }
+            debug!("Found a virtualenv named .venv at {}", dot_venv.display());
+            return Ok(dot_venv);
+        }
+    }
+
+    bail!("Couldn't find a virtualenv or conda environment.")
+}

crates/puffin-requirements/Cargo.toml

@@ -3,13 +3,8 @@ name = "puffin-requirements"
 version = "0.1.0"
 edition = "2021"
 
-# See more keys and their definitions at https://doc.rust-lang.org/cargo/reference/manifest.html
-
 [dependencies]
 anyhow = { version = "1.0.75" }
-clap = { version = "4.4.6", features = ["derive"] }
-colored = { version = "2.0.4" }
-insta = { version = "1.33.0" }
 mailparse = { version = "0.14.0" }
 memchr = { version = "2.6.4" }
 once_cell = { version = "1.18.0" }