Skip to main content

uv_resolver/
yanks.rs

1use std::sync::Arc;
2
3use rustc_hash::{FxHashMap, FxHashSet};
4
5use uv_distribution_types::{RequirementSource, ResolutionRecorder};
6use uv_normalize::PackageName;
7use uv_pep440::Version;
8
9use crate::{DependencyMode, Manifest, ResolverEnvironment};
10
11/// A set of package versions that are permitted, even if they're marked as yanked by the
12/// relevant index.
13#[derive(Debug, Default, Clone)]
14pub struct AllowedYanks {
15    versions: Arc<FxHashMap<PackageName, FxHashSet<Version>>>,
16    recorder: Option<ResolutionRecorder>,
17}
18
19impl AllowedYanks {
20    pub fn from_manifest(
21        manifest: &Manifest,
22        env: &ResolverEnvironment,
23        dependencies: DependencyMode,
24    ) -> Self {
25        let mut allowed_yanks = FxHashMap::<PackageName, FxHashSet<Version>>::default();
26
27        // Allow yanks for any pinned input requirements.
28        for requirement in manifest.candidate_selection_requirements(env, dependencies) {
29            let RequirementSource::Registry { specifier, .. } = &requirement.source else {
30                continue;
31            };
32            let [specifier] = specifier.as_ref() else {
33                continue;
34            };
35            if matches!(
36                specifier.operator(),
37                uv_pep440::Operator::Equal | uv_pep440::Operator::ExactEqual
38            ) {
39                allowed_yanks
40                    .entry(requirement.name.clone())
41                    .or_default()
42                    .insert(specifier.version().clone());
43            }
44        }
45
46        // Allow yanks for any packages that are already pinned in the lockfile.
47        for (name, preferences) in manifest.preferences.iter() {
48            allowed_yanks
49                .entry(name.clone())
50                .or_default()
51                .extend(preferences.map(|(.., version)| version.clone()));
52        }
53
54        Self {
55            versions: Arc::new(allowed_yanks),
56            recorder: manifest.recorder.clone(),
57        }
58    }
59
60    /// Returns `true` if the package-version is allowed, even if it's marked as yanked.
61    pub(crate) fn contains(&self, package_name: &PackageName, version: &Version) -> bool {
62        if let Some(recorder) = &self.recorder {
63            recorder.candidate_policy(package_name);
64        }
65        self.versions
66            .get(package_name)
67            .is_some_and(|versions| versions.contains(version))
68    }
69}