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}
198
199impl PreviewFeature {
200    /// Returns the string representation of a single preview feature flag.
201    fn as_str(self) -> &'static str {
202        match self {
203            Self::PythonInstallDefault => "python-install-default",
204            Self::PythonUpgrade => "python-upgrade",
205            Self::JsonOutput => "json-output",
206            Self::Pylock => "pylock",
207            Self::AddBounds => "add-bounds",
208            Self::PackageConflicts => "package-conflicts",
209            Self::ExtraBuildDependencies => "extra-build-dependencies",
210            Self::DetectModuleConflicts => "detect-module-conflicts",
211            Self::Format => "format",
212            Self::NativeAuth => "native-auth",
213            Self::S3Endpoint => "s3-endpoint",
214            Self::CacheSize => "cache-size",
215            Self::InitProjectFlag => "init-project-flag",
216            Self::WorkspaceMetadata => "workspace-metadata",
217            Self::WorkspaceDir => "workspace-dir",
218            Self::WorkspaceList => "workspace-list",
219            Self::SbomExport => "sbom-export",
220            Self::AuthHelper => "auth-helper",
221            Self::DirectPublish => "direct-publish",
222            Self::TargetWorkspaceDiscovery => "target-workspace-discovery",
223            Self::MetadataJson => "metadata-json",
224            Self::GcsEndpoint => "gcs-endpoint",
225            Self::AdjustUlimit => "adjust-ulimit",
226            Self::SpecialCondaEnvNames => "special-conda-env-names",
227            Self::RelocatableEnvsDefault => "relocatable-envs-default",
228            Self::PublishRequireNormalized => "publish-require-normalized",
229            Self::Audit => "audit",
230            Self::ProjectDirectoryMustExist => "project-directory-must-exist",
231        }
232    }
233}
234
235impl Display for PreviewFeature {
236    fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
237        write!(f, "{}", self.as_str())
238    }
239}
240
241#[derive(Debug, Error, Clone)]
242#[error("Unknown feature flag")]
243pub struct PreviewFeatureParseError;
244
245impl FromStr for PreviewFeature {
246    type Err = PreviewFeatureParseError;
247
248    fn from_str(s: &str) -> Result<Self, Self::Err> {
249        Ok(match s {
250            "python-install-default" => Self::PythonInstallDefault,
251            "python-upgrade" => Self::PythonUpgrade,
252            "json-output" => Self::JsonOutput,
253            "pylock" => Self::Pylock,
254            "add-bounds" => Self::AddBounds,
255            "package-conflicts" => Self::PackageConflicts,
256            "extra-build-dependencies" => Self::ExtraBuildDependencies,
257            "detect-module-conflicts" => Self::DetectModuleConflicts,
258            "format" => Self::Format,
259            "native-auth" => Self::NativeAuth,
260            "s3-endpoint" => Self::S3Endpoint,
261            "gcs-endpoint" => Self::GcsEndpoint,
262            "cache-size" => Self::CacheSize,
263            "init-project-flag" => Self::InitProjectFlag,
264            "workspace-metadata" => Self::WorkspaceMetadata,
265            "workspace-dir" => Self::WorkspaceDir,
266            "workspace-list" => Self::WorkspaceList,
267            "sbom-export" => Self::SbomExport,
268            "auth-helper" => Self::AuthHelper,
269            "direct-publish" => Self::DirectPublish,
270            "target-workspace-discovery" => Self::TargetWorkspaceDiscovery,
271            "metadata-json" => Self::MetadataJson,
272            "adjust-ulimit" => Self::AdjustUlimit,
273            "special-conda-env-names" => Self::SpecialCondaEnvNames,
274            "relocatable-envs-default" => Self::RelocatableEnvsDefault,
275            "publish-require-normalized" => Self::PublishRequireNormalized,
276            "audit" => Self::Audit,
277            "project-directory-must-exist" => Self::ProjectDirectoryMustExist,
278            _ => return Err(PreviewFeatureParseError),
279        })
280    }
281}
282
283#[derive(Clone, Copy, PartialEq, Eq, Default)]
284pub struct Preview {
285    flags: BitFlags<PreviewFeature>,
286}
287
288impl Debug for Preview {
289    fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
290        let flags: Vec<_> = self.flags.iter().collect();
291        f.debug_struct("Preview").field("flags", &flags).finish()
292    }
293}
294
295impl Preview {
296    pub fn new(flags: &[PreviewFeature]) -> Self {
297        Self {
298            flags: flags.iter().copied().fold(BitFlags::empty(), BitOr::bitor),
299        }
300    }
301
302    pub fn all() -> Self {
303        Self {
304            flags: BitFlags::all(),
305        }
306    }
307
308    pub fn from_args(preview: bool, no_preview: bool, preview_features: &[PreviewFeature]) -> Self {
309        if no_preview {
310            return Self::default();
311        }
312
313        if preview {
314            return Self::all();
315        }
316
317        Self::new(preview_features)
318    }
319
320    /// Check if a single feature is enabled.
321    pub fn is_enabled(&self, flag: PreviewFeature) -> bool {
322        self.flags.contains(flag)
323    }
324
325    /// Check if all preview feature rae enabled.
326    pub fn all_enabled(&self) -> bool {
327        self.flags.is_all()
328    }
329
330    /// Check if any preview feature is enabled.
331    pub fn any_enabled(&self) -> bool {
332        !self.flags.is_empty()
333    }
334}
335
336impl Display for Preview {
337    fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
338        if self.flags.is_empty() {
339            write!(f, "disabled")
340        } else if self.flags.is_all() {
341            write!(f, "enabled")
342        } else {
343            write!(
344                f,
345                "{}",
346                itertools::join(self.flags.iter().map(PreviewFeature::as_str), ",")
347            )
348        }
349    }
350}
351
352#[derive(Debug, Error, Clone)]
353pub enum PreviewParseError {
354    #[error("Empty string in preview features: {0}")]
355    Empty(String),
356}
357
358impl FromStr for Preview {
359    type Err = PreviewParseError;
360
361    fn from_str(s: &str) -> Result<Self, Self::Err> {
362        let mut flags = BitFlags::empty();
363
364        for part in s.split(',') {
365            let part = part.trim();
366            if part.is_empty() {
367                return Err(PreviewParseError::Empty(
368                    "Empty string in preview features".to_string(),
369                ));
370            }
371
372            match PreviewFeature::from_str(part) {
373                Ok(flag) => flags |= flag,
374                Err(_) => {
375                    warn_user_once!("Unknown preview feature: `{part}`");
376                }
377            }
378        }
379
380        Ok(Self { flags })
381    }
382}
383
384#[cfg(test)]
385mod tests {
386    use super::*;
387
388    #[test]
389    fn test_preview_feature_from_str() {
390        let features = PreviewFeature::from_str("python-install-default").unwrap();
391        assert_eq!(features, PreviewFeature::PythonInstallDefault);
392    }
393
394    #[test]
395    fn test_preview_from_str() {
396        // Test single feature
397        let preview = Preview::from_str("python-install-default").unwrap();
398        assert_eq!(preview.flags, PreviewFeature::PythonInstallDefault);
399
400        // Test multiple features
401        let preview = Preview::from_str("python-upgrade,json-output").unwrap();
402        assert!(preview.is_enabled(PreviewFeature::PythonUpgrade));
403        assert!(preview.is_enabled(PreviewFeature::JsonOutput));
404        assert_eq!(preview.flags.bits().count_ones(), 2);
405
406        // Test with whitespace
407        let preview = Preview::from_str("pylock , add-bounds").unwrap();
408        assert!(preview.is_enabled(PreviewFeature::Pylock));
409        assert!(preview.is_enabled(PreviewFeature::AddBounds));
410
411        // Test empty string error
412        assert!(Preview::from_str("").is_err());
413        assert!(Preview::from_str("pylock,").is_err());
414        assert!(Preview::from_str(",pylock").is_err());
415
416        // Test unknown feature (should be ignored with warning)
417        let preview = Preview::from_str("unknown-feature,pylock").unwrap();
418        assert!(preview.is_enabled(PreviewFeature::Pylock));
419        assert_eq!(preview.flags.bits().count_ones(), 1);
420    }
421
422    #[test]
423    fn test_preview_display() {
424        // Test disabled
425        let preview = Preview::default();
426        assert_eq!(preview.to_string(), "disabled");
427        let preview = Preview::new(&[]);
428        assert_eq!(preview.to_string(), "disabled");
429
430        // Test enabled (all features)
431        let preview = Preview::all();
432        assert_eq!(preview.to_string(), "enabled");
433
434        // Test single feature
435        let preview = Preview::new(&[PreviewFeature::PythonInstallDefault]);
436        assert_eq!(preview.to_string(), "python-install-default");
437
438        // Test multiple features
439        let preview = Preview::new(&[PreviewFeature::PythonUpgrade, PreviewFeature::Pylock]);
440        assert_eq!(preview.to_string(), "python-upgrade,pylock");
441    }
442
443    #[test]
444    fn test_preview_from_args() {
445        // Test no preview and no no_preview, and no features
446        let preview = Preview::from_args(false, false, &[]);
447        assert_eq!(preview.to_string(), "disabled");
448
449        // Test no_preview
450        let preview = Preview::from_args(true, true, &[]);
451        assert_eq!(preview.to_string(), "disabled");
452
453        // Test preview (all features)
454        let preview = Preview::from_args(true, false, &[]);
455        assert_eq!(preview.to_string(), "enabled");
456
457        // Test specific features
458        let features = vec![PreviewFeature::PythonUpgrade, PreviewFeature::JsonOutput];
459        let preview = Preview::from_args(false, false, &features);
460        assert!(preview.is_enabled(PreviewFeature::PythonUpgrade));
461        assert!(preview.is_enabled(PreviewFeature::JsonOutput));
462        assert!(!preview.is_enabled(PreviewFeature::Pylock));
463    }
464
465    #[test]
466    fn test_preview_feature_as_str() {
467        assert_eq!(
468            PreviewFeature::PythonInstallDefault.as_str(),
469            "python-install-default"
470        );
471        assert_eq!(PreviewFeature::PythonUpgrade.as_str(), "python-upgrade");
472        assert_eq!(PreviewFeature::JsonOutput.as_str(), "json-output");
473        assert_eq!(PreviewFeature::Pylock.as_str(), "pylock");
474        assert_eq!(PreviewFeature::AddBounds.as_str(), "add-bounds");
475        assert_eq!(
476            PreviewFeature::PackageConflicts.as_str(),
477            "package-conflicts"
478        );
479        assert_eq!(
480            PreviewFeature::ExtraBuildDependencies.as_str(),
481            "extra-build-dependencies"
482        );
483        assert_eq!(
484            PreviewFeature::DetectModuleConflicts.as_str(),
485            "detect-module-conflicts"
486        );
487        assert_eq!(PreviewFeature::Format.as_str(), "format");
488        assert_eq!(PreviewFeature::NativeAuth.as_str(), "native-auth");
489        assert_eq!(PreviewFeature::S3Endpoint.as_str(), "s3-endpoint");
490        assert_eq!(PreviewFeature::CacheSize.as_str(), "cache-size");
491        assert_eq!(
492            PreviewFeature::InitProjectFlag.as_str(),
493            "init-project-flag"
494        );
495        assert_eq!(
496            PreviewFeature::WorkspaceMetadata.as_str(),
497            "workspace-metadata"
498        );
499        assert_eq!(PreviewFeature::WorkspaceDir.as_str(), "workspace-dir");
500        assert_eq!(PreviewFeature::WorkspaceList.as_str(), "workspace-list");
501        assert_eq!(PreviewFeature::SbomExport.as_str(), "sbom-export");
502        assert_eq!(PreviewFeature::AuthHelper.as_str(), "auth-helper");
503        assert_eq!(PreviewFeature::DirectPublish.as_str(), "direct-publish");
504        assert_eq!(
505            PreviewFeature::TargetWorkspaceDiscovery.as_str(),
506            "target-workspace-discovery"
507        );
508        assert_eq!(PreviewFeature::MetadataJson.as_str(), "metadata-json");
509        assert_eq!(PreviewFeature::GcsEndpoint.as_str(), "gcs-endpoint");
510        assert_eq!(PreviewFeature::AdjustUlimit.as_str(), "adjust-ulimit");
511        assert_eq!(
512            PreviewFeature::SpecialCondaEnvNames.as_str(),
513            "special-conda-env-names"
514        );
515        assert_eq!(
516            PreviewFeature::RelocatableEnvsDefault.as_str(),
517            "relocatable-envs-default"
518        );
519        assert_eq!(
520            PreviewFeature::PublishRequireNormalized.as_str(),
521            "publish-require-normalized"
522        );
523        assert_eq!(
524            PreviewFeature::ProjectDirectoryMustExist.as_str(),
525            "project-directory-must-exist"
526        );
527    }
528
529    #[test]
530    fn test_global_preview() {
531        {
532            let _guard =
533                test::with_features(&[PreviewFeature::Pylock, PreviewFeature::WorkspaceMetadata]);
534            assert!(!is_enabled(PreviewFeature::InitProjectFlag));
535            assert!(is_enabled(PreviewFeature::Pylock));
536            assert!(is_enabled(PreviewFeature::WorkspaceMetadata));
537            assert!(!is_enabled(PreviewFeature::AuthHelper));
538        }
539        {
540            let _guard =
541                test::with_features(&[PreviewFeature::InitProjectFlag, PreviewFeature::AuthHelper]);
542            assert!(is_enabled(PreviewFeature::InitProjectFlag));
543            assert!(!is_enabled(PreviewFeature::Pylock));
544            assert!(!is_enabled(PreviewFeature::WorkspaceMetadata));
545            assert!(is_enabled(PreviewFeature::AuthHelper));
546        }
547    }
548
549    #[test]
550    #[should_panic(
551        expected = "Additional calls to `uv_preview::test::with_features` are not allowed while holding a `FeaturesGuard`"
552    )]
553    fn test_global_preview_panic_nested() {
554        let _guard =
555            test::with_features(&[PreviewFeature::Pylock, PreviewFeature::WorkspaceMetadata]);
556        let _guard2 =
557            test::with_features(&[PreviewFeature::InitProjectFlag, PreviewFeature::AuthHelper]);
558    }
559
560    #[test]
561    #[should_panic(expected = "uv_preview::test::with_features")]
562    fn test_global_preview_panic_uninitialized() {
563        let _preview = get();
564    }
565}