From 08f09e474329c93780bc99c49a3bba3a70eb27b5 Mon Sep 17 00:00:00 2001 From: Zanie Blue Date: Tue, 31 Oct 2023 11:59:40 -0500 Subject: [PATCH] Add support for `pip-compile --extra ` (#239) Adds support for `pip-compile --extra ...` 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 | 6 +- crates/puffin-cli/src/commands/pip_sync.rs | 2 +- .../puffin-cli/src/commands/pip_uninstall.rs | 2 +- crates/puffin-cli/src/main.rs | 6 ++ crates/puffin-cli/src/requirements.rs | 31 +++++-- crates/puffin-cli/tests/pip_compile.rs | 53 ++++++++++++ ...compile__compile_pyproject_toml_extra.snap | 28 +++++++ crates/puffin-package/src/extra_name.rs | 82 +++++++++++++++++++ crates/puffin-package/src/lib.rs | 1 + 9 files changed, 202 insertions(+), 9 deletions(-) create mode 100644 crates/puffin-cli/tests/snapshots/pip_compile__compile_pyproject_toml_extra.snap create mode 100644 crates/puffin-package/src/extra_name.rs diff --git a/crates/puffin-cli/src/commands/pip_compile.rs b/crates/puffin-cli/src/commands/pip_compile.rs index a8647302f..bb3d1a28f 100644 --- a/crates/puffin-cli/src/commands/pip_compile.rs +++ b/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, 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 = 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(); diff --git a/crates/puffin-cli/src/commands/pip_sync.rs b/crates/puffin-cli/src/commands/pip_sync.rs index 6c7dbfb55..1c74d17b3 100644 --- a/crates/puffin-cli/src/commands/pip_sync.rs +++ b/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")?; diff --git a/crates/puffin-cli/src/commands/pip_uninstall.rs b/crates/puffin-cli/src/commands/pip_uninstall.rs index a605d0bdc..005992840 100644 --- a/crates/puffin-cli/src/commands/pip_uninstall.rs +++ b/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()?; diff --git a/crates/puffin-cli/src/main.rs b/crates/puffin-cli/src/main.rs index 71bb6fed7..0f743e965 100644 --- a/crates/puffin-cli/src/main.rs +++ b/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, + /// Include optional dependencies in the given extra group name; may be provided more than once. + #[clap(long)] + extra: Vec, + #[clap(long, value_enum)] resolution: Option, @@ -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(), diff --git a/crates/puffin-cli/src/requirements.rs b/crates/puffin-cli/src/requirements.rs index bfe013acc..20cd16dc2 100644 --- a/crates/puffin-cli/src/requirements.rs +++ b/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 { + pub(crate) fn try_from_source( + source: &RequirementsSource, + extras: &[ExtraName], + ) -> Result { 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::(&contents) .with_context(|| format!("Failed to read `{}`", path.display()))?; - let requirements = pyproject_toml + let requirements: Vec = 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 { 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); } diff --git a/crates/puffin-cli/tests/pip_compile.rs b/crates/puffin-cli/tests/pip_compile.rs index eeb8a3933..b02235d94 100644 --- a/crates/puffin-cli/tests/pip_compile.rs +++ b/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(()) +} diff --git a/crates/puffin-cli/tests/snapshots/pip_compile__compile_pyproject_toml_extra.snap b/crates/puffin-cli/tests/snapshots/pip_compile__compile_pyproject_toml_extra.snap new file mode 100644 index 000000000..78e32bbac --- /dev/null +++ b/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] + diff --git a/crates/puffin-package/src/extra_name.rs b/crates/puffin-package/src/extra_name.rs new file mode 100644 index 000000000..920996127 --- /dev/null +++ b/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 = Lazy::new(|| Regex::new(r"[-_.]+").unwrap()); + +impl ExtraName { + /// See: + /// + pub fn normalize(name: impl AsRef) -> 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 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" + ); + } +} diff --git a/crates/puffin-package/src/lib.rs b/crates/puffin-package/src/lib.rs index c0f3bd132..7f57fe8ce 100644 --- a/crates/puffin-package/src/lib.rs +++ b/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;