Add support for pip-compile --extra <name> (#239)

Adds support for `pip-compile --extra <name> ...` which includes
optional dependencies in the specified group in the resolution.

Following precedent in `pip-compile`, if a given extra is not found,
there is no error. ~We could consider warning in this case.~ We should
probably add an error but it expands scope and will be considered
separately in #241

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

@@ -8,6 +8,7 @@ use colored::Colorize;
 use fs_err::File;
 use itertools::Itertools;
 use pubgrub::report::Reporter;
+use puffin_package::extra_name::ExtraName;
 use tracing::debug;
 
 use pep508_rs::Requirement;
@@ -31,6 +32,7 @@ const VERSION: &str = env!("CARGO_PKG_VERSION");
 pub(crate) async fn pip_compile(
     requirements: &[RequirementsSource],
     constraints: &[RequirementsSource],
+    extras: Vec<ExtraName>,
     output_file: Option<&Path>,
     resolution_mode: ResolutionMode,
     prerelease_mode: PreReleaseMode,
@@ -45,14 +47,14 @@ pub(crate) async fn pip_compile(
     let RequirementsSpecification {
         requirements,
         constraints,
-    } = RequirementsSpecification::try_from_sources(requirements, constraints)?;
+    } = RequirementsSpecification::try_from_sources(requirements, constraints, &extras)?;
     let preferences: Vec<Requirement> = output_file
         .filter(|_| upgrade_mode.is_prefer_pinned())
         .filter(|output_file| output_file.exists())
         .map(Path::to_path_buf)
         .map(RequirementsSource::from)
         .as_ref()
-        .map(RequirementsSpecification::try_from_source)
+        .map(|source| RequirementsSpecification::try_from_source(source, &extras))
         .transpose()?
         .map(|spec| spec.requirements)
         .unwrap_or_default();

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

@@ -34,7 +34,7 @@ pub(crate) async fn pip_sync(
     let RequirementsSpecification {
         requirements,
         constraints: _,
-    } = RequirementsSpecification::try_from_sources(sources, &[])?;
+    } = RequirementsSpecification::try_from_sources(sources, &[], &[])?;
 
     if requirements.is_empty() {
         writeln!(printer, "No requirements found")?;

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

@@ -25,7 +25,7 @@ pub(crate) async fn pip_uninstall(
     let RequirementsSpecification {
         requirements,
         constraints: _,
-    } = RequirementsSpecification::try_from_sources(sources, &[])?;
+    } = RequirementsSpecification::try_from_sources(sources, &[], &[])?;
 
     // Detect the current Python interpreter.
     let platform = Platform::current()?;

crates/puffin-cli/src/main.rs

@@ -4,6 +4,7 @@ use std::process::ExitCode;
 use clap::{Args, Parser, Subcommand};
 use colored::Colorize;
 use directories::ProjectDirs;
+use puffin_package::extra_name::ExtraName;
 use puffin_resolver::{PreReleaseMode, ResolutionMode};
 use url::Url;
 
@@ -71,6 +72,10 @@ struct PipCompileArgs {
     #[clap(short, long)]
     constraint: Vec<PathBuf>,
 
+    /// Include optional dependencies in the given extra group name; may be provided more than once.
+    #[clap(long)]
+    extra: Vec<ExtraName>,
+
     #[clap(long, value_enum)]
     resolution: Option<ResolutionMode>,
 
@@ -201,6 +206,7 @@ async fn main() -> ExitCode {
             commands::pip_compile(
                 &requirements,
                 &constraints,
+                args.extra,
                 args.output_file.as_deref(),
                 args.resolution.unwrap_or_default(),
                 args.prerelease.unwrap_or_default(),

crates/puffin-cli/src/requirements.rs

@@ -7,6 +7,7 @@ use anyhow::{Context, Result};
 use fs_err as fs;
 
 use pep508_rs::Requirement;
+use puffin_package::extra_name::ExtraName;
 use puffin_package::requirements_txt::RequirementsTxt;
 
 #[derive(Debug)]
@@ -45,7 +46,10 @@ pub(crate) struct RequirementsSpecification {
 
 impl RequirementsSpecification {
     /// Read the requirements and constraints from a source.
-    pub(crate) fn try_from_source(source: &RequirementsSource) -> Result<Self> {
+    pub(crate) fn try_from_source(
+        source: &RequirementsSource,
+        extras: &[ExtraName],
+    ) -> Result<Self> {
         Ok(match source {
             RequirementsSource::Name(name) => {
                 let requirement = Requirement::from_str(name)
@@ -70,11 +74,27 @@ impl RequirementsSpecification {
                 let contents = fs::read_to_string(path)?;
                 let pyproject_toml = toml::from_str::<pyproject_toml::PyProjectToml>(&contents)
                     .with_context(|| format!("Failed to read `{}`", path.display()))?;
-                let requirements = pyproject_toml
+                let requirements: Vec<Requirement> = pyproject_toml
                     .project
                     .into_iter()
-                    .flat_map(|project| project.dependencies.into_iter().flatten())
+                    .flat_map(|project| {
+                        project.dependencies.into_iter().flatten().chain(
+                            // Include any optional dependencies specified in `extras`
+                            project.optional_dependencies.into_iter().flat_map(
+                                |optional_dependencies| {
+                                    extras.iter().flat_map(move |extra| {
+                                        optional_dependencies
+                                            .get(extra.as_ref())
+                                            .cloned()
+                                            // undefined extra requests are ignored silently
+                                            .unwrap_or_default()
+                                    })
+                                },
+                            ),
+                        )
+                    })
                     .collect();
+
                 Self {
                     requirements,
                     constraints: vec![],
@@ -87,6 +107,7 @@ impl RequirementsSpecification {
     pub(crate) fn try_from_sources(
         requirements: &[RequirementsSource],
         constraints: &[RequirementsSource],
+        extras: &[ExtraName],
     ) -> Result<Self> {
         let mut spec = Self::default();
 
@@ -94,14 +115,14 @@ impl RequirementsSpecification {
         // A `requirements.txt` can contain a `-c constraints.txt` directive within it, so reading
         // a requirements file can also add constraints.
         for source in requirements {
-            let source = Self::try_from_source(source)?;
+            let source = Self::try_from_source(source, extras)?;
             spec.requirements.extend(source.requirements);
             spec.constraints.extend(source.constraints);
         }
 
         // Read all constraints, treating both requirements _and_ constraints as constraints.
         for source in constraints {
-            let source = Self::try_from_source(source)?;
+            let source = Self::try_from_source(source, extras)?;
             spec.constraints.extend(source.requirements);
             spec.constraints.extend(source.constraints);
         }

crates/puffin-cli/tests/pip_compile.rs

@@ -227,3 +227,56 @@ fn compile_constraints_inline() -> Result<()> {
 
     Ok(())
 }
+
+/// Resolve a package from an extra in a `pyproject.toml` file.
+#[test]
+fn compile_pyproject_toml_extra() -> Result<()> {
+    let temp_dir = assert_fs::TempDir::new()?;
+    let cache_dir = assert_fs::TempDir::new()?;
+    let venv = temp_dir.child(".venv");
+
+    Command::new(get_cargo_bin(BIN_NAME))
+        .arg("venv")
+        .arg(venv.as_os_str())
+        .arg("--cache-dir")
+        .arg(cache_dir.path())
+        .current_dir(&temp_dir)
+        .assert()
+        .success();
+    venv.assert(predicates::path::is_dir());
+
+    let pyproject_toml = temp_dir.child("pyproject.toml");
+    pyproject_toml.touch()?;
+    pyproject_toml.write_str(
+        r#"[build-system]
+requires = ["setuptools", "wheel"]
+
+[project]
+name = "project"
+dependencies = []
+optional-dependencies.foo = [
+    "django==5.0b1",
+]
+"#,
+    )?;
+
+    insta::with_settings!({
+        filters => vec![
+            (r"\d+(ms|s)", "[TIME]"),
+            (r"#    .* pip-compile", "#    [BIN_PATH] pip-compile"),
+            (r"--cache-dir .*", "--cache-dir [CACHE_DIR]"),
+        ]
+    }, {
+        assert_cmd_snapshot!(Command::new(get_cargo_bin(BIN_NAME))
+            .arg("pip-compile")
+            .arg("pyproject.toml")
+            .arg("--extra")
+            .arg("foo")
+            .arg("--cache-dir")
+            .arg(cache_dir.path())
+            .env("VIRTUAL_ENV", venv.as_os_str())
+            .current_dir(&temp_dir));
+    });
+
+    Ok(())
+}

crates/puffin-cli/tests/snapshots/pip_compile__compile_pyproject_toml_extra.snap

@@ -0,0 +1,28 @@
+---
+source: crates/puffin-cli/tests/pip_compile.rs
+info:
+  program: puffin
+  args:
+    - pip-compile
+    - pyproject.toml
+    - "--extra"
+    - foo
+    - "--cache-dir"
+    - /var/folders/bc/qlsk3t6x7c9fhhbvvcg68k9c0000gp/T/.tmpAYEAdM
+  env:
+    VIRTUAL_ENV: /var/folders/bc/qlsk3t6x7c9fhhbvvcg68k9c0000gp/T/.tmp1xuOcV/.venv
+---
+success: true
+exit_code: 0
+----- stdout -----
+# This file was autogenerated by Puffin v0.0.1 via the following command:
+#    [BIN_PATH] pip-compile pyproject.toml --extra foo --cache-dir [CACHE_DIR]
+asgiref==3.7.2
+    # via django
+django==5.0b1
+sqlparse==0.4.4
+    # via django
+
+----- stderr -----
+Resolved 3 packages in [TIME]
+

crates/puffin-package/src/extra_name.rs

@@ -0,0 +1,82 @@
+use std::fmt;
+use std::fmt::{Display, Formatter};
+
+use once_cell::sync::Lazy;
+use regex::Regex;
+
+#[derive(Debug, Clone, PartialEq, Eq, Hash, PartialOrd, Ord)]
+pub struct ExtraName(String);
+
+impl From<&ExtraName> for ExtraName {
+    fn from(extra_name: &ExtraName) -> Self {
+        extra_name.clone()
+    }
+}
+
+impl Display for ExtraName {
+    fn fmt(&self, f: &mut Formatter<'_>) -> fmt::Result {
+        self.0.fmt(f)
+    }
+}
+
+static NAME_NORMALIZE: Lazy<Regex> = Lazy::new(|| Regex::new(r"[-_.]+").unwrap());
+
+impl ExtraName {
+    /// See: <https://peps.python.org/pep-0685/#specification/>
+    ///      <https://packaging.python.org/en/latest/specifications/name-normalization/>
+    pub fn normalize(name: impl AsRef<str>) -> Self {
+        // TODO(charlie): Avoid allocating in the common case (when no normalization is required).
+        let mut normalized = NAME_NORMALIZE.replace_all(name.as_ref(), "-").to_string();
+        normalized.make_ascii_lowercase();
+        Self(normalized)
+    }
+}
+
+impl AsRef<str> for ExtraName {
+    fn as_ref(&self) -> &str {
+        self.0.as_ref()
+    }
+}
+
+impl From<&str> for ExtraName {
+    fn from(name: &str) -> Self {
+        Self::normalize(name)
+    }
+}
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+
+    #[test]
+    fn normalize() {
+        assert_eq!(
+            ExtraName::normalize("friendly-bard").as_ref(),
+            "friendly-bard"
+        );
+        assert_eq!(
+            ExtraName::normalize("Friendly-Bard").as_ref(),
+            "friendly-bard"
+        );
+        assert_eq!(
+            ExtraName::normalize("FRIENDLY-BARD").as_ref(),
+            "friendly-bard"
+        );
+        assert_eq!(
+            ExtraName::normalize("friendly.bard").as_ref(),
+            "friendly-bard"
+        );
+        assert_eq!(
+            ExtraName::normalize("friendly_bard").as_ref(),
+            "friendly-bard"
+        );
+        assert_eq!(
+            ExtraName::normalize("friendly--bard").as_ref(),
+            "friendly-bard"
+        );
+        assert_eq!(
+            ExtraName::normalize("FrIeNdLy-._.-bArD").as_ref(),
+            "friendly-bard"
+        );
+    }
+}

crates/puffin-package/src/lib.rs

@@ -1,4 +1,5 @@
 pub mod dist_info_name;
+pub mod extra_name;
 pub mod metadata;
 pub mod package_name;
 pub mod requirements_txt;