use std::borrow::Cow;
#[cfg(any(test, feature = "testing"))]
use std::ops::BitOr;
use std::sync::{Mutex, OnceLock};
use std::{
fmt::{Debug, Display, Formatter},
str::FromStr,
};
use enumflags2::{BitFlags, bitflags};
use thiserror::Error;
use uv_macros::PreviewMetadata;
use uv_warnings::warn_user_once;
enum PreviewState {
Provisional(Preview),
Final(Preview),
}
enum PreviewMode {
Normal(Mutex<PreviewState>),
#[cfg(feature = "testing")]
Test(std::sync::RwLock<Option<Preview>>),
}
static PREVIEW: OnceLock<PreviewMode> = OnceLock::new();
#[derive(Debug, Error)]
pub enum PreviewError {
#[error("The preview configuration has already been finalized")]
AlreadyFinalized,
#[error("The preview configuration has not been initialized yet")]
NotInitialized,
#[cfg(feature = "testing")]
#[error("The preview configuration is in test mode and {}::{} cannot be used", module_path!(), .0)]
InTest(&'static str),
}
pub fn set(preview: Preview) -> Result<(), PreviewError> {
let mode = PREVIEW.get_or_init(|| {
PreviewMode::Normal(Mutex::new(PreviewState::Provisional(Preview::default())))
});
match mode {
PreviewMode::Normal(mutex) => {
let mut state = mutex.lock().unwrap();
match &*state {
PreviewState::Provisional(_) => {
*state = PreviewState::Provisional(preview);
Ok(())
}
PreviewState::Final(_) => Err(PreviewError::AlreadyFinalized),
}
}
#[cfg(feature = "testing")]
PreviewMode::Test(_) => Err(PreviewError::InTest("set")),
}
}
pub fn finalize() -> Result<(), PreviewError> {
match PREVIEW.get().ok_or(PreviewError::NotInitialized)? {
PreviewMode::Normal(mutex) => {
let mut state = mutex.lock().unwrap();
match &*state {
PreviewState::Provisional(preview) => {
*state = PreviewState::Final(*preview);
Ok(())
}
PreviewState::Final(_) => Err(PreviewError::AlreadyFinalized),
}
}
#[cfg(feature = "testing")]
PreviewMode::Test(_) => Err(PreviewError::InTest("finalize")),
}
}
fn get() -> Preview {
match PREVIEW.get() {
Some(PreviewMode::Normal(mutex)) => match *mutex.lock().unwrap() {
PreviewState::Provisional(preview) => preview,
PreviewState::Final(preview) => preview,
},
#[cfg(feature = "testing")]
Some(PreviewMode::Test(rwlock)) => {
assert!(
test::HELD.get(),
"The preview configuration is in test mode but the current thread does not hold a `FeaturesGuard`\nHint: Use `{}::test::with_features` to get a `FeaturesGuard` and hold it when testing functions which rely on the global preview state",
module_path!()
);
rwlock
.read()
.unwrap()
.expect("FeaturesGuard is held but preview value is not set")
}
#[cfg(feature = "testing")]
None => panic!(
"The preview configuration has not been initialized\nHint: Use `{}::init` or `{}::test::with_features` to initialize it",
module_path!(),
module_path!()
),
#[cfg(not(feature = "testing"))]
None => panic!("The preview configuration has not been initialized"),
}
}
pub fn is_enabled(flag: PreviewFeature) -> bool {
get().is_enabled(flag)
}
#[cfg(feature = "testing")]
pub mod test {
use super::{PREVIEW, Preview, PreviewMode};
use std::cell::Cell;
use std::sync::{Mutex, MutexGuard, RwLock};
static MUTEX: Mutex<()> = Mutex::new(());
thread_local! {
pub(crate) static HELD: Cell<bool> = const { Cell::new(false) };
}
#[derive(Debug)]
#[expect(unused)]
pub struct FeaturesGuard(MutexGuard<'static, ()>);
pub fn with_features(features: &[super::PreviewFeature]) -> FeaturesGuard {
assert!(
!HELD.get(),
"Additional calls to `{}::with_features` are not allowed while holding a `FeaturesGuard`",
module_path!()
);
let guard = match MUTEX.lock() {
Ok(guard) => guard,
Err(err) => err.into_inner(),
};
HELD.set(true);
let state = PREVIEW.get_or_init(|| PreviewMode::Test(RwLock::new(None)));
match state {
PreviewMode::Test(rwlock) => {
*rwlock.write().unwrap() = Some(Preview::new(features));
}
PreviewMode::Normal(_) => {
panic!(
"Cannot use `{}::with_features` after `uv_preview::init` has been called",
module_path!()
);
}
}
FeaturesGuard(guard)
}
impl Drop for FeaturesGuard {
fn drop(&mut self) {
HELD.set(false);
match PREVIEW.get().unwrap() {
PreviewMode::Test(rwlock) => {
*rwlock.write().unwrap() = None;
}
PreviewMode::Normal(_) => {
unreachable!("FeaturesGuard should not exist when in Normal mode");
}
}
}
}
}
#[bitflags]
#[expect(
clippy::use_self,
reason = "enumflags2 refers to the enum by name when inferring bits"
)]
#[repr(u64)]
#[derive(Debug, Clone, Copy, PartialEq, Eq, PreviewMetadata)]
pub enum PreviewFeature {
PythonInstallDefault,
JsonOutput,
Pylock,
AddBounds,
PackageConflicts,
ExtraBuildDependencies,
DetectModuleConflicts,
#[preview(alias = "format")]
FormatCommand,
NativeAuth,
S3Endpoint,
CacheSize,
InitProjectFlag,
WorkspaceMetadata,
WorkspaceDir,
WorkspaceList,
SbomExport,
AuthHelper,
DirectPublish,
TargetWorkspaceDiscovery,
MetadataJson,
GcsEndpoint,
AdjustUlimit,
SpecialCondaEnvNames,
RelocatableEnvsDefault,
PublishRequireNormalized,
#[preview(alias = "audit")]
AuditCommand,
ProjectDirectoryMustExist,
IndexExcludeNewer,
AzureEndpoint,
TomlBackwardsCompatibility,
MalwareCheck,
VenvSafeClear,
#[preview(alias = "check")]
CheckCommand,
PackagedInit,
CentralizedProjectEnvs,
ToolInstallLocks,
WorkspaceListScripts,
NoDistutilsPatch,
IndexHashAlgorithm,
LockfileFormatCheck,
LockWithoutMetadata,
}
impl Display for PreviewFeature {
fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
write!(f, "{}", self.as_str())
}
}
#[derive(Debug, Error, Clone)]
#[error("Unknown feature flag")]
pub struct PreviewFeatureParseError;
impl FromStr for PreviewFeature {
type Err = PreviewFeatureParseError;
fn from_str(s: &str) -> Result<Self, Self::Err> {
Self::metadata()
.iter()
.find(|(feature, _, aliases)| feature.as_str() == s || aliases.contains(&s))
.map(|(feature, _, _)| *feature)
.ok_or(PreviewFeatureParseError)
}
}
#[derive(Debug, Error, Clone, Copy, PartialEq, Eq)]
#[error("preview feature name cannot be empty")]
pub struct EmptyPreviewFeatureNameError;
#[derive(Debug, Clone)]
pub enum MaybePreviewFeature {
Known(PreviewFeature),
Unknown(String),
}
impl FromStr for MaybePreviewFeature {
type Err = EmptyPreviewFeatureNameError;
fn from_str(s: &str) -> Result<Self, Self::Err> {
let s = s.trim();
if s.is_empty() {
return Err(EmptyPreviewFeatureNameError);
}
Ok(match PreviewFeature::from_str(s) {
Ok(feature) => Self::Known(feature),
Err(_) => Self::Unknown(s.to_string()),
})
}
}
impl<'de> serde::Deserialize<'de> for MaybePreviewFeature {
fn deserialize<D>(deserializer: D) -> Result<Self, D::Error>
where
D: serde::Deserializer<'de>,
{
let name: Cow<'de, str> = serde::Deserialize::deserialize(deserializer)?;
Self::from_str(&name).map_err(serde::de::Error::custom)
}
}
#[cfg(feature = "schemars")]
impl schemars::JsonSchema for MaybePreviewFeature {
fn schema_name() -> Cow<'static, str> {
Cow::Borrowed("PreviewFeature")
}
fn json_schema(_generator: &mut schemars::SchemaGenerator) -> schemars::Schema {
let choices: Vec<&str> = BitFlags::<PreviewFeature>::all()
.iter()
.map(PreviewFeature::as_str)
.collect();
schemars::json_schema!({
"type": "string",
"anyOf": [
{
"enum": choices,
},
{
"pattern": "\\S",
},
],
})
}
}
#[derive(Clone, Copy, PartialEq, Eq, Default)]
pub struct Preview {
flags: BitFlags<PreviewFeature>,
}
impl Debug for Preview {
fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
let flags: Vec<_> = self.flags.iter().collect();
f.debug_struct("Preview").field("flags", &flags).finish()
}
}
impl Preview {
#[cfg(any(test, feature = "testing"))]
fn new(flags: &[PreviewFeature]) -> Self {
Self {
flags: flags.iter().copied().fold(BitFlags::empty(), BitOr::bitor),
}
}
pub fn all() -> Self {
Self {
flags: BitFlags::all(),
}
}
pub fn is_enabled(&self, flag: PreviewFeature) -> bool {
self.flags.contains(flag)
}
pub fn all_enabled(&self) -> bool {
self.flags.is_all()
}
pub fn any_enabled(&self) -> bool {
!self.flags.is_empty()
}
pub fn from_feature_names<'a>(
feature_names: impl IntoIterator<Item = &'a MaybePreviewFeature>,
) -> Self {
let mut flags = BitFlags::empty();
for feature_name in feature_names {
match feature_name {
MaybePreviewFeature::Known(feature) => flags |= *feature,
MaybePreviewFeature::Unknown(feature_name) => {
warn_user_once!("Unknown preview feature: `{feature_name}`");
}
}
}
Self { flags }
}
}
impl Display for Preview {
fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
if self.flags.is_empty() {
write!(f, "disabled")
} else if self.flags.is_all() {
write!(f, "enabled")
} else {
write!(
f,
"{}",
itertools::join(self.flags.iter().map(PreviewFeature::as_str), ",")
)
}
}
}
impl FromStr for Preview {
type Err = EmptyPreviewFeatureNameError;
fn from_str(s: &str) -> Result<Self, Self::Err> {
let feature_names = s
.split(',')
.map(MaybePreviewFeature::from_str)
.collect::<Result<Vec<_>, _>>()?;
Ok(Self::from_feature_names(&feature_names))
}
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn test_preview_feature_from_str() {
for &(feature, _, aliases) in PreviewFeature::metadata() {
assert_eq!(PreviewFeature::from_str(feature.as_str()).unwrap(), feature);
for &alias in aliases {
assert_eq!(PreviewFeature::from_str(alias).unwrap(), feature);
}
}
}
#[test]
fn test_preview_from_str() {
let preview = Preview::from_str("python-install-default").unwrap();
assert_eq!(preview.flags, PreviewFeature::PythonInstallDefault);
let preview = Preview::from_str("json-output,pylock").unwrap();
assert!(preview.is_enabled(PreviewFeature::JsonOutput));
assert!(preview.is_enabled(PreviewFeature::Pylock));
assert_eq!(preview.flags.bits().count_ones(), 2);
let preview = Preview::from_str("tool-install-locks").unwrap();
assert!(preview.is_enabled(PreviewFeature::ToolInstallLocks));
let preview = Preview::from_str("pylock , add-bounds").unwrap();
assert!(preview.is_enabled(PreviewFeature::Pylock));
assert!(preview.is_enabled(PreviewFeature::AddBounds));
assert_eq!(Preview::from_str(""), Err(EmptyPreviewFeatureNameError));
assert!(Preview::from_str("pylock,").is_err());
assert!(Preview::from_str(",pylock").is_err());
let preview = Preview::from_str("unknown-feature,pylock").unwrap();
assert!(preview.is_enabled(PreviewFeature::Pylock));
assert_eq!(preview.flags.bits().count_ones(), 1);
}
#[test]
fn test_preview_display() {
let preview = Preview::default();
assert_eq!(preview.to_string(), "disabled");
let preview = Preview::new(&[]);
assert_eq!(preview.to_string(), "disabled");
let preview = Preview::all();
assert_eq!(preview.to_string(), "enabled");
let preview = Preview::new(&[PreviewFeature::PythonInstallDefault]);
assert_eq!(preview.to_string(), "python-install-default");
let preview = Preview::new(&[PreviewFeature::JsonOutput, PreviewFeature::Pylock]);
assert_eq!(preview.to_string(), "json-output,pylock");
}
#[test]
fn test_global_preview() {
{
let _guard =
test::with_features(&[PreviewFeature::Pylock, PreviewFeature::WorkspaceMetadata]);
assert!(!is_enabled(PreviewFeature::InitProjectFlag));
assert!(is_enabled(PreviewFeature::Pylock));
assert!(is_enabled(PreviewFeature::WorkspaceMetadata));
assert!(!is_enabled(PreviewFeature::AuthHelper));
}
{
let _guard =
test::with_features(&[PreviewFeature::InitProjectFlag, PreviewFeature::AuthHelper]);
assert!(is_enabled(PreviewFeature::InitProjectFlag));
assert!(!is_enabled(PreviewFeature::Pylock));
assert!(!is_enabled(PreviewFeature::WorkspaceMetadata));
assert!(is_enabled(PreviewFeature::AuthHelper));
}
}
#[test]
#[should_panic(
expected = "Additional calls to `uv_preview::test::with_features` are not allowed while holding a `FeaturesGuard`"
)]
fn test_global_preview_panic_nested() {
let _guard =
test::with_features(&[PreviewFeature::Pylock, PreviewFeature::WorkspaceMetadata]);
let _guard2 =
test::with_features(&[PreviewFeature::InitProjectFlag, PreviewFeature::AuthHelper]);
}
#[test]
#[should_panic(expected = "uv_preview::test::with_features")]
fn test_global_preview_panic_uninitialized() {
let _preview = get();
}
}