Skip to main content

uv_preview/
lib.rs

1use std::sync::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
12/// Indicates how the preview was initialised, to distinguish between normal
13/// code and unit tests.
14enum PreviewMode {
15    /// Initialised by a call to [`init`].
16    Normal(Preview),
17    /// Initialised by a call to [`test::with_features`].
18    #[cfg(feature = "testing")]
19    Test(std::sync::RwLock<Option<Preview>>),
20}
21
22static PREVIEW: OnceLock<PreviewMode> = OnceLock::new();
23
24/// Error returned when [`init`] is called more than once.
25#[derive(Debug, Error)]
26#[error("The preview configuration has already been initialized")]
27pub struct InitError;
28
29/// Initialize the global preview configuration.
30///
31/// This should be called once at startup with the resolved preview settings.
32pub fn init(preview: Preview) -> Result<(), InitError> {
33    PREVIEW
34        .set(PreviewMode::Normal(preview))
35        .map_err(|_| InitError)
36}
37
38/// Get the current global preview configuration.
39///
40/// # Panics
41///
42/// When called before [`init`] or (with the `testing` feature) when the
43/// current thread does not hold a [`test::with_features`] guard.
44pub fn get() -> Preview {
45    match PREVIEW.get() {
46        Some(PreviewMode::Normal(preview)) => *preview,
47        #[cfg(feature = "testing")]
48        Some(PreviewMode::Test(rwlock)) => {
49            assert!(
50                test::HELD.get(),
51                "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",
52                module_path!()
53            );
54            // The unwrap may panic only if the current thread had panicked
55            // while attempting to write the value and then recovered with
56            // `catch_unwind`. This seems unlikely.
57            rwlock
58                .read()
59                .unwrap()
60                .expect("FeaturesGuard is held but preview value is not set")
61        }
62        #[cfg(feature = "testing")]
63        None => panic!(
64            "The preview configuration has not been initialized\nHint: Use `{}::init` or `{}::test::with_features` to initialize it",
65            module_path!(),
66            module_path!()
67        ),
68        #[cfg(not(feature = "testing"))]
69        None => panic!("The preview configuration has not been initialized"),
70    }
71}
72
73/// Check if a specific preview feature is enabled globally.
74pub fn is_enabled(flag: PreviewFeature) -> bool {
75    get().is_enabled(flag)
76}
77
78/// Functions for unit tests, do not use from normal code!
79#[cfg(feature = "testing")]
80pub mod test {
81    use super::{PREVIEW, Preview, PreviewMode};
82    use std::cell::Cell;
83    use std::sync::{Mutex, MutexGuard, RwLock};
84
85    /// The global preview state test mutex. It does not guard any data but is
86    /// simply used to ensure tests which rely on the global preview state are
87    /// ran serially.
88    static MUTEX: Mutex<()> = Mutex::new(());
89
90    thread_local! {
91        /// Whether the current thread holds the global mutex.
92        ///
93        /// This is used to catch situations where a test forgets to set the
94        /// global test state but happens to work anyway because of another test
95        /// setting the state.
96        pub(crate) static HELD: Cell<bool> = const { Cell::new(false) };
97    }
98
99    /// A scope guard which ensures that the global preview state is configured
100    /// and consistent for the duration of its lifetime.
101    #[derive(Debug)]
102    #[expect(unused)]
103    pub struct FeaturesGuard(MutexGuard<'static, ()>);
104
105    /// Temporarily set the state of preview features for the duration of the
106    /// lifetime of the returned guard.
107    ///
108    /// Calls cannot be nested, and this function must be used to set the global
109    /// preview features when testing functionality which uses it, otherwise
110    /// that functionality will panic.
111    ///
112    /// The preview state will only be valid for the thread which calls this
113    /// function, it will not be valid for any other thread. This is a
114    /// consequence of how `HELD` is used to check for tests which are missing
115    /// the guard.
116    pub fn with_features(features: &[super::PreviewFeature]) -> FeaturesGuard {
117        assert!(
118            !HELD.get(),
119            "Additional calls to `{}::with_features` are not allowed while holding a `FeaturesGuard`",
120            module_path!()
121        );
122
123        let guard = match MUTEX.lock() {
124            Ok(guard) => guard,
125            // This is okay because the mutex isn't guarding any data, so when
126            // it gets poisoned, it just means a test thread died while holding
127            // it, so it's safe to just re-grab it from the PoisonError, there's
128            // no chance of any corruption.
129            Err(err) => err.into_inner(),
130        };
131
132        HELD.set(true);
133
134        let state = PREVIEW.get_or_init(|| PreviewMode::Test(RwLock::new(None)));
135        match state {
136            PreviewMode::Test(rwlock) => {
137                *rwlock.write().unwrap() = Some(Preview::new(features));
138            }
139            PreviewMode::Normal(_) => {
140                panic!(
141                    "Cannot use `{}::with_features` after `uv_preview::init` has been called",
142                    module_path!()
143                );
144            }
145        }
146        FeaturesGuard(guard)
147    }
148
149    impl Drop for FeaturesGuard {
150        fn drop(&mut self) {
151            HELD.set(false);
152
153            match PREVIEW.get().unwrap() {
154                PreviewMode::Test(rwlock) => {
155                    *rwlock.write().unwrap() = None;
156                }
157                PreviewMode::Normal(_) => {
158                    unreachable!("FeaturesGuard should not exist when in Normal mode");
159                }
160            }
161        }
162    }
163}
164
165#[bitflags]
166#[repr(u32)]
167#[derive(Debug, Clone, Copy, PartialEq, Eq)]
168pub enum PreviewFeature {
169    PythonInstallDefault = 1 << 0,
170    PythonUpgrade = 1 << 1,
171    JsonOutput = 1 << 2,
172    Pylock = 1 << 3,
173    AddBounds = 1 << 4,
174    PackageConflicts = 1 << 5,
175    ExtraBuildDependencies = 1 << 6,
176    DetectModuleConflicts = 1 << 7,
177    Format = 1 << 8,
178    NativeAuth = 1 << 9,
179    S3Endpoint = 1 << 10,
180    CacheSize = 1 << 11,
181    InitProjectFlag = 1 << 12,
182    WorkspaceMetadata = 1 << 13,
183    WorkspaceDir = 1 << 14,
184    WorkspaceList = 1 << 15,
185    SbomExport = 1 << 16,
186    AuthHelper = 1 << 17,
187    DirectPublish = 1 << 18,
188    TargetWorkspaceDiscovery = 1 << 19,
189    MetadataJson = 1 << 20,
190    GcsEndpoint = 1 << 21,
191    AdjustUlimit = 1 << 22,
192    SpecialCondaEnvNames = 1 << 23,
193    RelocatableEnvsDefault = 1 << 24,
194    PublishRequireNormalized = 1 << 25,
195    Audit = 1 << 26,
196    ProjectDirectoryMustExist = 1 << 27,
197    IndexExcludeNewer = 1 << 28,
198}
199
200impl PreviewFeature {
201    /// Returns the string representation of a single preview feature flag.
202    fn as_str(self) -> &'static str {
203        match self {
204            Self::PythonInstallDefault => "python-install-default",
205            Self::PythonUpgrade => "python-upgrade",
206            Self::JsonOutput => "json-output",
207            Self::Pylock => "pylock",
208            Self::AddBounds => "add-bounds",
209            Self::PackageConflicts => "package-conflicts",
210            Self::ExtraBuildDependencies => "extra-build-dependencies",
211            Self::DetectModuleConflicts => "detect-module-conflicts",
212            Self::Format => "format",
213            Self::NativeAuth => "native-auth",
214            Self::S3Endpoint => "s3-endpoint",
215            Self::CacheSize => "cache-size",
216            Self::InitProjectFlag => "init-project-flag",
217            Self::WorkspaceMetadata => "workspace-metadata",
218            Self::WorkspaceDir => "workspace-dir",
219            Self::WorkspaceList => "workspace-list",
220            Self::SbomExport => "sbom-export",
221            Self::AuthHelper => "auth-helper",
222            Self::DirectPublish => "direct-publish",
223            Self::TargetWorkspaceDiscovery => "target-workspace-discovery",
224            Self::MetadataJson => "metadata-json",
225            Self::GcsEndpoint => "gcs-endpoint",
226            Self::AdjustUlimit => "adjust-ulimit",
227            Self::SpecialCondaEnvNames => "special-conda-env-names",
228            Self::RelocatableEnvsDefault => "relocatable-envs-default",
229            Self::PublishRequireNormalized => "publish-require-normalized",
230            Self::Audit => "audit",
231            Self::ProjectDirectoryMustExist => "project-directory-must-exist",
232            Self::IndexExcludeNewer => "index-exclude-newer",
233        }
234    }
235}
236
237impl Display for PreviewFeature {
238    fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
239        write!(f, "{}", self.as_str())
240    }
241}
242
243#[derive(Debug, Error, Clone)]
244#[error("Unknown feature flag")]
245pub struct PreviewFeatureParseError;
246
247impl FromStr for PreviewFeature {
248    type Err = PreviewFeatureParseError;
249
250    fn from_str(s: &str) -> Result<Self, Self::Err> {
251        Ok(match s {
252            "python-install-default" => Self::PythonInstallDefault,
253            "python-upgrade" => Self::PythonUpgrade,
254            "json-output" => Self::JsonOutput,
255            "pylock" => Self::Pylock,
256            "add-bounds" => Self::AddBounds,
257            "package-conflicts" => Self::PackageConflicts,
258            "extra-build-dependencies" => Self::ExtraBuildDependencies,
259            "detect-module-conflicts" => Self::DetectModuleConflicts,
260            "format" => Self::Format,
261            "native-auth" => Self::NativeAuth,
262            "s3-endpoint" => Self::S3Endpoint,
263            "gcs-endpoint" => Self::GcsEndpoint,
264            "cache-size" => Self::CacheSize,
265            "init-project-flag" => Self::InitProjectFlag,
266            "workspace-metadata" => Self::WorkspaceMetadata,
267            "workspace-dir" => Self::WorkspaceDir,
268            "workspace-list" => Self::WorkspaceList,
269            "sbom-export" => Self::SbomExport,
270            "auth-helper" => Self::AuthHelper,
271            "direct-publish" => Self::DirectPublish,
272            "target-workspace-discovery" => Self::TargetWorkspaceDiscovery,
273            "metadata-json" => Self::MetadataJson,
274            "adjust-ulimit" => Self::AdjustUlimit,
275            "special-conda-env-names" => Self::SpecialCondaEnvNames,
276            "relocatable-envs-default" => Self::RelocatableEnvsDefault,
277            "publish-require-normalized" => Self::PublishRequireNormalized,
278            "audit" => Self::Audit,
279            "project-directory-must-exist" => Self::ProjectDirectoryMustExist,
280            "index-exclude-newer" => Self::IndexExcludeNewer,
281            _ => return Err(PreviewFeatureParseError),
282        })
283    }
284}
285
286#[derive(Clone, Copy, PartialEq, Eq, Default)]
287pub struct Preview {
288    flags: BitFlags<PreviewFeature>,
289}
290
291impl Debug for Preview {
292    fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
293        let flags: Vec<_> = self.flags.iter().collect();
294        f.debug_struct("Preview").field("flags", &flags).finish()
295    }
296}
297
298impl Preview {
299    pub fn new(flags: &[PreviewFeature]) -> Self {
300        Self {
301            flags: flags.iter().copied().fold(BitFlags::empty(), BitOr::bitor),
302        }
303    }
304
305    pub fn all() -> Self {
306        Self {
307            flags: BitFlags::all(),
308        }
309    }
310
311    pub fn from_args(preview: bool, no_preview: bool, preview_features: &[PreviewFeature]) -> Self {
312        if no_preview {
313            return Self::default();
314        }
315
316        if preview {
317            return Self::all();
318        }
319
320        Self::new(preview_features)
321    }
322
323    /// Check if a single feature is enabled.
324    pub fn is_enabled(&self, flag: PreviewFeature) -> bool {
325        self.flags.contains(flag)
326    }
327
328    /// Check if all preview feature rae enabled.
329    pub fn all_enabled(&self) -> bool {
330        self.flags.is_all()
331    }
332
333    /// Check if any preview feature is enabled.
334    pub fn any_enabled(&self) -> bool {
335        !self.flags.is_empty()
336    }
337}
338
339impl Display for Preview {
340    fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
341        if self.flags.is_empty() {
342            write!(f, "disabled")
343        } else if self.flags.is_all() {
344            write!(f, "enabled")
345        } else {
346            write!(
347                f,
348                "{}",
349                itertools::join(self.flags.iter().map(PreviewFeature::as_str), ",")
350            )
351        }
352    }
353}
354
355#[derive(Debug, Error, Clone)]
356pub enum PreviewParseError {
357    #[error("Empty string in preview features: {0}")]
358    Empty(String),
359}
360
361impl FromStr for Preview {
362    type Err = PreviewParseError;
363
364    fn from_str(s: &str) -> Result<Self, Self::Err> {
365        let mut flags = BitFlags::empty();
366
367        for part in s.split(',') {
368            let part = part.trim();
369            if part.is_empty() {
370                return Err(PreviewParseError::Empty(
371                    "Empty string in preview features".to_string(),
372                ));
373            }
374
375            match PreviewFeature::from_str(part) {
376                Ok(flag) => flags |= flag,
377                Err(_) => {
378                    warn_user_once!("Unknown preview feature: `{part}`");
379                }
380            }
381        }
382
383        Ok(Self { flags })
384    }
385}
386
387#[cfg(test)]
388mod tests {
389    use super::*;
390
391    #[test]
392    fn test_preview_feature_from_str() {
393        let features = PreviewFeature::from_str("python-install-default").unwrap();
394        assert_eq!(features, PreviewFeature::PythonInstallDefault);
395    }
396
397    #[test]
398    fn test_preview_from_str() {
399        // Test single feature
400        let preview = Preview::from_str("python-install-default").unwrap();
401        assert_eq!(preview.flags, PreviewFeature::PythonInstallDefault);
402
403        // Test multiple features
404        let preview = Preview::from_str("python-upgrade,json-output").unwrap();
405        assert!(preview.is_enabled(PreviewFeature::PythonUpgrade));
406        assert!(preview.is_enabled(PreviewFeature::JsonOutput));
407        assert_eq!(preview.flags.bits().count_ones(), 2);
408
409        // Test with whitespace
410        let preview = Preview::from_str("pylock , add-bounds").unwrap();
411        assert!(preview.is_enabled(PreviewFeature::Pylock));
412        assert!(preview.is_enabled(PreviewFeature::AddBounds));
413
414        // Test empty string error
415        assert!(Preview::from_str("").is_err());
416        assert!(Preview::from_str("pylock,").is_err());
417        assert!(Preview::from_str(",pylock").is_err());
418
419        // Test unknown feature (should be ignored with warning)
420        let preview = Preview::from_str("unknown-feature,pylock").unwrap();
421        assert!(preview.is_enabled(PreviewFeature::Pylock));
422        assert_eq!(preview.flags.bits().count_ones(), 1);
423    }
424
425    #[test]
426    fn test_preview_display() {
427        // Test disabled
428        let preview = Preview::default();
429        assert_eq!(preview.to_string(), "disabled");
430        let preview = Preview::new(&[]);
431        assert_eq!(preview.to_string(), "disabled");
432
433        // Test enabled (all features)
434        let preview = Preview::all();
435        assert_eq!(preview.to_string(), "enabled");
436
437        // Test single feature
438        let preview = Preview::new(&[PreviewFeature::PythonInstallDefault]);
439        assert_eq!(preview.to_string(), "python-install-default");
440
441        // Test multiple features
442        let preview = Preview::new(&[PreviewFeature::PythonUpgrade, PreviewFeature::Pylock]);
443        assert_eq!(preview.to_string(), "python-upgrade,pylock");
444    }
445
446    #[test]
447    fn test_preview_from_args() {
448        // Test no preview and no no_preview, and no features
449        let preview = Preview::from_args(false, false, &[]);
450        assert_eq!(preview.to_string(), "disabled");
451
452        // Test no_preview
453        let preview = Preview::from_args(true, true, &[]);
454        assert_eq!(preview.to_string(), "disabled");
455
456        // Test preview (all features)
457        let preview = Preview::from_args(true, false, &[]);
458        assert_eq!(preview.to_string(), "enabled");
459
460        // Test specific features
461        let features = vec![PreviewFeature::PythonUpgrade, PreviewFeature::JsonOutput];
462        let preview = Preview::from_args(false, false, &features);
463        assert!(preview.is_enabled(PreviewFeature::PythonUpgrade));
464        assert!(preview.is_enabled(PreviewFeature::JsonOutput));
465        assert!(!preview.is_enabled(PreviewFeature::Pylock));
466    }
467
468    #[test]
469    fn test_preview_feature_as_str() {
470        assert_eq!(
471            PreviewFeature::PythonInstallDefault.as_str(),
472            "python-install-default"
473        );
474        assert_eq!(PreviewFeature::PythonUpgrade.as_str(), "python-upgrade");
475        assert_eq!(PreviewFeature::JsonOutput.as_str(), "json-output");
476        assert_eq!(PreviewFeature::Pylock.as_str(), "pylock");
477        assert_eq!(PreviewFeature::AddBounds.as_str(), "add-bounds");
478        assert_eq!(
479            PreviewFeature::PackageConflicts.as_str(),
480            "package-conflicts"
481        );
482        assert_eq!(
483            PreviewFeature::ExtraBuildDependencies.as_str(),
484            "extra-build-dependencies"
485        );
486        assert_eq!(
487            PreviewFeature::DetectModuleConflicts.as_str(),
488            "detect-module-conflicts"
489        );
490        assert_eq!(PreviewFeature::Format.as_str(), "format");
491        assert_eq!(PreviewFeature::NativeAuth.as_str(), "native-auth");
492        assert_eq!(PreviewFeature::S3Endpoint.as_str(), "s3-endpoint");
493        assert_eq!(PreviewFeature::CacheSize.as_str(), "cache-size");
494        assert_eq!(
495            PreviewFeature::InitProjectFlag.as_str(),
496            "init-project-flag"
497        );
498        assert_eq!(
499            PreviewFeature::WorkspaceMetadata.as_str(),
500            "workspace-metadata"
501        );
502        assert_eq!(PreviewFeature::WorkspaceDir.as_str(), "workspace-dir");
503        assert_eq!(PreviewFeature::WorkspaceList.as_str(), "workspace-list");
504        assert_eq!(PreviewFeature::SbomExport.as_str(), "sbom-export");
505        assert_eq!(PreviewFeature::AuthHelper.as_str(), "auth-helper");
506        assert_eq!(PreviewFeature::DirectPublish.as_str(), "direct-publish");
507        assert_eq!(
508            PreviewFeature::TargetWorkspaceDiscovery.as_str(),
509            "target-workspace-discovery"
510        );
511        assert_eq!(PreviewFeature::MetadataJson.as_str(), "metadata-json");
512        assert_eq!(PreviewFeature::GcsEndpoint.as_str(), "gcs-endpoint");
513        assert_eq!(PreviewFeature::AdjustUlimit.as_str(), "adjust-ulimit");
514        assert_eq!(
515            PreviewFeature::SpecialCondaEnvNames.as_str(),
516            "special-conda-env-names"
517        );
518        assert_eq!(
519            PreviewFeature::RelocatableEnvsDefault.as_str(),
520            "relocatable-envs-default"
521        );
522        assert_eq!(
523            PreviewFeature::PublishRequireNormalized.as_str(),
524            "publish-require-normalized"
525        );
526        assert_eq!(
527            PreviewFeature::ProjectDirectoryMustExist.as_str(),
528            "project-directory-must-exist"
529        );
530        assert_eq!(
531            PreviewFeature::IndexExcludeNewer.as_str(),
532            "index-exclude-newer"
533        );
534    }
535
536    #[test]
537    fn test_global_preview() {
538        {
539            let _guard =
540                test::with_features(&[PreviewFeature::Pylock, PreviewFeature::WorkspaceMetadata]);
541            assert!(!is_enabled(PreviewFeature::InitProjectFlag));
542            assert!(is_enabled(PreviewFeature::Pylock));
543            assert!(is_enabled(PreviewFeature::WorkspaceMetadata));
544            assert!(!is_enabled(PreviewFeature::AuthHelper));
545        }
546        {
547            let _guard =
548                test::with_features(&[PreviewFeature::InitProjectFlag, PreviewFeature::AuthHelper]);
549            assert!(is_enabled(PreviewFeature::InitProjectFlag));
550            assert!(!is_enabled(PreviewFeature::Pylock));
551            assert!(!is_enabled(PreviewFeature::WorkspaceMetadata));
552            assert!(is_enabled(PreviewFeature::AuthHelper));
553        }
554    }
555
556    #[test]
557    #[should_panic(
558        expected = "Additional calls to `uv_preview::test::with_features` are not allowed while holding a `FeaturesGuard`"
559    )]
560    fn test_global_preview_panic_nested() {
561        let _guard =
562            test::with_features(&[PreviewFeature::Pylock, PreviewFeature::WorkspaceMetadata]);
563        let _guard2 =
564            test::with_features(&[PreviewFeature::InitProjectFlag, PreviewFeature::AuthHelper]);
565    }
566
567    #[test]
568    #[should_panic(expected = "uv_preview::test::with_features")]
569    fn test_global_preview_panic_uninitialized() {
570        let _preview = get();
571    }
572}