diff --git a/crates/puffin-cli/src/commands/mod.rs b/crates/puffin-cli/src/commands/mod.rs index 330d1e675..67e65c345 100644 --- a/crates/puffin-cli/src/commands/mod.rs +++ b/crates/puffin-cli/src/commands/mod.rs @@ -5,7 +5,7 @@ pub(crate) use add::add; pub(crate) use clean::clean; use distribution_types::InstalledMetadata; pub(crate) use freeze::freeze; -pub(crate) use pip_compile::{extra_name_with_clap_error, pip_compile}; +pub(crate) use pip_compile::{extra_name_with_clap_error, pip_compile, Upgrade}; pub(crate) use pip_install::pip_install; pub(crate) use pip_sync::pip_sync; pub(crate) use pip_uninstall::pip_uninstall; diff --git a/crates/puffin-cli/src/commands/pip_compile.rs b/crates/puffin-cli/src/commands/pip_compile.rs index 006f5aa41..5f3d0ead9 100644 --- a/crates/puffin-cli/src/commands/pip_compile.rs +++ b/crates/puffin-cli/src/commands/pip_compile.rs @@ -11,6 +11,7 @@ use anyhow::{anyhow, Context, Result}; use chrono::{DateTime, Utc}; use itertools::Itertools; use owo_colors::OwoColorize; +use rustc_hash::FxHashSet; use tempfile::tempdir_in; use tracing::debug; @@ -23,7 +24,7 @@ use puffin_client::{FlatIndex, FlatIndexClient, RegistryClientBuilder}; use puffin_dispatch::BuildDispatch; use puffin_installer::Downloader; use puffin_interpreter::{Interpreter, PythonVersion}; -use puffin_normalize::ExtraName; +use puffin_normalize::{ExtraName, PackageName}; use puffin_resolver::{ DisplayResolutionGraph, InMemoryIndex, Manifest, PreReleaseMode, ResolutionMode, ResolutionOptions, Resolver, @@ -48,7 +49,7 @@ pub(crate) async fn pip_compile( output_file: Option<&Path>, resolution_mode: ResolutionMode, prerelease_mode: PreReleaseMode, - upgrade_mode: UpgradeMode, + upgrade: Upgrade, generate_hashes: bool, index_locations: IndexLocations, setup_py: SetupPyStrategy, @@ -110,7 +111,8 @@ pub(crate) async fn pip_compile( } let preferences: Vec = output_file - .filter(|_| upgrade_mode.is_prefer_pinned()) + // As an optimization, skip reading the lockfile is we're upgrading all packages anyway. + .filter(|_| !upgrade.is_all()) .filter(|output_file| output_file.exists()) .map(Path::to_path_buf) .map(RequirementsSource::from) @@ -118,6 +120,17 @@ pub(crate) async fn pip_compile( .map(|source| RequirementsSpecification::from_source(source, &extras)) .transpose()? .map(|spec| spec.requirements) + .map(|requirements| match upgrade { + // Respect all pinned versions from the existing lockfile. + Upgrade::None => requirements, + // Ignore all pinned versions from the existing lockfile. + Upgrade::All => vec![], + // Ignore pinned versions for the specified packages. + Upgrade::Packages(packages) => requirements + .into_iter() + .filter(|requirement| !packages.contains(&requirement.name)) + .collect(), + }) .unwrap_or_default(); // Detect the current Python interpreter. @@ -325,28 +338,34 @@ pub(crate) async fn pip_compile( } /// Whether to allow package upgrades. -#[derive(Debug, Clone, Copy, PartialEq, Eq)] -pub(crate) enum UpgradeMode { - /// Allow package upgrades, ignoring the existing lockfile. - AllowUpgrades, +#[derive(Debug)] +pub(crate) enum Upgrade { /// Prefer pinned versions from the existing lockfile, if possible. - PreferPinned, + None, + + /// Allow package upgrades for all packages, ignoring the existing lockfile. + All, + + /// Allow package upgrades, but only for the specified packages. + Packages(FxHashSet), } -impl UpgradeMode { - fn is_prefer_pinned(self) -> bool { - self == Self::PreferPinned - } -} - -impl From for UpgradeMode { - fn from(value: bool) -> Self { - if value { - Self::AllowUpgrades +impl Upgrade { + /// Determine the upgrade strategy from the command-line arguments. + pub(crate) fn from_args(upgrade: bool, upgrade_package: Vec) -> Self { + if upgrade { + Self::All + } else if !upgrade_package.is_empty() { + Self::Packages(upgrade_package.into_iter().collect()) } else { - Self::PreferPinned + Self::None } } + + /// Returns `true` if all packages should be upgraded. + pub(crate) fn is_all(&self) -> bool { + matches!(self, Self::All) + } } pub(crate) fn extra_name_with_clap_error(arg: &str) -> Result { diff --git a/crates/puffin-cli/src/main.rs b/crates/puffin-cli/src/main.rs index c9588115f..4d6645c0a 100644 --- a/crates/puffin-cli/src/main.rs +++ b/crates/puffin-cli/src/main.rs @@ -17,7 +17,7 @@ use puffin_resolver::{PreReleaseMode, ResolutionMode}; use puffin_traits::SetupPyStrategy; use requirements::ExtrasSpecification; -use crate::commands::{extra_name_with_clap_error, ExitStatus}; +use crate::commands::{extra_name_with_clap_error, ExitStatus, Upgrade}; use crate::requirements::RequirementsSource; #[cfg(target_os = "windows")] @@ -188,6 +188,11 @@ struct PipCompileArgs { #[clap(long)] upgrade: bool, + /// Allow upgrades for a specific package, ignoring pinned versions in the existing output + /// file. + #[clap(long)] + upgrade_package: Vec, + /// Include distribution hashes in the output file. #[clap(long)] generate_hashes: bool, @@ -552,6 +557,7 @@ async fn inner() -> Result { } else { ExtrasSpecification::Some(&args.extra) }; + let upgrade = Upgrade::from_args(args.upgrade, args.upgrade_package); commands::pip_compile( &requirements, &constraints, @@ -560,7 +566,7 @@ async fn inner() -> Result { args.output_file.as_deref(), args.resolution, args.prerelease, - args.upgrade.into(), + upgrade, args.generate_hashes, index_urls, if args.legacy_setup_py { diff --git a/crates/puffin-cli/tests/pip_compile.rs b/crates/puffin-cli/tests/pip_compile.rs index 6741b6db0..980a19bf9 100644 --- a/crates/puffin-cli/tests/pip_compile.rs +++ b/crates/puffin-cli/tests/pip_compile.rs @@ -1,16 +1,18 @@ #![cfg(all(feature = "python", feature = "pypi"))] -use std::iter; use std::path::PathBuf; use std::process::Command; +use std::{fs, iter}; use anyhow::{bail, Context, Result}; use assert_cmd::prelude::*; use assert_fs::prelude::*; use assert_fs::TempDir; use indoc::indoc; +use insta::assert_snapshot; use insta_cmd::_macro_support::insta; use insta_cmd::{assert_cmd_snapshot, get_cargo_bin}; +use itertools::Itertools; use common::{create_venv_py312, BIN_NAME, INSTA_FILTERS}; @@ -3194,3 +3196,227 @@ fn find_links_url() -> Result<()> { Ok(()) } + +/// Use an existing resolution for `black==23.10.1`, with stale versions of `click` and `pathspec`. +/// Nothing should change. +#[test] +fn upgrade_none() -> Result<()> { + let temp_dir = TempDir::new()?; + let cache_dir = TempDir::new()?; + let venv = create_venv_py312(&temp_dir, &cache_dir); + + let requirements_in = temp_dir.child("requirements.in"); + requirements_in.write_str("black==23.10.1")?; + + let requirements_txt = temp_dir.child("requirements.txt"); + requirements_txt.write_str(indoc! {r" + black==23.10.1 + click==8.1.2 + # via black + mypy-extensions==1.0.0 + # via black + packaging==23.2 + # via black + pathspec==0.11.0 + # via black + platformdirs==4.0.0 + # via black + "})?; + + insta::with_settings!({ + filters => INSTA_FILTERS.to_vec() + }, { + assert_cmd_snapshot!(Command::new(get_cargo_bin(BIN_NAME)) + .arg("pip") + .arg("compile") + .arg("requirements.in") + .arg("--output-file") + .arg("requirements.txt") + .arg("--cache-dir") + .arg(cache_dir.path()) + .arg("--exclude-newer") + .arg(EXCLUDE_NEWER) + .env("VIRTUAL_ENV", venv.as_os_str()) + .current_dir(&temp_dir), @r###" + success: true + exit_code: 0 + ----- stdout ----- + + ----- stderr ----- + Resolved 6 packages in [TIME] + "###); + }); + + // Read the output requirements, but skip the header. + let resolution = fs::read_to_string(requirements_txt.path())? + .lines() + .skip_while(|line| line.trim_start().starts_with('#')) + .join("\n"); + assert_snapshot!(resolution, @r###" + black==23.10.1 + click==8.1.2 + # via black + mypy-extensions==1.0.0 + # via black + packaging==23.2 + # via black + pathspec==0.11.0 + # via black + platformdirs==4.0.0 + # via black + "###); + + Ok(()) +} + +/// Use an existing resolution for `black==23.10.1`, with stale versions of `click` and `pathspec`. +/// Both packages should be upgraded. +#[test] +fn upgrade_all() -> Result<()> { + let temp_dir = TempDir::new()?; + let cache_dir = TempDir::new()?; + let venv = create_venv_py312(&temp_dir, &cache_dir); + + let requirements_in = temp_dir.child("requirements.in"); + requirements_in.write_str("black==23.10.1")?; + + let requirements_txt = temp_dir.child("requirements.txt"); + requirements_txt.write_str(indoc! {r" + # This file was autogenerated by Puffin v0.0.1 via the following command: + # puffin pip compile requirements.in --python-version 3.12 --cache-dir [CACHE_DIR] + black==23.10.1 + click==8.1.2 + # via black + mypy-extensions==1.0.0 + # via black + packaging==23.2 + # via black + pathspec==0.11.0 + # via black + platformdirs==4.0.0 + # via black + "})?; + + insta::with_settings!({ + filters => INSTA_FILTERS.to_vec() + }, { + assert_cmd_snapshot!(Command::new(get_cargo_bin(BIN_NAME)) + .arg("pip") + .arg("compile") + .arg("requirements.in") + .arg("--output-file") + .arg("requirements.txt") + .arg("--upgrade") + .arg("--cache-dir") + .arg(cache_dir.path()) + .arg("--exclude-newer") + .arg(EXCLUDE_NEWER) + .env("VIRTUAL_ENV", venv.as_os_str()) + .current_dir(&temp_dir), @r###" + success: true + exit_code: 0 + ----- stdout ----- + + ----- stderr ----- + Resolved 6 packages in [TIME] + "###); + }); + + // Read the output requirements, but skip the header. + let resolution = fs::read_to_string(requirements_txt.path())? + .lines() + .skip_while(|line| line.trim_start().starts_with('#')) + .join("\n"); + assert_snapshot!(resolution, @r###" + black==23.10.1 + click==8.1.7 + # via black + mypy-extensions==1.0.0 + # via black + packaging==23.2 + # via black + pathspec==0.11.2 + # via black + platformdirs==4.0.0 + # via black + "###); + + Ok(()) +} + +/// Use an existing resolution for `black==23.10.1`, with stale versions of `click` and `pathspec`. +/// Only `click` should be upgraded. +#[test] +fn upgrade_package() -> Result<()> { + let temp_dir = TempDir::new()?; + let cache_dir = TempDir::new()?; + let venv = create_venv_py312(&temp_dir, &cache_dir); + + let requirements_in = temp_dir.child("requirements.in"); + requirements_in.write_str("black==23.10.1")?; + + let requirements_txt = temp_dir.child("requirements.txt"); + requirements_txt.write_str(indoc! {r" + # This file was autogenerated by Puffin v0.0.1 via the following command: + # puffin pip compile requirements.in --python-version 3.12 --cache-dir [CACHE_DIR] + black==23.10.1 + click==8.1.2 + # via black + mypy-extensions==1.0.0 + # via black + packaging==23.2 + # via black + pathspec==0.11.0 + # via black + platformdirs==4.0.0 + # via black + "})?; + + insta::with_settings!({ + filters => INSTA_FILTERS.to_vec() + }, { + assert_cmd_snapshot!(Command::new(get_cargo_bin(BIN_NAME)) + .arg("pip") + .arg("compile") + .arg("requirements.in") + .arg("--output-file") + .arg("requirements.txt") + .arg("--upgrade-package") + .arg("click") + .arg("--cache-dir") + .arg(cache_dir.path()) + .arg("--exclude-newer") + .arg(EXCLUDE_NEWER) + .env("VIRTUAL_ENV", venv.as_os_str()) + .current_dir(&temp_dir), @r###" + success: true + exit_code: 0 + ----- stdout ----- + + ----- stderr ----- + Resolved 6 packages in [TIME] + "###); + }); + + // Read the output requirements, but skip the header. + let resolution = fs::read_to_string(requirements_txt.path())? + .lines() + .skip_while(|line| line.trim_start().starts_with('#')) + .join("\n"); + assert_snapshot!(resolution, @r###" + black==23.10.1 + click==8.1.7 + # via black + mypy-extensions==1.0.0 + # via black + packaging==23.2 + # via black + pathspec==0.11.0 + # via black + platformdirs==4.0.0 + # via black + "### + ); + + Ok(()) +}