uv_configuration/
overrides.rs

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

uv_configuration/
overrides.rs
