1use std::borrow::Cow;
2use std::sync::{Mutex, OnceLock};
3use std::{
4 fmt::{Debug, Display, Formatter},
5 ops::BitOr,
6 str::FromStr,
7};
8
9use enumflags2::{BitFlags, bitflags};
10use thiserror::Error;
11use uv_warnings::warn_user_once;
12
13enum PreviewState {
15 Provisional(Preview),
16 Final(Preview),
17}
18
19enum PreviewMode {
22 Normal(Mutex<PreviewState>),
24 #[cfg(feature = "testing")]
26 Test(std::sync::RwLock<Option<Preview>>),
27}
28
29static PREVIEW: OnceLock<PreviewMode> = OnceLock::new();
30
31#[derive(Debug, Error)]
33pub enum PreviewError {
34 #[error("The preview configuration has already been finalized")]
36 AlreadyFinalized,
37
38 #[error("The preview configuration has not been initialized yet")]
40 NotInitialized,
41
42 #[cfg(feature = "testing")]
44 #[error("The preview configuration is in test mode and {}::{} cannot be used", module_path!(), .0)]
45 InTest(&'static str),
46}
47
48pub fn set(preview: Preview) -> Result<(), PreviewError> {
52 let mode = PREVIEW.get_or_init(|| {
53 PreviewMode::Normal(Mutex::new(PreviewState::Provisional(Preview::default())))
54 });
55 match mode {
56 PreviewMode::Normal(mutex) => {
57 let mut state = mutex.lock().unwrap();
60 match &*state {
61 PreviewState::Provisional(_) => {
62 *state = PreviewState::Provisional(preview);
63 Ok(())
64 }
65 PreviewState::Final(_) => Err(PreviewError::AlreadyFinalized),
66 }
67 }
68 #[cfg(feature = "testing")]
69 PreviewMode::Test(_) => Err(PreviewError::InTest("set")),
70 }
71}
72
73pub fn finalize() -> Result<(), PreviewError> {
74 match PREVIEW.get().ok_or(PreviewError::NotInitialized)? {
75 PreviewMode::Normal(mutex) => {
76 let mut state = mutex.lock().unwrap();
79 match &*state {
80 PreviewState::Provisional(preview) => {
81 *state = PreviewState::Final(*preview);
82 Ok(())
83 }
84 PreviewState::Final(_) => Err(PreviewError::AlreadyFinalized),
85 }
86 }
87 #[cfg(feature = "testing")]
88 PreviewMode::Test(_) => Err(PreviewError::InTest("finalize")),
89 }
90}
91
92fn get() -> Preview {
99 match PREVIEW.get() {
100 Some(PreviewMode::Normal(mutex)) => match *mutex.lock().unwrap() {
101 PreviewState::Provisional(preview) => preview,
102 PreviewState::Final(preview) => preview,
103 },
104 #[cfg(feature = "testing")]
105 Some(PreviewMode::Test(rwlock)) => {
106 assert!(
107 test::HELD.get(),
108 "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",
109 module_path!()
110 );
111 rwlock
115 .read()
116 .unwrap()
117 .expect("FeaturesGuard is held but preview value is not set")
118 }
119 #[cfg(feature = "testing")]
120 None => panic!(
121 "The preview configuration has not been initialized\nHint: Use `{}::init` or `{}::test::with_features` to initialize it",
122 module_path!(),
123 module_path!()
124 ),
125 #[cfg(not(feature = "testing"))]
126 None => panic!("The preview configuration has not been initialized"),
127 }
128}
129
130pub fn is_enabled(flag: PreviewFeature) -> bool {
132 get().is_enabled(flag)
133}
134
135#[cfg(feature = "testing")]
137pub mod test {
138 use super::{PREVIEW, Preview, PreviewMode};
139 use std::cell::Cell;
140 use std::sync::{Mutex, MutexGuard, RwLock};
141
142 static MUTEX: Mutex<()> = Mutex::new(());
146
147 thread_local! {
148 pub(crate) static HELD: Cell<bool> = const { Cell::new(false) };
154 }
155
156 #[derive(Debug)]
159 #[expect(unused)]
160 pub struct FeaturesGuard(MutexGuard<'static, ()>);
161
162 pub fn with_features(features: &[super::PreviewFeature]) -> FeaturesGuard {
174 assert!(
175 !HELD.get(),
176 "Additional calls to `{}::with_features` are not allowed while holding a `FeaturesGuard`",
177 module_path!()
178 );
179
180 let guard = match MUTEX.lock() {
181 Ok(guard) => guard,
182 Err(err) => err.into_inner(),
187 };
188
189 HELD.set(true);
190
191 let state = PREVIEW.get_or_init(|| PreviewMode::Test(RwLock::new(None)));
192 match state {
193 PreviewMode::Test(rwlock) => {
194 *rwlock.write().unwrap() = Some(Preview::new(features));
195 }
196 PreviewMode::Normal(_) => {
197 panic!(
198 "Cannot use `{}::with_features` after `uv_preview::init` has been called",
199 module_path!()
200 );
201 }
202 }
203 FeaturesGuard(guard)
204 }
205
206 impl Drop for FeaturesGuard {
207 fn drop(&mut self) {
208 HELD.set(false);
209
210 match PREVIEW.get().unwrap() {
211 PreviewMode::Test(rwlock) => {
212 *rwlock.write().unwrap() = None;
213 }
214 PreviewMode::Normal(_) => {
215 unreachable!("FeaturesGuard should not exist when in Normal mode");
216 }
217 }
218 }
219 }
220}
221
222#[bitflags]
223#[repr(u64)]
224#[derive(Debug, Clone, Copy, PartialEq, Eq)]
225pub enum PreviewFeature {
226 PythonInstallDefault = 1 << 0,
227 PythonUpgrade = 1 << 1,
228 JsonOutput = 1 << 2,
229 Pylock = 1 << 3,
230 AddBounds = 1 << 4,
231 PackageConflicts = 1 << 5,
232 ExtraBuildDependencies = 1 << 6,
233 DetectModuleConflicts = 1 << 7,
234 Format = 1 << 8,
235 NativeAuth = 1 << 9,
236 S3Endpoint = 1 << 10,
237 CacheSize = 1 << 11,
238 InitProjectFlag = 1 << 12,
239 WorkspaceMetadata = 1 << 13,
240 WorkspaceDir = 1 << 14,
241 WorkspaceList = 1 << 15,
242 SbomExport = 1 << 16,
243 AuthHelper = 1 << 17,
244 DirectPublish = 1 << 18,
245 TargetWorkspaceDiscovery = 1 << 19,
246 MetadataJson = 1 << 20,
247 GcsEndpoint = 1 << 21,
248 AdjustUlimit = 1 << 22,
249 SpecialCondaEnvNames = 1 << 23,
250 RelocatableEnvsDefault = 1 << 24,
251 PublishRequireNormalized = 1 << 25,
252 Audit = 1 << 26,
253 ProjectDirectoryMustExist = 1 << 27,
254 IndexExcludeNewer = 1 << 28,
255 AzureEndpoint = 1 << 29,
256 TomlBackwardsCompatibility = 1 << 30,
257 MalwareCheck = 1 << 31,
258 VenvSafeClear = 1 << 32,
259 Check = 1 << 33,
260 PackagedInit = 1 << 34,
261}
262
263impl PreviewFeature {
264 fn as_str(self) -> &'static str {
266 match self {
267 Self::PythonInstallDefault => "python-install-default",
268 Self::PythonUpgrade => "python-upgrade",
269 Self::JsonOutput => "json-output",
270 Self::Pylock => "pylock",
271 Self::AddBounds => "add-bounds",
272 Self::PackageConflicts => "package-conflicts",
273 Self::ExtraBuildDependencies => "extra-build-dependencies",
274 Self::DetectModuleConflicts => "detect-module-conflicts",
275 Self::Format => "format-command",
276 Self::NativeAuth => "native-auth",
277 Self::S3Endpoint => "s3-endpoint",
278 Self::CacheSize => "cache-size",
279 Self::InitProjectFlag => "init-project-flag",
280 Self::WorkspaceMetadata => "workspace-metadata",
281 Self::WorkspaceDir => "workspace-dir",
282 Self::WorkspaceList => "workspace-list",
283 Self::SbomExport => "sbom-export",
284 Self::AuthHelper => "auth-helper",
285 Self::DirectPublish => "direct-publish",
286 Self::TargetWorkspaceDiscovery => "target-workspace-discovery",
287 Self::MetadataJson => "metadata-json",
288 Self::GcsEndpoint => "gcs-endpoint",
289 Self::AdjustUlimit => "adjust-ulimit",
290 Self::SpecialCondaEnvNames => "special-conda-env-names",
291 Self::RelocatableEnvsDefault => "relocatable-envs-default",
292 Self::PublishRequireNormalized => "publish-require-normalized",
293 Self::Audit => "audit-command",
294 Self::ProjectDirectoryMustExist => "project-directory-must-exist",
295 Self::IndexExcludeNewer => "index-exclude-newer",
296 Self::AzureEndpoint => "azure-endpoint",
297 Self::TomlBackwardsCompatibility => "toml-backwards-compatibility",
298 Self::MalwareCheck => "malware-check",
299 Self::VenvSafeClear => "venv-safe-clear",
300 Self::Check => "check-command",
301 Self::PackagedInit => "packaged-init",
302 }
303 }
304}
305
306impl Display for PreviewFeature {
307 fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
308 write!(f, "{}", self.as_str())
309 }
310}
311
312#[derive(Debug, Error, Clone)]
313#[error("Unknown feature flag")]
314pub struct PreviewFeatureParseError;
315
316impl FromStr for PreviewFeature {
317 type Err = PreviewFeatureParseError;
318
319 fn from_str(s: &str) -> Result<Self, Self::Err> {
320 Ok(match s {
321 "python-install-default" => Self::PythonInstallDefault,
322 "python-upgrade" => Self::PythonUpgrade,
323 "json-output" => Self::JsonOutput,
324 "pylock" => Self::Pylock,
325 "add-bounds" => Self::AddBounds,
326 "package-conflicts" => Self::PackageConflicts,
327 "extra-build-dependencies" => Self::ExtraBuildDependencies,
328 "detect-module-conflicts" => Self::DetectModuleConflicts,
329 "format" | "format-command" => Self::Format,
330 "native-auth" => Self::NativeAuth,
331 "s3-endpoint" => Self::S3Endpoint,
332 "gcs-endpoint" => Self::GcsEndpoint,
333 "cache-size" => Self::CacheSize,
334 "init-project-flag" => Self::InitProjectFlag,
335 "workspace-metadata" => Self::WorkspaceMetadata,
336 "workspace-dir" => Self::WorkspaceDir,
337 "workspace-list" => Self::WorkspaceList,
338 "sbom-export" => Self::SbomExport,
339 "auth-helper" => Self::AuthHelper,
340 "direct-publish" => Self::DirectPublish,
341 "target-workspace-discovery" => Self::TargetWorkspaceDiscovery,
342 "metadata-json" => Self::MetadataJson,
343 "adjust-ulimit" => Self::AdjustUlimit,
344 "special-conda-env-names" => Self::SpecialCondaEnvNames,
345 "relocatable-envs-default" => Self::RelocatableEnvsDefault,
346 "publish-require-normalized" => Self::PublishRequireNormalized,
347 "audit" | "audit-command" => Self::Audit,
348 "project-directory-must-exist" => Self::ProjectDirectoryMustExist,
349 "index-exclude-newer" => Self::IndexExcludeNewer,
350 "azure-endpoint" => Self::AzureEndpoint,
351 "toml-backwards-compatibility" => Self::TomlBackwardsCompatibility,
352 "malware-check" => Self::MalwareCheck,
353 "venv-safe-clear" => Self::VenvSafeClear,
354 "check" | "check-command" => Self::Check,
355 "packaged-init" => Self::PackagedInit,
356 _ => return Err(PreviewFeatureParseError),
357 })
358 }
359}
360
361#[derive(Debug, Error, Clone, Copy, PartialEq, Eq)]
362#[error("preview feature name cannot be empty")]
363pub struct EmptyPreviewFeatureNameError;
364
365#[derive(Debug, Clone)]
367pub enum MaybePreviewFeature {
368 Known(PreviewFeature),
369 Unknown(String),
370}
371
372impl FromStr for MaybePreviewFeature {
373 type Err = EmptyPreviewFeatureNameError;
374
375 fn from_str(s: &str) -> Result<Self, Self::Err> {
376 let s = s.trim();
377 if s.is_empty() {
378 return Err(EmptyPreviewFeatureNameError);
379 }
380
381 Ok(match PreviewFeature::from_str(s) {
382 Ok(feature) => Self::Known(feature),
383 Err(_) => Self::Unknown(s.to_string()),
384 })
385 }
386}
387
388impl<'de> serde::Deserialize<'de> for MaybePreviewFeature {
389 fn deserialize<D>(deserializer: D) -> Result<Self, D::Error>
390 where
391 D: serde::Deserializer<'de>,
392 {
393 let name: Cow<'de, str> = serde::Deserialize::deserialize(deserializer)?;
394 Self::from_str(&name).map_err(serde::de::Error::custom)
395 }
396}
397
398#[cfg(feature = "schemars")]
399impl schemars::JsonSchema for MaybePreviewFeature {
400 fn schema_name() -> Cow<'static, str> {
401 Cow::Borrowed("PreviewFeature")
402 }
403
404 fn json_schema(_generator: &mut schemars::SchemaGenerator) -> schemars::Schema {
405 let choices: Vec<&str> = BitFlags::<PreviewFeature>::all()
408 .iter()
409 .map(PreviewFeature::as_str)
410 .collect();
411 schemars::json_schema!({
412 "type": "string",
413 "anyOf": [
414 {
415 "enum": choices,
416 },
417 {
418 "pattern": "\\S",
419 },
420 ],
421 })
422 }
423}
424
425#[derive(Clone, Copy, PartialEq, Eq, Default)]
426pub struct Preview {
427 flags: BitFlags<PreviewFeature>,
428}
429
430impl Debug for Preview {
431 fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
432 let flags: Vec<_> = self.flags.iter().collect();
433 f.debug_struct("Preview").field("flags", &flags).finish()
434 }
435}
436
437impl Preview {
438 pub fn new(flags: &[PreviewFeature]) -> Self {
439 Self {
440 flags: flags.iter().copied().fold(BitFlags::empty(), BitOr::bitor),
441 }
442 }
443
444 pub fn all() -> Self {
445 Self {
446 flags: BitFlags::all(),
447 }
448 }
449
450 pub fn is_enabled(&self, flag: PreviewFeature) -> bool {
452 self.flags.contains(flag)
453 }
454
455 pub fn all_enabled(&self) -> bool {
457 self.flags.is_all()
458 }
459
460 pub fn any_enabled(&self) -> bool {
462 !self.flags.is_empty()
463 }
464
465 pub fn from_feature_names<'a>(
467 feature_names: impl IntoIterator<Item = &'a MaybePreviewFeature>,
468 ) -> Self {
469 let mut flags = BitFlags::empty();
470
471 for feature_name in feature_names {
472 match feature_name {
473 MaybePreviewFeature::Known(feature) => flags |= *feature,
474 MaybePreviewFeature::Unknown(feature_name) => {
475 warn_user_once!("Unknown preview feature: `{feature_name}`");
476 }
477 }
478 }
479
480 Self { flags }
481 }
482}
483
484impl Display for Preview {
485 fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
486 if self.flags.is_empty() {
487 write!(f, "disabled")
488 } else if self.flags.is_all() {
489 write!(f, "enabled")
490 } else {
491 write!(
492 f,
493 "{}",
494 itertools::join(self.flags.iter().map(PreviewFeature::as_str), ",")
495 )
496 }
497 }
498}
499
500impl FromStr for Preview {
501 type Err = EmptyPreviewFeatureNameError;
502
503 fn from_str(s: &str) -> Result<Self, Self::Err> {
504 let feature_names = s
505 .split(',')
506 .map(MaybePreviewFeature::from_str)
507 .collect::<Result<Vec<_>, _>>()?;
508
509 Ok(Self::from_feature_names(&feature_names))
510 }
511}
512
513#[cfg(test)]
514mod tests {
515 use super::*;
516
517 #[test]
518 fn test_preview_feature_from_str() {
519 let features = PreviewFeature::from_str("python-install-default").unwrap();
520 assert_eq!(features, PreviewFeature::PythonInstallDefault);
521 }
522
523 #[test]
524 fn test_preview_from_str() {
525 let preview = Preview::from_str("python-install-default").unwrap();
527 assert_eq!(preview.flags, PreviewFeature::PythonInstallDefault);
528
529 let preview = Preview::from_str("python-upgrade,json-output").unwrap();
531 assert!(preview.is_enabled(PreviewFeature::PythonUpgrade));
532 assert!(preview.is_enabled(PreviewFeature::JsonOutput));
533 assert_eq!(preview.flags.bits().count_ones(), 2);
534
535 let preview = Preview::from_str("pylock , add-bounds").unwrap();
537 assert!(preview.is_enabled(PreviewFeature::Pylock));
538 assert!(preview.is_enabled(PreviewFeature::AddBounds));
539
540 assert_eq!(Preview::from_str(""), Err(EmptyPreviewFeatureNameError));
542 assert!(Preview::from_str("pylock,").is_err());
543 assert!(Preview::from_str(",pylock").is_err());
544
545 let preview = Preview::from_str("unknown-feature,pylock").unwrap();
547 assert!(preview.is_enabled(PreviewFeature::Pylock));
548 assert_eq!(preview.flags.bits().count_ones(), 1);
549 }
550
551 #[test]
552 fn test_preview_display() {
553 let preview = Preview::default();
555 assert_eq!(preview.to_string(), "disabled");
556 let preview = Preview::new(&[]);
557 assert_eq!(preview.to_string(), "disabled");
558
559 let preview = Preview::all();
561 assert_eq!(preview.to_string(), "enabled");
562
563 let preview = Preview::new(&[PreviewFeature::PythonInstallDefault]);
565 assert_eq!(preview.to_string(), "python-install-default");
566
567 let preview = Preview::new(&[PreviewFeature::PythonUpgrade, PreviewFeature::Pylock]);
569 assert_eq!(preview.to_string(), "python-upgrade,pylock");
570 }
571
572 #[test]
573 fn test_preview_feature_as_str() {
574 assert_eq!(
575 PreviewFeature::PythonInstallDefault.as_str(),
576 "python-install-default"
577 );
578 assert_eq!(PreviewFeature::PythonUpgrade.as_str(), "python-upgrade");
579 assert_eq!(PreviewFeature::JsonOutput.as_str(), "json-output");
580 assert_eq!(PreviewFeature::Pylock.as_str(), "pylock");
581 assert_eq!(PreviewFeature::AddBounds.as_str(), "add-bounds");
582 assert_eq!(
583 PreviewFeature::PackageConflicts.as_str(),
584 "package-conflicts"
585 );
586 assert_eq!(
587 PreviewFeature::ExtraBuildDependencies.as_str(),
588 "extra-build-dependencies"
589 );
590 assert_eq!(
591 PreviewFeature::DetectModuleConflicts.as_str(),
592 "detect-module-conflicts"
593 );
594 assert_eq!(PreviewFeature::Format.as_str(), "format-command");
595 assert_eq!(PreviewFeature::NativeAuth.as_str(), "native-auth");
596 assert_eq!(PreviewFeature::S3Endpoint.as_str(), "s3-endpoint");
597 assert_eq!(PreviewFeature::CacheSize.as_str(), "cache-size");
598 assert_eq!(
599 PreviewFeature::InitProjectFlag.as_str(),
600 "init-project-flag"
601 );
602 assert_eq!(
603 PreviewFeature::WorkspaceMetadata.as_str(),
604 "workspace-metadata"
605 );
606 assert_eq!(PreviewFeature::WorkspaceDir.as_str(), "workspace-dir");
607 assert_eq!(PreviewFeature::WorkspaceList.as_str(), "workspace-list");
608 assert_eq!(PreviewFeature::SbomExport.as_str(), "sbom-export");
609 assert_eq!(PreviewFeature::AuthHelper.as_str(), "auth-helper");
610 assert_eq!(PreviewFeature::DirectPublish.as_str(), "direct-publish");
611 assert_eq!(
612 PreviewFeature::TargetWorkspaceDiscovery.as_str(),
613 "target-workspace-discovery"
614 );
615 assert_eq!(PreviewFeature::MetadataJson.as_str(), "metadata-json");
616 assert_eq!(PreviewFeature::GcsEndpoint.as_str(), "gcs-endpoint");
617 assert_eq!(PreviewFeature::AdjustUlimit.as_str(), "adjust-ulimit");
618 assert_eq!(
619 PreviewFeature::SpecialCondaEnvNames.as_str(),
620 "special-conda-env-names"
621 );
622 assert_eq!(
623 PreviewFeature::RelocatableEnvsDefault.as_str(),
624 "relocatable-envs-default"
625 );
626 assert_eq!(
627 PreviewFeature::PublishRequireNormalized.as_str(),
628 "publish-require-normalized"
629 );
630 assert_eq!(
631 PreviewFeature::ProjectDirectoryMustExist.as_str(),
632 "project-directory-must-exist"
633 );
634 assert_eq!(
635 PreviewFeature::IndexExcludeNewer.as_str(),
636 "index-exclude-newer"
637 );
638 assert_eq!(PreviewFeature::AzureEndpoint.as_str(), "azure-endpoint");
639 assert_eq!(
640 PreviewFeature::TomlBackwardsCompatibility.as_str(),
641 "toml-backwards-compatibility"
642 );
643 assert_eq!(PreviewFeature::MalwareCheck.as_str(), "malware-check");
644 assert_eq!(PreviewFeature::VenvSafeClear.as_str(), "venv-safe-clear");
645 assert_eq!(PreviewFeature::Audit.as_str(), "audit-command");
646 assert_eq!(PreviewFeature::Check.as_str(), "check-command");
647 }
648
649 #[test]
650 fn test_global_preview() {
651 {
652 let _guard =
653 test::with_features(&[PreviewFeature::Pylock, PreviewFeature::WorkspaceMetadata]);
654 assert!(!is_enabled(PreviewFeature::InitProjectFlag));
655 assert!(is_enabled(PreviewFeature::Pylock));
656 assert!(is_enabled(PreviewFeature::WorkspaceMetadata));
657 assert!(!is_enabled(PreviewFeature::AuthHelper));
658 }
659 {
660 let _guard =
661 test::with_features(&[PreviewFeature::InitProjectFlag, PreviewFeature::AuthHelper]);
662 assert!(is_enabled(PreviewFeature::InitProjectFlag));
663 assert!(!is_enabled(PreviewFeature::Pylock));
664 assert!(!is_enabled(PreviewFeature::WorkspaceMetadata));
665 assert!(is_enabled(PreviewFeature::AuthHelper));
666 }
667 }
668
669 #[test]
670 #[should_panic(
671 expected = "Additional calls to `uv_preview::test::with_features` are not allowed while holding a `FeaturesGuard`"
672 )]
673 fn test_global_preview_panic_nested() {
674 let _guard =
675 test::with_features(&[PreviewFeature::Pylock, PreviewFeature::WorkspaceMetadata]);
676 let _guard2 =
677 test::with_features(&[PreviewFeature::InitProjectFlag, PreviewFeature::AuthHelper]);
678 }
679
680 #[test]
681 #[should_panic(expected = "uv_preview::test::with_features")]
682 fn test_global_preview_panic_uninitialized() {
683 let _preview = get();
684 }
685}