Add an --update-package setting to allow individual package upgrades (#953)

Closes #950.

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;

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<Requirement> = 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)]
+#[derive(Debug)]
-pub(crate) enum UpgradeMode {
+pub(crate) enum Upgrade {
-    /// Allow package upgrades, ignoring the existing lockfile.
-    AllowUpgrades,
     /// 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<PackageName>),
 }
 
-impl UpgradeMode {
+impl Upgrade {
-    fn is_prefer_pinned(self) -> bool {
+    /// Determine the upgrade strategy from the command-line arguments.
-        self == Self::PreferPinned
+    pub(crate) fn from_args(upgrade: bool, upgrade_package: Vec<PackageName>) -> Self {
-    }
+        if upgrade {
-}
+            Self::All
-
+        } else if !upgrade_package.is_empty() {
-impl From<bool> for UpgradeMode {
+            Self::Packages(upgrade_package.into_iter().collect())
-    fn from(value: bool) -> Self {
-        if value {
-            Self::AllowUpgrades
         } 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<ExtraName> {

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<PackageName>,
+
     /// Include distribution hashes in the output file.
     #[clap(long)]
     generate_hashes: bool,
@@ -552,6 +557,7 @@ async fn inner() -> Result<ExitStatus> {
             } 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<ExitStatus> {
                 args.output_file.as_deref(),
                 args.resolution,
                 args.prerelease,
-                args.upgrade.into(),
+                upgrade,
                 args.generate_hashes,
                 index_urls,
                 if args.legacy_setup_py {

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(())
+}