uv_resolver/
preferences.rs

1use std::path::Path;
2use std::str::FromStr;
3
4use rustc_hash::FxHashMap;
5use tracing::trace;
6
7use uv_distribution_types::{IndexUrl, InstalledDist, InstalledDistKind};
8use uv_normalize::PackageName;
9use uv_pep440::{Operator, Version};
10use uv_pep508::{MarkerTree, VerbatimUrl, VersionOrUrl};
11use uv_pypi_types::{HashDigest, HashDigests, HashError};
12use uv_requirements_txt::{RequirementEntry, RequirementsTxtRequirement};
13
14use crate::lock::PylockTomlPackage;
15use crate::universal_marker::UniversalMarker;
16use crate::{LockError, ResolverEnvironment};
17
18#[derive(thiserror::Error, Debug)]
19pub enum PreferenceError {
20    #[error(transparent)]
21    Hash(#[from] HashError),
22}
23
24/// A pinned requirement, as extracted from a `requirements.txt` file.
25#[derive(Clone, Debug)]
26pub struct Preference {
27    name: PackageName,
28    version: Version,
29    /// The markers on the requirement itself (those after the semicolon).
30    marker: MarkerTree,
31    /// The index URL of the package, if any.
32    index: PreferenceIndex,
33    /// If coming from a package with diverging versions, the markers of the forks this preference
34    /// is part of, otherwise `None`.
35    fork_markers: Vec<UniversalMarker>,
36    hashes: HashDigests,
37    /// The source of the preference.
38    source: PreferenceSource,
39}
40
41impl Preference {
42    /// Create a [`Preference`] from a [`RequirementEntry`].
43    pub fn from_entry(entry: RequirementEntry) -> Result<Option<Self>, PreferenceError> {
44        let RequirementsTxtRequirement::Named(requirement) = entry.requirement else {
45            return Ok(None);
46        };
47
48        let Some(VersionOrUrl::VersionSpecifier(specifier)) = requirement.version_or_url.as_ref()
49        else {
50            trace!("Excluding {requirement} from preferences due to non-version specifier");
51            return Ok(None);
52        };
53
54        let [specifier] = specifier.as_ref() else {
55            trace!("Excluding {requirement} from preferences due to multiple version specifiers");
56            return Ok(None);
57        };
58
59        if *specifier.operator() != Operator::Equal {
60            trace!("Excluding {requirement} from preferences due to inexact version specifier");
61            return Ok(None);
62        }
63
64        Ok(Some(Self {
65            name: requirement.name,
66            version: specifier.version().clone(),
67            marker: requirement.marker,
68            // `requirements.txt` doesn't have fork annotations.
69            fork_markers: vec![],
70            // `requirements.txt` doesn't allow a requirement to specify an explicit index.
71            index: PreferenceIndex::Any,
72            hashes: entry
73                .hashes
74                .iter()
75                .map(String::as_str)
76                .map(HashDigest::from_str)
77                .collect::<Result<_, _>>()?,
78            source: PreferenceSource::RequirementsTxt,
79        }))
80    }
81
82    /// Create a [`Preference`] from a locked distribution.
83    pub fn from_lock(
84        package: &crate::lock::Package,
85        install_path: &Path,
86    ) -> Result<Option<Self>, LockError> {
87        let Some(version) = package.version() else {
88            return Ok(None);
89        };
90        Ok(Some(Self {
91            name: package.id.name.clone(),
92            version: version.clone(),
93            marker: MarkerTree::TRUE,
94            index: PreferenceIndex::from(package.index(install_path)?),
95            fork_markers: package.fork_markers().to_vec(),
96            hashes: HashDigests::empty(),
97            source: PreferenceSource::Lock,
98        }))
99    }
100
101    /// Create a [`Preference`] from a locked distribution.
102    pub fn from_pylock_toml(package: &PylockTomlPackage) -> Result<Option<Self>, LockError> {
103        let Some(version) = package.version.as_ref() else {
104            return Ok(None);
105        };
106        Ok(Some(Self {
107            name: package.name.clone(),
108            version: version.clone(),
109            marker: MarkerTree::TRUE,
110            index: PreferenceIndex::from(
111                package
112                    .index
113                    .as_ref()
114                    .map(|index| IndexUrl::from(VerbatimUrl::from(index.clone()))),
115            ),
116            // `pylock.toml` doesn't have fork annotations.
117            fork_markers: vec![],
118            hashes: HashDigests::empty(),
119            source: PreferenceSource::Lock,
120        }))
121    }
122
123    /// Create a [`Preference`] from an installed distribution.
124    pub fn from_installed(dist: &InstalledDist) -> Option<Self> {
125        let InstalledDistKind::Registry(dist) = &dist.kind else {
126            return None;
127        };
128        Some(Self {
129            name: dist.name.clone(),
130            version: dist.version.clone(),
131            marker: MarkerTree::TRUE,
132            index: PreferenceIndex::Any,
133            fork_markers: vec![],
134            hashes: HashDigests::empty(),
135            source: PreferenceSource::Environment,
136        })
137    }
138
139    /// Return the [`PackageName`] of the package for this [`Preference`].
140    pub fn name(&self) -> &PackageName {
141        &self.name
142    }
143
144    /// Return the [`Version`] of the package for this [`Preference`].
145    pub fn version(&self) -> &Version {
146        &self.version
147    }
148}
149
150#[derive(Debug, Clone)]
151pub enum PreferenceIndex {
152    /// The preference should match to any index.
153    Any,
154    /// The preference should match to an implicit index.
155    Implicit,
156    /// The preference should match to a specific index.
157    Explicit(IndexUrl),
158}
159
160impl PreferenceIndex {
161    /// Returns `true` if the preference matches the given explicit [`IndexUrl`].
162    pub(crate) fn matches(&self, index: &IndexUrl) -> bool {
163        match self {
164            Self::Any => true,
165            Self::Implicit => false,
166            Self::Explicit(preference) => preference == index,
167        }
168    }
169}
170
171impl From<Option<IndexUrl>> for PreferenceIndex {
172    fn from(index: Option<IndexUrl>) -> Self {
173        match index {
174            Some(index) => Self::Explicit(index),
175            None => Self::Implicit,
176        }
177    }
178}
179
180#[derive(Debug, Clone, Copy, PartialEq, Eq)]
181pub(crate) enum PreferenceSource {
182    /// The preference is from an installed package in the environment.
183    Environment,
184    /// The preference is from a `uv.ock` file.
185    Lock,
186    /// The preference is from a `requirements.txt` file.
187    RequirementsTxt,
188    /// The preference is from the current solve.
189    Resolver,
190}
191
192#[derive(Debug, Clone)]
193pub(crate) struct Entry {
194    marker: UniversalMarker,
195    index: PreferenceIndex,
196    pin: Pin,
197    source: PreferenceSource,
198}
199
200impl Entry {
201    /// Return the [`UniversalMarker`] associated with the entry.
202    pub(crate) fn marker(&self) -> &UniversalMarker {
203        &self.marker
204    }
205
206    /// Return the [`IndexUrl`] associated with the entry, if any.
207    pub(crate) fn index(&self) -> &PreferenceIndex {
208        &self.index
209    }
210
211    /// Return the pinned data associated with the entry.
212    pub(crate) fn pin(&self) -> &Pin {
213        &self.pin
214    }
215
216    /// Return the source of the entry.
217    pub(crate) fn source(&self) -> PreferenceSource {
218        self.source
219    }
220}
221
222/// A set of pinned packages that should be preserved during resolution, if possible.
223///
224/// The marker is the marker of the fork that resolved to the pin, if any.
225///
226/// Preferences should be prioritized first by whether their marker matches and then by the order
227/// they are stored, so that a lockfile has higher precedence than sibling forks.
228#[derive(Debug, Clone, Default)]
229pub struct Preferences(FxHashMap<PackageName, Vec<Entry>>);
230
231impl Preferences {
232    /// Create a map of pinned packages from an iterator of [`Preference`] entries.
233    ///
234    /// The provided [`ResolverEnvironment`] will be used to filter the preferences
235    /// to an applicable subset.
236    pub fn from_iter(
237        preferences: impl IntoIterator<Item = Preference>,
238        env: &ResolverEnvironment,
239    ) -> Self {
240        let mut map = FxHashMap::<PackageName, Vec<_>>::default();
241        for preference in preferences {
242            // Filter non-matching preferences when resolving for an environment.
243            if let Some(markers) = env.marker_environment() {
244                if !preference.marker.evaluate(markers, &[]) {
245                    trace!("Excluding {preference} from preferences due to unmatched markers");
246                    continue;
247                }
248
249                if !preference.fork_markers.is_empty() {
250                    if !preference
251                        .fork_markers
252                        .iter()
253                        .any(|marker| marker.evaluate_no_extras(markers))
254                    {
255                        trace!(
256                            "Excluding {preference} from preferences due to unmatched fork markers"
257                        );
258                        continue;
259                    }
260                }
261            }
262
263            // Flatten the list of markers into individual entries.
264            if preference.fork_markers.is_empty() {
265                map.entry(preference.name).or_default().push(Entry {
266                    marker: UniversalMarker::TRUE,
267                    index: preference.index,
268                    pin: Pin {
269                        version: preference.version,
270                        hashes: preference.hashes,
271                    },
272                    source: preference.source,
273                });
274            } else {
275                for fork_marker in preference.fork_markers {
276                    map.entry(preference.name.clone()).or_default().push(Entry {
277                        marker: fork_marker,
278                        index: preference.index.clone(),
279                        pin: Pin {
280                            version: preference.version.clone(),
281                            hashes: preference.hashes.clone(),
282                        },
283                        source: preference.source,
284                    });
285                }
286            }
287        }
288
289        Self(map)
290    }
291
292    /// Insert a preference at the back.
293    pub(crate) fn insert(
294        &mut self,
295        package_name: PackageName,
296        index: Option<IndexUrl>,
297        markers: UniversalMarker,
298        pin: impl Into<Pin>,
299        source: PreferenceSource,
300    ) {
301        self.0.entry(package_name).or_default().push(Entry {
302            marker: markers,
303            index: PreferenceIndex::from(index),
304            pin: pin.into(),
305            source,
306        });
307    }
308
309    /// Returns an iterator over the preferences.
310    pub fn iter(
311        &self,
312    ) -> impl Iterator<
313        Item = (
314            &PackageName,
315            impl Iterator<Item = (&UniversalMarker, &PreferenceIndex, &Version)>,
316        ),
317    > {
318        self.0.iter().map(|(name, preferences)| {
319            (
320                name,
321                preferences
322                    .iter()
323                    .map(|entry| (&entry.marker, &entry.index, entry.pin.version())),
324            )
325        })
326    }
327
328    /// Return the pinned version for a package, if any.
329    pub(crate) fn get(&self, package_name: &PackageName) -> &[Entry] {
330        self.0
331            .get(package_name)
332            .map(Vec::as_slice)
333            .unwrap_or_default()
334    }
335
336    /// Return the hashes for a package, if the version matches that of the pin.
337    pub(crate) fn match_hashes(
338        &self,
339        package_name: &PackageName,
340        version: &Version,
341    ) -> Option<&[HashDigest]> {
342        self.0
343            .get(package_name)
344            .into_iter()
345            .flatten()
346            .find(|entry| entry.pin.version() == version)
347            .map(|entry| entry.pin.hashes())
348    }
349}
350
351impl std::fmt::Display for Preference {
352    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
353        write!(f, "{}=={}", self.name, self.version)
354    }
355}
356
357/// The pinned data associated with a package in a locked `requirements.txt` file (e.g., `flask==1.2.3`).
358#[derive(Debug, Clone)]
359pub(crate) struct Pin {
360    version: Version,
361    hashes: HashDigests,
362}
363
364impl Pin {
365    /// Return the version of the pinned package.
366    pub(crate) fn version(&self) -> &Version {
367        &self.version
368    }
369
370    /// Return the hashes of the pinned package.
371    pub(crate) fn hashes(&self) -> &[HashDigest] {
372        self.hashes.as_slice()
373    }
374}
375
376impl From<Version> for Pin {
377    fn from(version: Version) -> Self {
378        Self {
379            version,
380            hashes: HashDigests::empty(),
381        }
382    }
383}
uv_resolver/preferences.rs

uv_resolver/
preferences.rs
