use std::borrow::Cow;
use rustc_hash::FxHashMap;
use uv_configuration::{Prerelease, PrereleaseMode, specifier_opts_into_prereleases};
use uv_distribution_types::{Requirement, RequirementSource};
use uv_normalize::PackageName;
use crate::resolver::ForkSet;
use crate::{DependencyMode, Manifest, ResolverEnvironment};
#[derive(Debug, Clone)]
pub(crate) struct PrereleaseStrategy {
default: PrereleasePolicy,
package: FxHashMap<PackageName, PrereleasePolicy>,
}
#[derive(Debug, Clone)]
enum PrereleasePolicy {
Disallow,
Allow,
IfNecessary,
Explicit(ForkSet),
}
impl PrereleaseStrategy {
#[allow(deprecated)]
pub(crate) fn from_prerelease(
prerelease: &Prerelease,
manifest: &Manifest,
env: &ResolverEnvironment,
dependencies: DependencyMode,
) -> Self {
Self {
default: Self::policy(prerelease.global, manifest, env, dependencies),
package: prerelease
.package
.iter()
.map(|(name, mode)| {
(
name.clone(),
Self::policy(*mode, manifest, env, dependencies),
)
})
.collect(),
}
}
#[allow(deprecated)]
fn policy(
mode: PrereleaseMode,
manifest: &Manifest,
env: &ResolverEnvironment,
dependencies: DependencyMode,
) -> PrereleasePolicy {
match mode {
PrereleaseMode::Disallow => PrereleasePolicy::Disallow,
PrereleaseMode::Allow => PrereleasePolicy::Allow,
PrereleaseMode::IfNecessary | PrereleaseMode::IfNecessaryOrExplicit => {
PrereleasePolicy::IfNecessary
}
PrereleaseMode::Explicit => PrereleasePolicy::Explicit(Self::explicit_packages(
manifest.candidate_selection_requirements(env, dependencies),
)),
}
}
fn explicit_packages<'a>(requirements: impl Iterator<Item = Cow<'a, Requirement>>) -> ForkSet {
let mut packages = ForkSet::default();
for requirement in requirements {
let RequirementSource::Registry { specifier, .. } = &requirement.source else {
continue;
};
if specifier.iter().any(specifier_opts_into_prereleases) {
packages.add(&requirement, ());
}
}
packages
}
pub(crate) fn selection(
&self,
package_name: &PackageName,
env: &ResolverEnvironment,
) -> PrereleaseSelection {
match self.package.get(package_name).unwrap_or(&self.default) {
PrereleasePolicy::Disallow => PrereleaseSelection::Disallow,
PrereleasePolicy::Allow => PrereleaseSelection::Allow,
PrereleasePolicy::IfNecessary => PrereleaseSelection::PreferStable,
PrereleasePolicy::Explicit(packages) => {
if packages.contains(package_name, env) {
PrereleaseSelection::PreferStable
} else {
PrereleaseSelection::Disallow
}
}
}
}
}
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub(crate) enum PrereleaseSelection {
Disallow,
Allow,
PreferStable,
}