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#[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 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 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 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}