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