Allow constraints to be provided in --upgrade-package (#4952)

## Summary

Allows, e.g., `--upgrade-package flask<3.0.0`.

Closes https://github.com/astral-sh/uv/issues/1964.
This commit is contained in:
Charlie Marsh 2024-07-09 20:09:13 -07:00 committed by GitHub
parent 5b6dffe522
commit 23eb42deed
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
15 changed files with 175 additions and 34 deletions

View file

@ -1,8 +1,10 @@
use either::Either;
use pep508_rs::MarkerTree;
use pypi_types::Requirement;
use rustc_hash::{FxBuildHasher, FxHashMap};
use std::borrow::Cow;
use either::Either;
use rustc_hash::FxHashMap;
use pep508_rs::MarkerTree;
use pypi_types::{Requirement, RequirementSource};
use uv_normalize::PackageName;
/// A set of constraints for a set of requirements.
@ -11,10 +13,16 @@ pub struct Constraints(FxHashMap<PackageName, Vec<Requirement>>);
impl Constraints {
/// Create a new set of constraints from a set of requirements.
pub fn from_requirements(requirements: Vec<Requirement>) -> Self {
let mut constraints: FxHashMap<PackageName, Vec<Requirement>> =
FxHashMap::with_capacity_and_hasher(requirements.len(), FxBuildHasher);
pub fn from_requirements(requirements: impl Iterator<Item = Requirement>) -> Self {
let mut constraints: FxHashMap<PackageName, Vec<Requirement>> = FxHashMap::default();
for requirement in requirements {
// Skip empty constraints.
if let RequirementSource::Registry { specifier, .. } = &requirement.source {
if specifier.is_empty() {
continue;
}
}
constraints
.entry(requirement.name.clone())
.or_default()

View file

@ -1,6 +1,8 @@
use either::Either;
use pep508_rs::PackageName;
use rustc_hash::FxHashSet;
use pypi_types::Requirement;
use rustc_hash::FxHashMap;
/// Whether to reinstall packages.
#[derive(Debug, Default, Clone)]
@ -54,12 +56,12 @@ pub enum Upgrade {
All,
/// Allow package upgrades, but only for the specified packages.
Packages(FxHashSet<PackageName>),
Packages(FxHashMap<PackageName, Vec<Requirement>>),
}
impl Upgrade {
/// Determine the upgrade strategy from the command-line arguments.
pub fn from_args(upgrade: Option<bool>, upgrade_package: Vec<PackageName>) -> Self {
pub fn from_args(upgrade: Option<bool>, upgrade_package: Vec<Requirement>) -> Self {
match upgrade {
Some(true) => Self::All,
Some(false) => Self::None,
@ -67,7 +69,15 @@ impl Upgrade {
if upgrade_package.is_empty() {
Self::None
} else {
Self::Packages(upgrade_package.into_iter().collect())
Self::Packages(upgrade_package.into_iter().fold(
FxHashMap::default(),
|mut map, requirement| {
map.entry(requirement.name.clone())
.or_default()
.push(requirement);
map
},
))
}
}
}
@ -82,4 +92,28 @@ impl Upgrade {
pub fn is_all(&self) -> bool {
matches!(self, Self::All)
}
/// Returns `true` if the specified package should be upgraded.
pub fn contains(&self, package_name: &PackageName) -> bool {
match &self {
Self::None => false,
Self::All => true,
Self::Packages(packages) => packages.contains_key(package_name),
}
}
/// Returns an iterator over the constraints.
///
/// When upgrading, users can provide bounds on the upgrade (e.g., `--upgrade-package flask<3`).
pub fn constraints(&self) -> impl Iterator<Item = &Requirement> {
if let Self::Packages(packages) = self {
Either::Right(
packages
.values()
.flat_map(|requirements| requirements.iter()),
)
} else {
Either::Left(std::iter::empty())
}
}
}