uv_resolver/
yanks.rs

1use std::sync::Arc;
2
3use rustc_hash::{FxHashMap, FxHashSet};
4
5use uv_distribution_types::RequirementSource;
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(Arc<FxHashMap<PackageName, FxHashSet<Version>>>);
15
16impl AllowedYanks {
17    pub fn from_manifest(
18        manifest: &Manifest,
19        env: &ResolverEnvironment,
20        dependencies: DependencyMode,
21    ) -> Self {
22        let mut allowed_yanks = FxHashMap::<PackageName, FxHashSet<Version>>::default();
23
24        // Allow yanks for any pinned input requirements.
25        for requirement in manifest.requirements(env, dependencies) {
26            let RequirementSource::Registry { specifier, .. } = &requirement.source else {
27                continue;
28            };
29            let [specifier] = specifier.as_ref() else {
30                continue;
31            };
32            if matches!(
33                specifier.operator(),
34                uv_pep440::Operator::Equal | uv_pep440::Operator::ExactEqual
35            ) {
36                allowed_yanks
37                    .entry(requirement.name.clone())
38                    .or_default()
39                    .insert(specifier.version().clone());
40            }
41        }
42
43        // Allow yanks for any packages that are already pinned in the lockfile.
44        for (name, preferences) in manifest.preferences.iter() {
45            allowed_yanks
46                .entry(name.clone())
47                .or_default()
48                .extend(preferences.map(|(.., version)| version.clone()));
49        }
50
51        Self(Arc::new(allowed_yanks))
52    }
53
54    /// Returns `true` if the package-version is allowed, even if it's marked as yanked.
55    pub fn contains(&self, package_name: &PackageName, version: &Version) -> bool {
56        self.0
57            .get(package_name)
58            .is_some_and(|versions| versions.contains(version))
59    }
60}