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}
260
261impl PreviewFeature {
262 fn as_str(self) -> &'static str {
264 match self {
265 Self::PythonInstallDefault => "python-install-default",
266 Self::PythonUpgrade => "python-upgrade",
267 Self::JsonOutput => "json-output",
268 Self::Pylock => "pylock",
269 Self::AddBounds => "add-bounds",
270 Self::PackageConflicts => "package-conflicts",
271 Self::ExtraBuildDependencies => "extra-build-dependencies",
272 Self::DetectModuleConflicts => "detect-module-conflicts",
273 Self::Format => "format-command",
274 Self::NativeAuth => "native-auth",
275 Self::S3Endpoint => "s3-endpoint",
276 Self::CacheSize => "cache-size",
277 Self::InitProjectFlag => "init-project-flag",
278 Self::WorkspaceMetadata => "workspace-metadata",
279 Self::WorkspaceDir => "workspace-dir",
280 Self::WorkspaceList => "workspace-list",
281 Self::SbomExport => "sbom-export",
282 Self::AuthHelper => "auth-helper",
283 Self::DirectPublish => "direct-publish",
284 Self::TargetWorkspaceDiscovery => "target-workspace-discovery",
285 Self::MetadataJson => "metadata-json",
286 Self::GcsEndpoint => "gcs-endpoint",
287 Self::AdjustUlimit => "adjust-ulimit",
288 Self::SpecialCondaEnvNames => "special-conda-env-names",
289 Self::RelocatableEnvsDefault => "relocatable-envs-default",
290 Self::PublishRequireNormalized => "publish-require-normalized",
291 Self::Audit => "audit-command",
292 Self::ProjectDirectoryMustExist => "project-directory-must-exist",
293 Self::IndexExcludeNewer => "index-exclude-newer",
294 Self::AzureEndpoint => "azure-endpoint",
295 Self::TomlBackwardsCompatibility => "toml-backwards-compatibility",
296 Self::MalwareCheck => "malware-check",
297 Self::VenvSafeClear => "venv-safe-clear",
298 Self::Check => "check-command",
299 }
300 }
301}
302
303impl Display for PreviewFeature {
304 fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
305 write!(f, "{}", self.as_str())
306 }
307}
308
309#[derive(Debug, Error, Clone)]
310#[error("Unknown feature flag")]
311pub struct PreviewFeatureParseError;
312
313impl FromStr for PreviewFeature {
314 type Err = PreviewFeatureParseError;
315
316 fn from_str(s: &str) -> Result<Self, Self::Err> {
317 Ok(match s {
318 "python-install-default" => Self::PythonInstallDefault,
319 "python-upgrade" => Self::PythonUpgrade,
320 "json-output" => Self::JsonOutput,
321 "pylock" => Self::Pylock,
322 "add-bounds" => Self::AddBounds,
323 "package-conflicts" => Self::PackageConflicts,
324 "extra-build-dependencies" => Self::ExtraBuildDependencies,
325 "detect-module-conflicts" => Self::DetectModuleConflicts,
326 "format" | "format-command" => Self::Format,
327 "native-auth" => Self::NativeAuth,
328 "s3-endpoint" => Self::S3Endpoint,
329 "gcs-endpoint" => Self::GcsEndpoint,
330 "cache-size" => Self::CacheSize,
331 "init-project-flag" => Self::InitProjectFlag,
332 "workspace-metadata" => Self::WorkspaceMetadata,
333 "workspace-dir" => Self::WorkspaceDir,
334 "workspace-list" => Self::WorkspaceList,
335 "sbom-export" => Self::SbomExport,
336 "auth-helper" => Self::AuthHelper,
337 "direct-publish" => Self::DirectPublish,
338 "target-workspace-discovery" => Self::TargetWorkspaceDiscovery,
339 "metadata-json" => Self::MetadataJson,
340 "adjust-ulimit" => Self::AdjustUlimit,
341 "special-conda-env-names" => Self::SpecialCondaEnvNames,
342 "relocatable-envs-default" => Self::RelocatableEnvsDefault,
343 "publish-require-normalized" => Self::PublishRequireNormalized,
344 "audit" | "audit-command" => Self::Audit,
345 "project-directory-must-exist" => Self::ProjectDirectoryMustExist,
346 "index-exclude-newer" => Self::IndexExcludeNewer,
347 "azure-endpoint" => Self::AzureEndpoint,
348 "toml-backwards-compatibility" => Self::TomlBackwardsCompatibility,
349 "malware-check" => Self::MalwareCheck,
350 "venv-safe-clear" => Self::VenvSafeClear,
351 "check" | "check-command" => Self::Check,
352 _ => return Err(PreviewFeatureParseError),
353 })
354 }
355}
356
357#[derive(Debug, Error, Clone, Copy, PartialEq, Eq)]
358#[error("preview feature name cannot be empty")]
359pub struct EmptyPreviewFeatureNameError;
360
361#[derive(Debug, Clone)]
363pub enum MaybePreviewFeature {
364 Known(PreviewFeature),
365 Unknown(String),
366}
367
368impl FromStr for MaybePreviewFeature {
369 type Err = EmptyPreviewFeatureNameError;
370
371 fn from_str(s: &str) -> Result<Self, Self::Err> {
372 let s = s.trim();
373 if s.is_empty() {
374 return Err(EmptyPreviewFeatureNameError);
375 }
376
377 Ok(match PreviewFeature::from_str(s) {
378 Ok(feature) => Self::Known(feature),
379 Err(_) => Self::Unknown(s.to_string()),
380 })
381 }
382}
383
384#[derive(Clone, Copy, PartialEq, Eq, Default)]
385pub struct Preview {
386 flags: BitFlags<PreviewFeature>,
387}
388
389impl Debug for Preview {
390 fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
391 let flags: Vec<_> = self.flags.iter().collect();
392 f.debug_struct("Preview").field("flags", &flags).finish()
393 }
394}
395
396impl Preview {
397 pub fn new(flags: &[PreviewFeature]) -> Self {
398 Self {
399 flags: flags.iter().copied().fold(BitFlags::empty(), BitOr::bitor),
400 }
401 }
402
403 pub fn all() -> Self {
404 Self {
405 flags: BitFlags::all(),
406 }
407 }
408
409 pub fn from_args(
411 preview: bool,
412 no_preview: bool,
413 preview_features: &[MaybePreviewFeature],
414 ) -> Self {
415 if no_preview {
416 return Self::default();
417 }
418
419 if preview {
420 return Self::all();
421 }
422
423 Self::from_feature_names(preview_features)
424 }
425
426 pub fn is_enabled(&self, flag: PreviewFeature) -> bool {
428 self.flags.contains(flag)
429 }
430
431 pub fn all_enabled(&self) -> bool {
433 self.flags.is_all()
434 }
435
436 pub fn any_enabled(&self) -> bool {
438 !self.flags.is_empty()
439 }
440
441 fn from_feature_names<'a>(
442 feature_names: impl IntoIterator<Item = &'a MaybePreviewFeature>,
443 ) -> Self {
444 let mut flags = BitFlags::empty();
445
446 for feature_name in feature_names {
447 match feature_name {
448 MaybePreviewFeature::Known(feature) => flags |= *feature,
449 MaybePreviewFeature::Unknown(feature_name) => {
450 warn_user_once!("Unknown preview feature: `{feature_name}`");
451 }
452 }
453 }
454
455 Self { flags }
456 }
457}
458
459impl Display for Preview {
460 fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
461 if self.flags.is_empty() {
462 write!(f, "disabled")
463 } else if self.flags.is_all() {
464 write!(f, "enabled")
465 } else {
466 write!(
467 f,
468 "{}",
469 itertools::join(self.flags.iter().map(PreviewFeature::as_str), ",")
470 )
471 }
472 }
473}
474
475impl FromStr for Preview {
476 type Err = EmptyPreviewFeatureNameError;
477
478 fn from_str(s: &str) -> Result<Self, Self::Err> {
479 let feature_names = s
480 .split(',')
481 .map(MaybePreviewFeature::from_str)
482 .collect::<Result<Vec<_>, _>>()?;
483
484 Ok(Self::from_feature_names(&feature_names))
485 }
486}
487
488#[cfg(test)]
489mod tests {
490 use super::*;
491
492 #[test]
493 fn test_preview_feature_from_str() {
494 let features = PreviewFeature::from_str("python-install-default").unwrap();
495 assert_eq!(features, PreviewFeature::PythonInstallDefault);
496 }
497
498 #[test]
499 fn test_preview_from_str() {
500 let preview = Preview::from_str("python-install-default").unwrap();
502 assert_eq!(preview.flags, PreviewFeature::PythonInstallDefault);
503
504 let preview = Preview::from_str("python-upgrade,json-output").unwrap();
506 assert!(preview.is_enabled(PreviewFeature::PythonUpgrade));
507 assert!(preview.is_enabled(PreviewFeature::JsonOutput));
508 assert_eq!(preview.flags.bits().count_ones(), 2);
509
510 let preview = Preview::from_str("pylock , add-bounds").unwrap();
512 assert!(preview.is_enabled(PreviewFeature::Pylock));
513 assert!(preview.is_enabled(PreviewFeature::AddBounds));
514
515 assert_eq!(Preview::from_str(""), Err(EmptyPreviewFeatureNameError));
517 assert!(Preview::from_str("pylock,").is_err());
518 assert!(Preview::from_str(",pylock").is_err());
519
520 let preview = Preview::from_str("unknown-feature,pylock").unwrap();
522 assert!(preview.is_enabled(PreviewFeature::Pylock));
523 assert_eq!(preview.flags.bits().count_ones(), 1);
524 }
525
526 #[test]
527 fn test_preview_display() {
528 let preview = Preview::default();
530 assert_eq!(preview.to_string(), "disabled");
531 let preview = Preview::new(&[]);
532 assert_eq!(preview.to_string(), "disabled");
533
534 let preview = Preview::all();
536 assert_eq!(preview.to_string(), "enabled");
537
538 let preview = Preview::new(&[PreviewFeature::PythonInstallDefault]);
540 assert_eq!(preview.to_string(), "python-install-default");
541
542 let preview = Preview::new(&[PreviewFeature::PythonUpgrade, PreviewFeature::Pylock]);
544 assert_eq!(preview.to_string(), "python-upgrade,pylock");
545 }
546
547 #[test]
548 fn test_preview_from_args() {
549 let preview = Preview::from_args(false, false, &[]);
551 assert_eq!(preview.to_string(), "disabled");
552
553 let preview = Preview::from_args(true, true, &[]);
555 assert_eq!(preview.to_string(), "disabled");
556
557 let preview = Preview::from_args(true, false, &[]);
559 assert_eq!(preview.to_string(), "enabled");
560
561 let features = ["python-upgrade", "json-output"].map(|name| name.parse().unwrap());
563 let preview = Preview::from_args(false, false, &features);
564 assert!(preview.is_enabled(PreviewFeature::PythonUpgrade));
565 assert!(preview.is_enabled(PreviewFeature::JsonOutput));
566 assert!(!preview.is_enabled(PreviewFeature::Pylock));
567
568 let features = ["unknown-feature", "pylock"].map(|name| name.parse().unwrap());
570 let preview = Preview::from_args(false, false, &features);
571 assert!(preview.is_enabled(PreviewFeature::Pylock));
572 assert_eq!(preview.flags.bits().count_ones(), 1);
573 }
574
575 #[test]
576 fn test_preview_feature_as_str() {
577 assert_eq!(
578 PreviewFeature::PythonInstallDefault.as_str(),
579 "python-install-default"
580 );
581 assert_eq!(PreviewFeature::PythonUpgrade.as_str(), "python-upgrade");
582 assert_eq!(PreviewFeature::JsonOutput.as_str(), "json-output");
583 assert_eq!(PreviewFeature::Pylock.as_str(), "pylock");
584 assert_eq!(PreviewFeature::AddBounds.as_str(), "add-bounds");
585 assert_eq!(
586 PreviewFeature::PackageConflicts.as_str(),
587 "package-conflicts"
588 );
589 assert_eq!(
590 PreviewFeature::ExtraBuildDependencies.as_str(),
591 "extra-build-dependencies"
592 );
593 assert_eq!(
594 PreviewFeature::DetectModuleConflicts.as_str(),
595 "detect-module-conflicts"
596 );
597 assert_eq!(PreviewFeature::Format.as_str(), "format-command");
598 assert_eq!(PreviewFeature::NativeAuth.as_str(), "native-auth");
599 assert_eq!(PreviewFeature::S3Endpoint.as_str(), "s3-endpoint");
600 assert_eq!(PreviewFeature::CacheSize.as_str(), "cache-size");
601 assert_eq!(
602 PreviewFeature::InitProjectFlag.as_str(),
603 "init-project-flag"
604 );
605 assert_eq!(
606 PreviewFeature::WorkspaceMetadata.as_str(),
607 "workspace-metadata"
608 );
609 assert_eq!(PreviewFeature::WorkspaceDir.as_str(), "workspace-dir");
610 assert_eq!(PreviewFeature::WorkspaceList.as_str(), "workspace-list");
611 assert_eq!(PreviewFeature::SbomExport.as_str(), "sbom-export");
612 assert_eq!(PreviewFeature::AuthHelper.as_str(), "auth-helper");
613 assert_eq!(PreviewFeature::DirectPublish.as_str(), "direct-publish");
614 assert_eq!(
615 PreviewFeature::TargetWorkspaceDiscovery.as_str(),
616 "target-workspace-discovery"
617 );
618 assert_eq!(PreviewFeature::MetadataJson.as_str(), "metadata-json");
619 assert_eq!(PreviewFeature::GcsEndpoint.as_str(), "gcs-endpoint");
620 assert_eq!(PreviewFeature::AdjustUlimit.as_str(), "adjust-ulimit");
621 assert_eq!(
622 PreviewFeature::SpecialCondaEnvNames.as_str(),
623 "special-conda-env-names"
624 );
625 assert_eq!(
626 PreviewFeature::RelocatableEnvsDefault.as_str(),
627 "relocatable-envs-default"
628 );
629 assert_eq!(
630 PreviewFeature::PublishRequireNormalized.as_str(),
631 "publish-require-normalized"
632 );
633 assert_eq!(
634 PreviewFeature::ProjectDirectoryMustExist.as_str(),
635 "project-directory-must-exist"
636 );
637 assert_eq!(
638 PreviewFeature::IndexExcludeNewer.as_str(),
639 "index-exclude-newer"
640 );
641 assert_eq!(PreviewFeature::AzureEndpoint.as_str(), "azure-endpoint");
642 assert_eq!(
643 PreviewFeature::TomlBackwardsCompatibility.as_str(),
644 "toml-backwards-compatibility"
645 );
646 assert_eq!(PreviewFeature::MalwareCheck.as_str(), "malware-check");
647 assert_eq!(PreviewFeature::VenvSafeClear.as_str(), "venv-safe-clear");
648 assert_eq!(PreviewFeature::Audit.as_str(), "audit-command");
649 assert_eq!(PreviewFeature::Check.as_str(), "check-command");
650 }
651
652 #[test]
653 fn test_global_preview() {
654 {
655 let _guard =
656 test::with_features(&[PreviewFeature::Pylock, PreviewFeature::WorkspaceMetadata]);
657 assert!(!is_enabled(PreviewFeature::InitProjectFlag));
658 assert!(is_enabled(PreviewFeature::Pylock));
659 assert!(is_enabled(PreviewFeature::WorkspaceMetadata));
660 assert!(!is_enabled(PreviewFeature::AuthHelper));
661 }
662 {
663 let _guard =
664 test::with_features(&[PreviewFeature::InitProjectFlag, PreviewFeature::AuthHelper]);
665 assert!(is_enabled(PreviewFeature::InitProjectFlag));
666 assert!(!is_enabled(PreviewFeature::Pylock));
667 assert!(!is_enabled(PreviewFeature::WorkspaceMetadata));
668 assert!(is_enabled(PreviewFeature::AuthHelper));
669 }
670 }
671
672 #[test]
673 #[should_panic(
674 expected = "Additional calls to `uv_preview::test::with_features` are not allowed while holding a `FeaturesGuard`"
675 )]
676 fn test_global_preview_panic_nested() {
677 let _guard =
678 test::with_features(&[PreviewFeature::Pylock, PreviewFeature::WorkspaceMetadata]);
679 let _guard2 =
680 test::with_features(&[PreviewFeature::InitProjectFlag, PreviewFeature::AuthHelper]);
681 }
682
683 #[test]
684 #[should_panic(expected = "uv_preview::test::with_features")]
685 fn test_global_preview_panic_uninitialized() {
686 let _preview = get();
687 }
688}