uv_configuration/
constraints.rs

1use std::borrow::Cow;
2
3use either::Either;
4use rustc_hash::FxHashMap;
5
6use uv_distribution_types::{Requirement, RequirementSource};
7use uv_normalize::PackageName;
8use uv_pep508::MarkerTree;
9
10/// A set of constraints for a set of requirements.
11#[derive(Debug, Default, Clone)]
12pub struct Constraints(FxHashMap<PackageName, Vec<Requirement>>);
13
14impl Constraints {
15    /// Create a new set of constraints from a set of requirements.
16    pub fn from_requirements(requirements: impl Iterator<Item = Requirement>) -> Self {
17        let mut constraints: FxHashMap<PackageName, Vec<Requirement>> = FxHashMap::default();
18        for requirement in requirements {
19            // Skip empty constraints.
20            if let RequirementSource::Registry { specifier, .. } = &requirement.source {
21                if specifier.is_empty() {
22                    continue;
23                }
24            }
25
26            constraints
27                .entry(requirement.name.clone())
28                .or_default()
29                .push(Requirement {
30                    // We add and apply constraints independent of their extras.
31                    extras: Box::new([]),
32                    ..requirement
33                });
34        }
35        Self(constraints)
36    }
37
38    /// Return an iterator over all [`Requirement`]s in the constraint set.
39    pub fn requirements(&self) -> impl Iterator<Item = &Requirement> {
40        self.0.values().flat_map(|requirements| requirements.iter())
41    }
42
43    /// Get the constraints for a package.
44    pub fn get(&self, name: &PackageName) -> Option<&Vec<Requirement>> {
45        self.0.get(name)
46    }
47
48    /// Apply the constraints to a set of requirements.
49    ///
50    /// NB: Change this method together with [`Overrides::apply`].
51    pub fn apply<'a>(
52        &'a self,
53        requirements: impl IntoIterator<Item = Cow<'a, Requirement>>,
54    ) -> impl Iterator<Item = Cow<'a, Requirement>> {
55        requirements.into_iter().flat_map(|requirement| {
56            let Some(constraints) = self.get(&requirement.name) else {
57                // Case 1: No constraint(s).
58                return Either::Left(std::iter::once(requirement));
59            };
60
61            // ASSUMPTION: There is one `extra = "..."`, and it's either the only marker or part
62            // of the main conjunction.
63            let Some(extra_expression) = requirement.marker.top_level_extra() else {
64                // Case 2: A non-optional dependency with constraint(s).
65                return Either::Right(Either::Right(
66                    std::iter::once(requirement).chain(constraints.iter().map(Cow::Borrowed)),
67                ));
68            };
69
70            // Case 3: An optional dependency with constraint(s).
71            //
72            // When the original requirement is an optional dependency, the constraint(s) need to
73            // be optional for the same extra, otherwise we activate extras that should be inactive.
74            Either::Right(Either::Left(std::iter::once(requirement).chain(
75                constraints.iter().cloned().map(move |constraint| {
76                    // Add the extra to the override marker.
77                    let mut joint_marker = MarkerTree::expression(extra_expression.clone());
78                    joint_marker.and(constraint.marker);
79                    Cow::Owned(Requirement {
80                        marker: joint_marker,
81                        ..constraint
82                    })
83                }),
84            )))
85        })
86    }
87}
uv_configuration/constraints.rs

uv_configuration/
constraints.rs
