vtcode_commons/
color_policy.rs1use once_cell::sync::Lazy;
8use std::ffi::OsString;
9use std::sync::atomic::{AtomicBool, AtomicU8, Ordering};
10
11#[derive(Clone, Copy, Debug, PartialEq, Eq)]
13pub enum ColorOutputPolicySource {
14 DefaultAuto,
16 NoColorEnv,
18 CliNoColor,
20 CliColorNever,
22 CliColorAlways,
24 ConfigOverride,
26}
27
28#[derive(Clone, Copy, Debug, PartialEq, Eq)]
30pub struct ColorOutputPolicy {
31 pub enabled: bool,
32 pub source: ColorOutputPolicySource,
33}
34
35const SOURCE_DEFAULT_AUTO: u8 = 0;
36const SOURCE_NO_COLOR_ENV: u8 = 1;
37const SOURCE_CLI_NO_COLOR: u8 = 2;
38const SOURCE_CLI_COLOR_NEVER: u8 = 3;
39const SOURCE_CLI_COLOR_ALWAYS: u8 = 4;
40const SOURCE_CONFIG_OVERRIDE: u8 = 5;
41
42static POLICY_ENABLED: AtomicBool = AtomicBool::new(true);
43static POLICY_SOURCE: AtomicU8 = AtomicU8::new(SOURCE_DEFAULT_AUTO);
44
45static INIT_FROM_ENV: Lazy<()> = Lazy::new(|| {
46 let default_policy = detect_policy_from_env();
47 set_color_output_policy(default_policy);
48});
49
50fn detect_policy_from_env() -> ColorOutputPolicy {
51 if no_color_env_active() {
52 ColorOutputPolicy {
53 enabled: false,
54 source: ColorOutputPolicySource::NoColorEnv,
55 }
56 } else {
57 ColorOutputPolicy {
58 enabled: true,
59 source: ColorOutputPolicySource::DefaultAuto,
60 }
61 }
62}
63
64fn encode_source(source: ColorOutputPolicySource) -> u8 {
65 match source {
66 ColorOutputPolicySource::DefaultAuto => SOURCE_DEFAULT_AUTO,
67 ColorOutputPolicySource::NoColorEnv => SOURCE_NO_COLOR_ENV,
68 ColorOutputPolicySource::CliNoColor => SOURCE_CLI_NO_COLOR,
69 ColorOutputPolicySource::CliColorNever => SOURCE_CLI_COLOR_NEVER,
70 ColorOutputPolicySource::CliColorAlways => SOURCE_CLI_COLOR_ALWAYS,
71 ColorOutputPolicySource::ConfigOverride => SOURCE_CONFIG_OVERRIDE,
72 }
73}
74
75fn decode_source(value: u8) -> ColorOutputPolicySource {
76 match value {
77 SOURCE_NO_COLOR_ENV => ColorOutputPolicySource::NoColorEnv,
78 SOURCE_CLI_NO_COLOR => ColorOutputPolicySource::CliNoColor,
79 SOURCE_CLI_COLOR_NEVER => ColorOutputPolicySource::CliColorNever,
80 SOURCE_CLI_COLOR_ALWAYS => ColorOutputPolicySource::CliColorAlways,
81 SOURCE_CONFIG_OVERRIDE => ColorOutputPolicySource::ConfigOverride,
82 _ => ColorOutputPolicySource::DefaultAuto,
83 }
84}
85
86fn no_color_env_active_from(value: Option<OsString>) -> bool {
87 value.map(|v| !v.is_empty()).unwrap_or(false)
88}
89
90pub fn no_color_env_active() -> bool {
92 no_color_env_active_from(std::env::var_os("NO_COLOR"))
93}
94
95pub fn current_color_output_policy() -> ColorOutputPolicy {
97 Lazy::force(&INIT_FROM_ENV);
98 ColorOutputPolicy {
99 enabled: POLICY_ENABLED.load(Ordering::Relaxed),
100 source: decode_source(POLICY_SOURCE.load(Ordering::Relaxed)),
101 }
102}
103
104pub fn set_color_output_policy(policy: ColorOutputPolicy) {
106 POLICY_ENABLED.store(policy.enabled, Ordering::Relaxed);
107 POLICY_SOURCE.store(encode_source(policy.source), Ordering::Relaxed);
108}
109
110pub fn reset_color_output_policy_from_env() {
112 set_color_output_policy(detect_policy_from_env());
113}
114
115pub fn color_output_enabled() -> bool {
117 current_color_output_policy().enabled
118}
119
120#[cfg(test)]
121mod tests {
122 use super::no_color_env_active_from;
123 use std::ffi::OsString;
124
125 #[test]
126 fn no_color_requires_non_empty_value() {
127 assert!(!no_color_env_active_from(None));
128 assert!(!no_color_env_active_from(Some(OsString::from(""))));
129 assert!(no_color_env_active_from(Some(OsString::from("1"))));
130 }
131}