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