Towards using prepare_metadata_for_build_wheel in the resolver (#616)

Make `prepare_metadata_for_build_wheel` accessible across the puffin
codebase by splitting the built call into a setup, a metadata and a
wheel call. This does not actually use the hook yet, but it's the
required refactoring for it.

Part of #599.

crates/puffin-build/src/lib.rs

@@ -4,9 +4,11 @@
 
 use std::env;
 use std::fmt::{Display, Formatter};
+use std::future::Future;
 use std::io;
 use std::io::BufRead;
 use std::path::{Path, PathBuf};
+use std::pin::Pin;
 use std::process::Output;
 use std::str::FromStr;
 use std::sync::Arc;
@@ -30,7 +32,7 @@ use zip::ZipArchive;
 
 use pep508_rs::Requirement;
 use puffin_interpreter::{Interpreter, Virtualenv};
-use puffin_traits::BuildContext;
+use puffin_traits::{BuildContext, SourceBuildTrait};
 
 /// e.g. `pygraphviz/graphviz_wrap.c:3020:10: fatal error: graphviz/cgraph.h: No such file or directory`
 static MISSING_HEADER_RE: Lazy<Regex> = Lazy::new(|| {
@@ -370,15 +372,20 @@ impl SourceBuild {
     }
 
     /// Try calling `prepare_metadata_for_build_wheel` to get the metadata without executing the
-    /// actual build
+    /// actual build.
     ///
     /// TODO(konstin): Return the actual metadata instead of the dist-info dir
-    pub async fn get_metadata_without_build(&mut self) -> Result<Option<&Path>, Error> {
+    pub async fn get_metadata_without_build(&mut self) -> Result<Option<PathBuf>, Error> {
         // setup.py builds don't support this.
         let Some(pep517_backend) = &self.pep517_backend else {
             return Ok(None);
         };
 
+        // We've already called this method, but return the existing result is easier than erroring
+        if let Some(metadata_dir) = &self.metadata_directory {
+            return Ok(Some(metadata_dir.clone()));
+        }
+
         let metadata_directory = self.temp_dir.path().join("metadata_directory");
         fs::create_dir(&metadata_directory)?;
 
@@ -430,7 +437,7 @@ impl SourceBuild {
             return Ok(None);
         }
         self.metadata_directory = Some(metadata_directory.join(message));
-        return Ok(self.metadata_directory.as_deref());
+        Ok(self.metadata_directory.clone())
     }
 
     /// Build a source distribution from an archive (`.zip` or `.tar.gz`), return the location of the
@@ -549,6 +556,21 @@ impl SourceBuild {
     }
 }
 
+impl SourceBuildTrait for SourceBuild {
+    fn metadata<'a>(
+        &'a mut self,
+    ) -> Pin<Box<dyn Future<Output = anyhow::Result<Option<PathBuf>>> + Send + 'a>> {
+        Box::pin(async { Ok(self.get_metadata_without_build().await?) })
+    }
+
+    fn wheel<'a>(
+        &'a self,
+        wheel_dir: &'a Path,
+    ) -> Pin<Box<dyn Future<Output = anyhow::Result<String>> + Send + 'a>> {
+        Box::pin(async { Ok(self.build(wheel_dir).await?) })
+    }
+}
+
 fn escape_path_for_python(path: &Path) -> String {
     path.to_string_lossy()
         .replace('\\', "\\\\")

crates/puffin-dev/src/build.rs

@@ -62,6 +62,7 @@ pub(crate) async fn build(args: BuildArgs) -> Result<PathBuf> {
         false,
         IndexUrls::default(),
     );
+
     let builder = SourceBuild::setup(
         &args.sdist,
         args.subdirectory.as_deref(),

crates/puffin-dispatch/src/lib.rs

@@ -66,6 +66,8 @@ impl BuildDispatch {
 }
 
 impl BuildContext for BuildDispatch {
+    type SourceDistBuilder = SourceBuild;
+
     fn cache(&self) -> &Cache {
         &self.cache
     }
@@ -213,13 +215,12 @@ impl BuildContext for BuildDispatch {
     }
 
     #[instrument(skip_all, fields(source_dist = source_dist, subdirectory = ?subdirectory))]
-    fn build_source<'a>(
+    fn setup_build<'a>(
         &'a self,
         source: &'a Path,
         subdirectory: Option<&'a Path>,
-        wheel_dir: &'a Path,
         source_dist: &'a str,
-    ) -> Pin<Box<dyn Future<Output = Result<String>> + Send + 'a>> {
+    ) -> Pin<Box<dyn Future<Output = Result<SourceBuild>> + Send + 'a>> {
         Box::pin(async move {
             if self.no_build {
                 bail!("Building source distributions is disabled");
@@ -234,7 +235,7 @@ impl BuildContext for BuildDispatch {
                 BuildKind::Wheel,
             )
             .await?;
-            Ok(builder.build(wheel_dir).await?)
+            Ok(builder)
         })
     }
 }

crates/puffin-distribution/src/source_dist.rs

@@ -31,7 +31,7 @@ use puffin_client::{CachedClient, CachedClientError, DataWithCachePolicy};
 use puffin_fs::write_atomic;
 use puffin_git::{Fetch, GitSource};
 use puffin_normalize::PackageName;
-use puffin_traits::BuildContext;
+use puffin_traits::{BuildContext, SourceBuildTrait};
 use pypi_types::Metadata21;
 
 use crate::locks::LockedFile;
@@ -643,12 +643,10 @@ impl<'a, T: BuildContext> SourceDistCachedBuilder<'a, T> {
         fs::create_dir_all(&cache_entry.dir).await?;
         let disk_filename = self
             .build_context
-            .build_source(
-                source_dist,
-                subdirectory,
-                &cache_entry.dir,
-                &dist.to_string(),
-            )
+            .setup_build(source_dist, subdirectory, &dist.to_string())
+            .await
+            .map_err(|err| SourceDistError::Build(Box::new(dist.clone()), err))?
+            .wheel(&cache_entry.dir)
             .await
             .map_err(|err| SourceDistError::Build(Box::new(dist.clone()), err))?;
 

crates/puffin-resolver/tests/resolver.rs

@@ -21,7 +21,7 @@ use puffin_interpreter::{Interpreter, Virtualenv};
 use puffin_resolver::{
     Graph, Manifest, PreReleaseMode, ResolutionMode, ResolutionOptions, Resolver,
 };
-use puffin_traits::BuildContext;
+use puffin_traits::{BuildContext, SourceBuildTrait};
 
 // Exclude any packages uploaded after this date.
 static EXCLUDE_NEWER: Lazy<DateTime<Utc>> = Lazy::new(|| {
@@ -36,6 +36,8 @@ struct DummyContext {
 }
 
 impl BuildContext for DummyContext {
+    type SourceDistBuilder = DummyBuilder;
+
     fn cache(&self) -> &Cache {
         &self.cache
     }
@@ -63,12 +65,28 @@ impl BuildContext for DummyContext {
         panic!("The test should not need to build source distributions")
     }
 
-    fn build_source<'a>(
+    fn setup_build<'a>(
         &'a self,
-        _sdist: &'a Path,
+        _source: &'a Path,
         _subdirectory: Option<&'a Path>,
-        _wheel_dir: &'a Path,
         _package_id: &'a str,
+    ) -> Pin<Box<dyn Future<Output = Result<Self::SourceDistBuilder>> + Send + 'a>> {
+        Box::pin(async { Ok(DummyBuilder) })
+    }
+}
+
+struct DummyBuilder;
+
+impl SourceBuildTrait for DummyBuilder {
+    fn metadata<'a>(
+        &'a mut self,
+    ) -> Pin<Box<dyn Future<Output = Result<Option<PathBuf>>> + Send + 'a>> {
+        panic!("The test should not need to build source distributions")
+    }
+
+    fn wheel<'a>(
+        &'a self,
+        _wheel_dir: &'a Path,
     ) -> Pin<Box<dyn Future<Output = Result<String>> + Send + 'a>> {
         panic!("The test should not need to build source distributions")
     }

crates/puffin-traits/src/lib.rs

@@ -1,7 +1,7 @@
 //! Avoid cyclic crate dependencies between resolver, installer and builder.
 
 use std::future::Future;
-use std::path::Path;
+use std::path::{Path, PathBuf};
 use std::pin::Pin;
 
 use anyhow::Result;
@@ -22,7 +22,7 @@ mod once_map;
 /// them and then build. The installer, the resolver and the source distribution builder are each in
 /// their own crate. To avoid circular crate dependencies, this type dispatches between the three
 /// crates with its three main methods ([`BuildContext::resolve`], [`BuildContext::install`] and
-/// [`BuildContext::build_source`]).
+/// [`BuildContext::setup_build`]).
 ///
 /// The overall main crate structure looks like this:
 ///
@@ -53,6 +53,8 @@ mod once_map;
 
 // TODO(konstin): Proper error types
 pub trait BuildContext {
+    type SourceDistBuilder: SourceBuildTrait + Send + Sync;
+
     fn cache(&self) -> &Cache;
 
     /// All (potentially nested) source distribution builds use the same base python and can reuse
@@ -62,7 +64,7 @@ pub trait BuildContext {
     /// The system (or conda) python interpreter to create venvs.
     fn base_python(&self) -> &Path;
 
-    /// Whether source distribution building is disabled. This [`BuildContext::build_source`] calls
+    /// Whether source distribution building is disabled. This [`BuildContext::setup_build`] calls
     /// will fail in this case. This method exists to avoid fetching source distributions if we know
     /// we can't build them
     fn no_build(&self) -> bool {
@@ -83,16 +85,42 @@ pub trait BuildContext {
         venv: &'a Virtualenv,
     ) -> Pin<Box<dyn Future<Output = Result<()>> + Send + 'a>>;
 
-    /// Build a source distribution into a wheel from an archive.
+    /// Setup a source distribution build by installing the required dependencies. A wrapper for
+    /// `puffin_build::SourceBuild::setup`.
     ///
-    /// Returns the filename of the built wheel inside the given `wheel_dir`.
+    /// For PEP 517 builds, this calls `get_requires_for_build_wheel`.
     ///
     /// `package_id` is for error reporting only.
-    fn build_source<'a>(
+    fn setup_build<'a>(
         &'a self,
         source: &'a Path,
         subdirectory: Option<&'a Path>,
-        wheel_dir: &'a Path,
         package_id: &'a str,
+    ) -> Pin<Box<dyn Future<Output = Result<Self::SourceDistBuilder>> + Send + 'a>>;
+}
+
+/// A wrapper for `puffin_build::SourceBuild` to avoid cyclical crate dependencies.
+///
+/// You can either call only `wheel()` to build the wheel directly, call only `metadata()` to get
+/// the metadata without performing the actual or first call `metadata()` and then `wheel()`.
+pub trait SourceBuildTrait {
+    /// A wrapper for `puffin_build::SourceBuild::get_metadata_without_build`.
+    ///
+    /// For PEP 517 builds, this calls `prepare_metadata_for_build_wheel`
+    ///
+    /// Returns the metadata directory if we're having a PEP 517 build and the
+    /// `prepare_metadata_for_build_wheel` hook exists
+    fn metadata<'a>(
+        &'a mut self,
+    ) -> Pin<Box<dyn Future<Output = Result<Option<PathBuf>>> + Send + 'a>>;
+
+    /// A wrapper for `puffin_build::SourceBuild::build`.
+    ///
+    /// For PEP 517 builds, this calls `build_wheel`.
+    ///
+    /// Returns the filename of the built wheel inside the given `wheel_dir`.
+    fn wheel<'a>(
+        &'a self,
+        wheel_dir: &'a Path,
     ) -> Pin<Box<dyn Future<Output = Result<String>> + Send + 'a>>;
 }