|
| 1 | +use std::env; |
| 2 | +use std::sync::LazyLock; |
| 3 | + |
| 4 | +#[derive(Debug, PartialEq, Clone, Copy)] |
| 5 | +pub enum ColorLevel { |
| 6 | + None, |
| 7 | + Basic, |
| 8 | + Enhanced, |
| 9 | + TrueColor, |
| 10 | +} |
| 11 | + |
| 12 | +impl ColorLevel { |
| 13 | + pub fn to_bool(&self) -> bool { |
| 14 | + matches!(self, ColorLevel::Basic | ColorLevel::Enhanced | ColorLevel::TrueColor) |
| 15 | + } |
| 16 | + |
| 17 | + pub fn new() -> Self { |
| 18 | + Self::None |
| 19 | + } |
| 20 | + |
| 21 | + pub fn update(self, new_level: Self) -> Self { |
| 22 | + match (&self, new_level) { |
| 23 | + (ColorLevel::None, level) => level, |
| 24 | + (_, level) if level > self => level, |
| 25 | + _ => self, |
| 26 | + } |
| 27 | + } |
| 28 | + |
| 29 | + pub fn force_none() -> Self { |
| 30 | + Self::None |
| 31 | + } |
| 32 | +} |
| 33 | + |
| 34 | +impl Default for ColorLevel { |
| 35 | + fn default() -> Self { |
| 36 | + Self::None |
| 37 | + } |
| 38 | +} |
| 39 | + |
| 40 | +impl PartialOrd for ColorLevel { |
| 41 | + fn partial_cmp(&self, other: &Self) -> Option<std::cmp::Ordering> { |
| 42 | + Some(match (self, other) { |
| 43 | + (Self::None, _) => std::cmp::Ordering::Less, |
| 44 | + (_, Self::None) => std::cmp::Ordering::Greater, |
| 45 | + (Self::Basic, Self::Basic) => std::cmp::Ordering::Equal, |
| 46 | + (Self::Basic, Self::Enhanced) => std::cmp::Ordering::Less, |
| 47 | + (Self::Basic, Self::TrueColor) => std::cmp::Ordering::Less, |
| 48 | + (Self::Enhanced, Self::Basic) => std::cmp::Ordering::Greater, |
| 49 | + (Self::Enhanced, Self::Enhanced) => std::cmp::Ordering::Equal, |
| 50 | + (Self::Enhanced, Self::TrueColor) => std::cmp::Ordering::Less, |
| 51 | + (Self::TrueColor, Self::Basic) => std::cmp::Ordering::Greater, |
| 52 | + (Self::TrueColor, Self::Enhanced) => std::cmp::Ordering::Greater, |
| 53 | + (Self::TrueColor, Self::TrueColor) => std::cmp::Ordering::Equal, |
| 54 | + }) |
| 55 | + } |
| 56 | +} |
| 57 | + |
| 58 | +static COLOR_LEVEL: LazyLock<ColorLevel> = LazyLock::new(determine_color_level); |
| 59 | + |
| 60 | +pub fn supports_color() -> bool { |
| 61 | + COLOR_LEVEL.to_bool() |
| 62 | +} |
| 63 | + |
| 64 | +fn determine_color_level() -> ColorLevel { |
| 65 | + |
| 66 | + // Check FORCE_COLOR environment variable first |
| 67 | + if let Some(color) = force_color_check() { |
| 68 | + return color; |
| 69 | + } |
| 70 | + |
| 71 | + // Windows checks |
| 72 | + if let Some(color) = windows_color_check() { |
| 73 | + return color; |
| 74 | + } |
| 75 | + |
| 76 | + // CI check |
| 77 | + if let Some(color) = ci_color_check() { |
| 78 | + return color; |
| 79 | + } |
| 80 | + |
| 81 | + // Term check |
| 82 | + if let Some(color) = term_color_check() { |
| 83 | + return color; |
| 84 | + } |
| 85 | + |
| 86 | + // Term Program check |
| 87 | + if let Some(color) = term_program_check() { |
| 88 | + return color; |
| 89 | + } |
| 90 | + |
| 91 | + // Final fallback based on whether the stream is a TTY |
| 92 | + if let Some(color) = tty_color_check() { |
| 93 | + return color; |
| 94 | + } |
| 95 | + |
| 96 | + ColorLevel::None |
| 97 | +} |
| 98 | + |
| 99 | +fn force_color_check() -> Option<ColorLevel> { |
| 100 | + env::var("FORCE_COLOR").ok().and_then(|val| match val.as_str() { |
| 101 | + "0" => Some(ColorLevel::None), |
| 102 | + "1" => Some(ColorLevel::Basic), |
| 103 | + "2" => Some(ColorLevel::Enhanced), |
| 104 | + "3" => Some(ColorLevel::TrueColor), |
| 105 | + _ => None, |
| 106 | + }) |
| 107 | +} |
| 108 | + |
| 109 | +fn windows_color_check() -> Option<ColorLevel> { |
| 110 | + if std::env::consts::OS == "windows" { |
| 111 | + return Some(ColorLevel::Basic); |
| 112 | + } |
| 113 | + None |
| 114 | +} |
| 115 | + |
| 116 | +fn ci_color_check() -> Option<ColorLevel> { |
| 117 | + if env::var("CI").is_ok() { |
| 118 | + if env::var("GITHUB_ACTIONS").is_ok() || env::var("GITEA_ACTIONS").is_ok() { |
| 119 | + return Some(ColorLevel::TrueColor); |
| 120 | + } |
| 121 | + let ci_providers = ["TRAVIS", "CIRCLECI", "APPVEYOR", "GITLAB_CI", "BUILDKITE", "DRONE", "codeship"]; |
| 122 | + if ci_providers.iter().any(|ci| env::var(ci).is_ok()) { |
| 123 | + return Some(ColorLevel::Basic); |
| 124 | + } |
| 125 | + } |
| 126 | + None |
| 127 | +} |
| 128 | + |
| 129 | +fn term_color_check() -> Option<ColorLevel> { |
| 130 | + match env::var("TERM").as_deref() { |
| 131 | + Ok("dumb") => None, |
| 132 | + Ok("xterm-kitty") | Ok("truecolor") | Ok("ansi") => Some(ColorLevel::TrueColor), |
| 133 | + Ok(term) if term.ends_with("-256color") => Some(ColorLevel::Enhanced), |
| 134 | + Ok(term) if term.starts_with("xterm") || term.starts_with("screen") => Some(ColorLevel::Basic), |
| 135 | + Ok("cygwin") => Some(ColorLevel::Basic), |
| 136 | + _ => None, |
| 137 | + } |
| 138 | +} |
| 139 | + |
| 140 | +fn term_program_check() -> Option<ColorLevel> { |
| 141 | + match env::var("TERM_PROGRAM").as_deref() { |
| 142 | + Ok("iTerm.app") => Some(ColorLevel::TrueColor), |
| 143 | + Ok("Apple_Terminal") => Some(ColorLevel::Enhanced), |
| 144 | + _ => None, |
| 145 | + } |
| 146 | +} |
| 147 | + |
| 148 | +fn tty_color_check() -> Option<ColorLevel> { |
| 149 | + if atty::is(atty::Stream::Stdout) { |
| 150 | + return Some(ColorLevel::TrueColor); |
| 151 | + } |
| 152 | + None |
| 153 | +} |
| 154 | + |
| 155 | +#[cfg(test)] |
| 156 | +mod tests { |
| 157 | + use super::*; |
| 158 | + use std::env; |
| 159 | + |
| 160 | + // test to_bool |
| 161 | + #[test] |
| 162 | + fn test_to_bool() { |
| 163 | + assert!(!ColorLevel::None.to_bool(), "None should not be true"); |
| 164 | + assert!(ColorLevel::Basic.to_bool(), "Basic should be true"); |
| 165 | + assert!(ColorLevel::Enhanced.to_bool(), "Enhanced should be true"); |
| 166 | + assert!(ColorLevel::TrueColor.to_bool(), "TrueColor should be true"); |
| 167 | + } |
| 168 | + |
| 169 | + // test new |
| 170 | + #[test] |
| 171 | + fn test_new() { |
| 172 | + assert_eq!(ColorLevel::None, ColorLevel::new(), "None should be the default color level"); |
| 173 | + } |
| 174 | + |
| 175 | + // test update |
| 176 | + #[test] |
| 177 | + fn test_update() { |
| 178 | + assert_eq!(ColorLevel::Basic, ColorLevel::None.update(ColorLevel::Basic), "Updating from None to Basic should change the color level"); |
| 179 | + assert_eq!(ColorLevel::Enhanced, ColorLevel::Basic.update(ColorLevel::Enhanced), "Updating from Basic to Enhanced should change the color level"); |
| 180 | + assert_eq!(ColorLevel::TrueColor, ColorLevel::Enhanced.update(ColorLevel::TrueColor), "Updating from Enhanced to TrueColor should change the color level"); |
| 181 | + // should never downgrade |
| 182 | + // updating to none should not change the color level |
| 183 | + assert_eq!(ColorLevel::Basic, ColorLevel::Basic.update(ColorLevel::None), "Updating from Basic to None should not change the color level"); |
| 184 | + assert_eq!(ColorLevel::Enhanced, ColorLevel::Enhanced.update(ColorLevel::None), "Updating from Enhanced to None should not change the color level"); |
| 185 | + assert_eq!(ColorLevel::TrueColor, ColorLevel::TrueColor.update(ColorLevel::None), "Updating from TrueColor to None should not change the color level"); |
| 186 | + // updating to a lower level should not change the color level |
| 187 | + assert_eq!(ColorLevel::Enhanced, ColorLevel::Enhanced.update(ColorLevel::Basic), "Updating from Enhanced to Basic should not change the color level"); |
| 188 | + assert_eq!(ColorLevel::TrueColor, ColorLevel::TrueColor.update(ColorLevel::Enhanced), "Updating from TrueColor to Enhanced should not change the color level"); |
| 189 | + |
| 190 | + } |
| 191 | + |
| 192 | + #[test] |
| 193 | + fn test_force_color_levels() { |
| 194 | + use std::env; |
| 195 | + |
| 196 | + // Define test cases as tuples: (env_var_value, expected_color_level) |
| 197 | + let test_cases = [ |
| 198 | + ("0", ColorLevel::None), |
| 199 | + ("1", ColorLevel::Basic), |
| 200 | + ("2", ColorLevel::Enhanced), |
| 201 | + ("3", ColorLevel::TrueColor), |
| 202 | + // Add a case for an invalid test that should fail |
| 203 | + ("0", ColorLevel::Basic), // this case is intentionally wrong |
| 204 | + ]; |
| 205 | + |
| 206 | + for (value, expected) in test_cases.iter() { |
| 207 | + env::set_var("FORCE_COLOR", value); |
| 208 | + let color_level = determine_color_level(); |
| 209 | + |
| 210 | + // If the expected color level is ColorLevel::Basic, this case should fail |
| 211 | + if *value == "0" && *expected == ColorLevel::Basic { |
| 212 | + assert_ne!(color_level, *expected, "Expected failure for FORCE_COLOR = {}", value); |
| 213 | + } else { |
| 214 | + assert_eq!(color_level, *expected, "Unexpected color level for FORCE_COLOR = {}", value); |
| 215 | + } |
| 216 | + |
| 217 | + env::remove_var("FORCE_COLOR"); |
| 218 | + } |
| 219 | + } |
| 220 | + |
| 221 | + // Test CI color level |
| 222 | + #[test] |
| 223 | + fn test_ci_color_levels() { |
| 224 | + // Define a list of CI providers and their expected ColorLevel |
| 225 | + let ci_providers = [ |
| 226 | + ("GITHUB_ACTIONS", ColorLevel::TrueColor), |
| 227 | + ("GITEA_ACTIONS", ColorLevel::TrueColor), |
| 228 | + ("TRAVIS", ColorLevel::Basic), |
| 229 | + ("CIRCLECI", ColorLevel::Basic), |
| 230 | + ("APPVEYOR", ColorLevel::Basic), |
| 231 | + ("GITLAB_CI", ColorLevel::Basic), |
| 232 | + ("BUILDKITE", ColorLevel::Basic), |
| 233 | + ("DRONE", ColorLevel::Basic), |
| 234 | + ("codeship", ColorLevel::Basic), |
| 235 | + ]; |
| 236 | + |
| 237 | + // Outer loop: Set the main "CI" environment variable to enable CI mode |
| 238 | + env::set_var("CI", "true"); |
| 239 | + |
| 240 | + for (provider, expected_level) in ci_providers.iter() { |
| 241 | + // Set the specific CI provider environment variable |
| 242 | + env::set_var(provider, "true"); |
| 243 | + |
| 244 | + // Check if the color level matches the expected level |
| 245 | + assert_eq!( |
| 246 | + &ci_color_check().unwrap_or(ColorLevel::None), |
| 247 | + expected_level, |
| 248 | + "Unexpected color level for CI provider {}", |
| 249 | + provider |
| 250 | + ); |
| 251 | + |
| 252 | + // Clean up by removing the provider-specific environment variable |
| 253 | + env::remove_var(provider); |
| 254 | + } |
| 255 | + |
| 256 | + // Remove the "CI" environment variable to clean up |
| 257 | + env::remove_var("CI"); |
| 258 | + } |
| 259 | + |
| 260 | +} |
0 commit comments