From a43328d914daa4d602200c70328f2433b92a89da Mon Sep 17 00:00:00 2001
From: Charlie Marsh <charlie.r.marsh@gmail.com>
Date: Fri, 6 Oct 2023 00:47:45 -0400
Subject: [PATCH] Support wheel installation (#19)

Closes https://github.com/astral-sh/puffin/issues/8.
---
 crates/puffin-cli/Cargo.toml              | 10 +++-
 crates/puffin-cli/src/commands/compile.rs | 18 ++++--
 crates/puffin-cli/src/commands/install.rs | 58 +++++++++++++++---
 crates/puffin-cli/src/logging.rs          | 16 +++--
 crates/puffin-client/Cargo.toml           |  3 +-
 crates/puffin-client/src/api.rs           | 24 ++++++++
 crates/puffin-client/src/lib.rs           |  2 +-
 crates/puffin-interpreter/src/lib.rs      | 17 +++++-
 crates/puffin-resolve/src/lib.rs          | 72 +++++++++++++++--------
 9 files changed, 174 insertions(+), 46 deletions(-)

diff --git a/crates/puffin-cli/Cargo.toml b/crates/puffin-cli/Cargo.toml
index 5016ab16f..e164d5e3b 100644
--- a/crates/puffin-cli/Cargo.toml
+++ b/crates/puffin-cli/Cargo.toml
@@ -11,17 +11,21 @@ puffin-package = { path = "../puffin-package" }
 puffin-resolve = { path = "../puffin-resolve" }
 
 anyhow = { version = "1.0.75" }
-clap = { version = "4.4.6", features = ["derive"] }
-colored = { version = "2.0.4" }
 async-std = { version = "1.12.0", features = [
     "attributes",
     "tokio1",
     "unstable",
 ] }
+clap = { version = "4.4.6", features = ["derive"] }
+colored = { version = "2.0.4" }
+directories = { version = "5.0.1" }
 futures = { version = "0.3.28" }
+install-wheel-rs = { version = "0.0.1" }
 pep508_rs = { version = "0.2.3" }
 pep440_rs = { version = "0.3.12" }
 tracing = { version = "0.1.37" }
 tracing-tree = { version = "0.2.5" }
 tracing-subscriber = { version = "0.3.17", features = ["env-filter"] }
-directories = "5.0.1"
+rayon = { version = "1.8.0" }
+url = { version = "2.4.1" }
+tempfile = { version = "3.8.0" }
diff --git a/crates/puffin-cli/src/commands/compile.rs b/crates/puffin-cli/src/commands/compile.rs
index aea032b50..d881bb5e1 100644
--- a/crates/puffin-cli/src/commands/compile.rs
+++ b/crates/puffin-cli/src/commands/compile.rs
@@ -2,6 +2,7 @@ use std::path::Path;
 use std::str::FromStr;
 
 use anyhow::Result;
+use puffin_client::PypiClientBuilder;
 use tracing::debug;
 
 use puffin_interpreter::PythonExecutable;
@@ -32,13 +33,22 @@ pub(crate) async fn compile(src: &Path, cache: Option<&Path>) -> Result<ExitStat
     // Determine the compatible platform tags.
     let tags = Tags::from_env(&platform, python.version())?;
 
-    // Resolve the dependencies.
-    let resolution = resolve(&requirements, markers, &tags, cache).await?;
+    // Instantiate a client.
+    let client = {
+        let mut pypi_client = PypiClientBuilder::default();
+        if let Some(cache) = cache {
+            pypi_client = pypi_client.cache(cache);
+        }
+        pypi_client.build()
+    };
 
-    for (name, version) in resolution.iter() {
+    // Resolve the dependencies.
+    let resolution = resolve(&requirements, markers, &tags, &client).await?;
+
+    for (name, package) in resolution.iter() {
         #[allow(clippy::print_stdout)]
         {
-            println!("{name}=={version}");
+            println!("{}=={}", name, package.version());
         }
     }
 
diff --git a/crates/puffin-cli/src/commands/install.rs b/crates/puffin-cli/src/commands/install.rs
index b0d6876eb..63ff460b9 100644
--- a/crates/puffin-cli/src/commands/install.rs
+++ b/crates/puffin-cli/src/commands/install.rs
@@ -2,8 +2,12 @@ use std::path::Path;
 use std::str::FromStr;
 
 use anyhow::Result;
+use async_std::fs::File;
 use tracing::debug;
+use url::Url;
 
+use install_wheel_rs::{install_wheel, InstallLocation};
+use puffin_client::PypiClientBuilder;
 use puffin_interpreter::PythonExecutable;
 use puffin_platform::tags::Tags;
 use puffin_platform::Platform;
@@ -32,14 +36,54 @@ pub(crate) async fn install(src: &Path, cache: Option<&Path>) -> Result<ExitStat
     // Determine the compatible platform tags.
     let tags = Tags::from_env(&platform, python.version())?;
 
-    // Resolve the dependencies.
-    let resolution = resolve(&requirements, markers, &tags, cache).await?;
-
-    for (name, version) in resolution.iter() {
-        #[allow(clippy::print_stdout)]
-        {
-            println!("{name}=={version}");
+    // Instantiate a client.
+    let client = {
+        let mut pypi_client = PypiClientBuilder::default();
+        if let Some(cache) = cache {
+            pypi_client = pypi_client.cache(cache);
         }
+        pypi_client.build()
+    };
+
+    // Resolve the dependencies.
+    let resolution = resolve(&requirements, markers, &tags, &client).await?;
+
+    // Create a temporary directory, in which we'll store the wheels.
+    let tmp_dir = tempfile::tempdir()?;
+
+    // Download each wheel.
+    // TODO(charlie): Store these in a content-addressed cache.
+    // TODO(charlie): Use channels to efficiently stream-and-install.
+    for (name, package) in resolution.iter() {
+        let url = Url::parse(package.url())?;
+        let reader = client.stream_external(&url).await?;
+
+        // TODO(charlie): Stream the unzip.
+        let mut writer = File::create(tmp_dir.path().join(format!("{name}.whl"))).await?;
+        async_std::io::copy(reader, &mut writer).await?;
+    }
+
+    // Install each wheel.
+    // TODO(charlie): Use channels to efficiently stream-and-install.
+    let location = InstallLocation::Venv {
+        venv_base: python.venv().to_path_buf(),
+        python_version: python.simple_version(),
+    };
+    let locked_dir = location.acquire_lock()?;
+    for (name, package) in resolution.iter() {
+        let path = tmp_dir.path().join(format!("{name}.whl"));
+        let filename = install_wheel_rs::WheelFilename::from_str(package.filename())?;
+
+        // TODO(charlie): Should this be async?
+        install_wheel(
+            &locked_dir,
+            std::fs::File::open(path)?,
+            filename,
+            false,
+            &[],
+            "",
+            python.executable(),
+        )?;
     }
 
     Ok(ExitStatus::Success)
diff --git a/crates/puffin-cli/src/logging.rs b/crates/puffin-cli/src/logging.rs
index 3931cf290..6e1f1dc88 100644
--- a/crates/puffin-cli/src/logging.rs
+++ b/crates/puffin-cli/src/logging.rs
@@ -1,18 +1,26 @@
 use anyhow::Result;
+use tracing::level_filters::LevelFilter;
+use tracing_subscriber::filter::Targets;
 use tracing_subscriber::layer::SubscriberExt;
 use tracing_subscriber::{EnvFilter, Layer, Registry};
 use tracing_tree::time::Uptime;
 
 pub(crate) fn setup_logging() -> Result<()> {
+    let targets = Targets::new()
+        .with_target("hyper", LevelFilter::WARN)
+        .with_target("reqwest", LevelFilter::WARN)
+        .with_target("async_io", LevelFilter::WARN)
+        .with_target("async_std", LevelFilter::WARN)
+        .with_target("blocking", LevelFilter::OFF)
+        .with_default(LevelFilter::TRACE);
+
     let subscriber = Registry::default().with(
         tracing_tree::HierarchicalLayer::default()
-            .with_indent_lines(true)
-            .with_indent_amount(2)
-            .with_bracketed_fields(true)
             .with_targets(true)
             .with_writer(|| Box::new(std::io::stderr()))
             .with_timer(Uptime::default())
-            .with_filter(EnvFilter::from_default_env()),
+            .with_filter(EnvFilter::from_default_env())
+            .with_filter(targets),
     );
     tracing::subscriber::set_global_default(subscriber)?;
 
diff --git a/crates/puffin-client/Cargo.toml b/crates/puffin-client/Cargo.toml
index 0bda52b23..a62536014 100644
--- a/crates/puffin-client/Cargo.toml
+++ b/crates/puffin-client/Cargo.toml
@@ -7,7 +7,7 @@ edition = "2021"
 puffin-package = { path = "../puffin-package" }
 
 http-cache-reqwest = { version = "0.11.3" }
-reqwest = { version = "0.11.22", features = ["json", "gzip", "brotli"] }
+reqwest = { version = "0.11.22", features = ["json", "gzip", "brotli", "stream"] }
 reqwest-middleware = { version = "0.2.3" }
 reqwest-retry = { version = "0.3.0" }
 serde = { version = "1.0.188" }
@@ -15,4 +15,5 @@ serde_json = { version = "1.0.107" }
 thiserror = { version = "1.0.49" }
 url = { version = "2.4.1" }
 tracing = { version = "0.1.37" }
+futures = "0.3.28"
 
diff --git a/crates/puffin-client/src/api.rs b/crates/puffin-client/src/api.rs
index 4d0d7e621..43c06c52c 100644
--- a/crates/puffin-client/src/api.rs
+++ b/crates/puffin-client/src/api.rs
@@ -1,5 +1,6 @@
 use std::fmt::Debug;
 
+use futures::{AsyncRead, StreamExt, TryStreamExt};
 use reqwest::StatusCode;
 use serde::{Deserialize, Serialize};
 use tracing::trace;
@@ -12,6 +13,7 @@ use crate::client::PypiClient;
 use crate::error::PypiClientError;
 
 impl PypiClient {
+    /// Fetch a package from the `PyPI` simple API.
     pub async fn simple(
         &self,
         package_name: impl AsRef<str>,
@@ -62,6 +64,7 @@ impl PypiClient {
             .await?)
     }
 
+    /// Fetch the metadata from a wheel file.
     pub async fn file(&self, file: File) -> Result<Metadata21, PypiClientError> {
         // Send to the proxy.
         let url = self.proxy.join(
@@ -101,6 +104,27 @@ impl PypiClient {
             .text()
             .await?)
     }
+
+    /// Stream a file from an external URL.
+    pub async fn stream_external(
+        &self,
+        url: &Url,
+    ) -> Result<Box<dyn AsyncRead + Unpin + Send + Sync>, PypiClientError> {
+        Ok(Box::new(
+            // TODO(charlie): Use an uncached client.
+            self.client
+                .get(url.to_string())
+                .send()
+                .await?
+                .error_for_status()?
+                .bytes_stream()
+                .map(|r| match r {
+                    Ok(bytes) => Ok(bytes),
+                    Err(err) => Err(std::io::Error::new(std::io::ErrorKind::Other, err)),
+                })
+                .into_async_read(),
+        ))
+    }
 }
 
 #[derive(Debug, Clone, Serialize, Deserialize)]
diff --git a/crates/puffin-client/src/lib.rs b/crates/puffin-client/src/lib.rs
index 6427814d3..2bbac2061 100644
--- a/crates/puffin-client/src/lib.rs
+++ b/crates/puffin-client/src/lib.rs
@@ -1,5 +1,5 @@
 pub use api::{File, SimpleJson};
-pub use client::PypiClientBuilder;
+pub use client::{PypiClient, PypiClientBuilder};
 
 mod api;
 mod client;
diff --git a/crates/puffin-interpreter/src/lib.rs b/crates/puffin-interpreter/src/lib.rs
index 7257f0f92..0c82417fa 100644
--- a/crates/puffin-interpreter/src/lib.rs
+++ b/crates/puffin-interpreter/src/lib.rs
@@ -15,6 +15,7 @@ mod virtual_env;
 /// A Python executable and its associated platform markers.
 #[derive(Debug)]
 pub struct PythonExecutable {
+    venv: PathBuf,
     executable: PathBuf,
     markers: MarkerEnvironment,
 }
@@ -24,15 +25,21 @@ impl PythonExecutable {
     pub fn from_env(platform: &Platform) -> Result<Self> {
         let platform = PythonPlatform::from(platform);
         let venv = virtual_env::detect_virtual_env(&platform)?;
-        let executable = platform.venv_python(venv);
+        let executable = platform.venv_python(&venv);
         let markers = markers::detect_markers(&executable)?;
 
         Ok(Self {
+            venv,
             executable,
             markers,
         })
     }
 
+    /// Returns the path to the Python virtual environment.
+    pub fn venv(&self) -> &Path {
+        self.venv.as_path()
+    }
+
     /// Returns the path to the Python executable.
     pub fn executable(&self) -> &Path {
         self.executable.as_path()
@@ -47,4 +54,12 @@ impl PythonExecutable {
     pub fn version(&self) -> &Version {
         &self.markers.python_version.version
     }
+
+    /// Returns the Python version as a simple tuple.
+    pub fn simple_version(&self) -> (u8, u8) {
+        (
+            u8::try_from(self.version().release[0]).expect("invalid major version"),
+            u8::try_from(self.version().release[1]).expect("invalid minor version"),
+        )
+    }
 }
diff --git a/crates/puffin-resolve/src/lib.rs b/crates/puffin-resolve/src/lib.rs
index 8216246b4..9ff32c2c7 100644
--- a/crates/puffin-resolve/src/lib.rs
+++ b/crates/puffin-resolve/src/lib.rs
@@ -1,5 +1,5 @@
 use std::collections::{HashMap, HashSet};
-use std::path::Path;
+
 use std::str::FromStr;
 
 use anyhow::Result;
@@ -9,7 +9,7 @@ use pep440_rs::Version;
 use pep508_rs::{MarkerEnvironment, Requirement, VersionOrUrl};
 use tracing::debug;
 
-use puffin_client::{File, PypiClientBuilder, SimpleJson};
+use puffin_client::{File, PypiClient, SimpleJson};
 use puffin_package::metadata::Metadata21;
 use puffin_package::package_name::PackageName;
 use puffin_package::requirements::Requirements;
@@ -17,30 +17,41 @@ use puffin_package::wheel::WheelFilename;
 use puffin_platform::tags::Tags;
 
 #[derive(Debug)]
-pub struct Resolution(HashMap<PackageName, Version>);
+pub struct Resolution(HashMap<PackageName, PinnedPackage>);
 
 impl Resolution {
-    pub fn iter(&self) -> impl Iterator<Item = (&PackageName, &Version)> {
+    pub fn iter(&self) -> impl Iterator<Item = (&PackageName, &PinnedPackage)> {
         self.0.iter()
     }
 }
 
+#[derive(Debug)]
+pub struct PinnedPackage {
+    metadata: Metadata21,
+    file: File,
+}
+
+impl PinnedPackage {
+    pub fn filename(&self) -> &str {
+        &self.file.filename
+    }
+
+    pub fn url(&self) -> &str {
+        &self.file.url
+    }
+
+    pub fn version(&self) -> &Version {
+        &self.metadata.version
+    }
+}
+
 /// Resolve a set of requirements into a set of pinned versions.
 pub async fn resolve(
     requirements: &Requirements,
     markers: &MarkerEnvironment,
     tags: &Tags,
-    cache: Option<&Path>,
+    client: &PypiClient,
 ) -> Result<Resolution> {
-    // Instantiate a client.
-    let pypi_client = {
-        let mut pypi_client = PypiClientBuilder::default();
-        if let Some(cache) = cache {
-            pypi_client = pypi_client.cache(cache);
-        }
-        pypi_client.build()
-    };
-
     // A channel to fetch package metadata (e.g., given `flask`, fetch all versions) and version
     // metadata (e.g., given `flask==1.0.0`, fetch the metadata for that version).
     let (package_sink, package_stream) = futures::channel::mpsc::unbounded();
@@ -49,14 +60,16 @@ pub async fn resolve(
     let mut package_stream = package_stream
         .map(|request: Request| match request {
             Request::Package(requirement) => Either::Left(
-                pypi_client
+                client
+                    // TODO(charlie): Remove this clone.
                     .simple(requirement.name.clone())
-                    .map_ok(move |metadata| Response::Package(metadata, requirement)),
+                    .map_ok(move |metadata| Response::Package(requirement, metadata)),
             ),
             Request::Version(requirement, file) => Either::Right(
-                pypi_client
-                    .file(file)
-                    .map_ok(move |metadata| Response::Version(metadata, requirement)),
+                client
+                    // TODO(charlie): Remove this clone.
+                    .file(file.clone())
+                    .map_ok(move |metadata| Response::Version(requirement, file, metadata)),
             ),
         })
         .buffer_unordered(32)
@@ -71,13 +84,14 @@ pub async fn resolve(
     }
 
     // Resolve the requirements.
-    let mut resolution: HashMap<PackageName, Version> = HashMap::with_capacity(requirements.len());
+    let mut resolution: HashMap<PackageName, PinnedPackage> =
+        HashMap::with_capacity(requirements.len());
 
     while let Some(chunk) = package_stream.next().await {
         for result in chunk {
             let result: Response = result?;
             match result {
-                Response::Package(metadata, requirement) => {
+                Response::Package(requirement, metadata) => {
                     // TODO(charlie): Support URLs. Right now, we treat a URL as an unpinned dependency.
                     let specifiers =
                         requirement
@@ -112,7 +126,7 @@ pub async fn resolve(
 
                     package_sink.unbounded_send(Request::Version(requirement, file.clone()))?;
                 }
-                Response::Version(metadata, requirement) => {
+                Response::Version(requirement, file, metadata) => {
                     debug!(
                         "--> selected version {} for {}",
                         metadata.version, requirement
@@ -121,12 +135,20 @@ pub async fn resolve(
                     // Add to the resolved set.
                     let normalized_name = PackageName::normalize(&requirement.name);
                     in_flight.remove(&normalized_name);
-                    resolution.insert(normalized_name, metadata.version);
+                    resolution.insert(
+                        normalized_name,
+                        PinnedPackage {
+                            // TODO(charlie): Remove this clone.
+                            metadata: metadata.clone(),
+                            file,
+                        },
+                    );
 
                     // Enqueue its dependencies.
                     for dependency in metadata.requires_dist {
                         if !dependency.evaluate_markers(
                             markers,
+                            // TODO(charlie): Remove this clone.
                             requirement.extras.clone().unwrap_or_default(),
                         ) {
                             debug!("--> ignoring {dependency} due to environment mismatch");
@@ -167,6 +189,6 @@ enum Request {
 
 #[derive(Debug)]
 enum Response {
-    Package(SimpleJson, Requirement),
-    Version(Metadata21, Requirement),
+    Package(Requirement, SimpleJson),
+    Version(Requirement, File, Metadata21),
 }