From 41652cc9c38e7475a1fd1b2431baad0180442e2e Mon Sep 17 00:00:00 2001
From: Heath Stewart <heaths@microsoft.com>
Date: Sat, 4 May 2024 09:11:00 -0700
Subject: [PATCH] Improve color detection across platforms

Fixes #882 by improving `style::available_color_count()`:
- Checks ANSI support on Windows.
- Uses the COLORTERM environment variable and falls back to the TERM environment variable.
- Supports "xterm-24bit" and "truecolor" values which return `u16::MAX`.
---
 Cargo.toml   |   1 +
 src/style.rs | 106 +++++++++++++++++++++++++++++++++++++++++++++++++--
 2 files changed, 104 insertions(+), 3 deletions(-)

diff --git a/Cargo.toml b/Cargo.toml
index 2c53c14b..68ca5df1 100644
--- a/Cargo.toml
+++ b/Cargo.toml
@@ -89,6 +89,7 @@ futures-timer = "3.0"
 async-std = "1.12"
 serde_json = "1.0"
 serial_test = "2.0.0"
+temp-env = "0.3.6"
 
 #
 # Examples
diff --git a/src/style.rs b/src/style.rs
index a2618365..b67a421b 100644
--- a/src/style.rs
+++ b/src/style.rs
@@ -161,9 +161,23 @@ pub fn style<D: Display>(val: D) -> StyledContent<D> {
 ///
 /// This does not always provide a good result.
 pub fn available_color_count() -> u16 {
-    env::var("TERM")
-        .map(|x| if x.contains("256color") { 256 } else { 8 })
-        .unwrap_or(8)
+    #[cfg(windows)]
+    {
+        // Check if we're running in a pseudo TTY, which supports true color.
+        // Fall back to env vars otherwise for other terminals on Windows.
+        if crate::ansi_support::supports_ansi() {
+            return u16::MAX;
+        }
+    }
+
+    const DEFAULT: u16 = 8;
+    env::var("COLORTERM")
+        .or_else(|_| env::var("TERM"))
+        .map_or(DEFAULT, |x| match x {
+            _ if x.contains("24bit") || x.contains("truecolor") => u16::MAX,
+            _ if x.contains("256") => 256,
+            _ => DEFAULT,
+        })
 }
 
 /// Forces colored output on or off globally, overriding NO_COLOR.
@@ -519,3 +533,89 @@ impl_display!(for ResetColor);
 fn parse_next_u8<'a>(iter: &mut impl Iterator<Item = &'a str>) -> Option<u8> {
     iter.next().and_then(|s| s.parse().ok())
 }
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+
+    // On Windows many env var tests will fail so we need to conditionally check for ANSI support.
+    // This allows other terminals on Windows to still assert env var support.
+    macro_rules! skip_windows_ansi_supported {
+        () => {
+            #[cfg(windows)]
+            {
+                if crate::ansi_support::supports_ansi() {
+                    return;
+                }
+            }
+        };
+    }
+
+    #[cfg_attr(windows, test)]
+    #[cfg(windows)]
+    fn windows_always_truecolor() {
+        // This should always be true on supported Windows 10+,
+        // but downlevel Windows clients and other terminals may fail `cargo test` otherwise.
+        if crate::ansi_support::supports_ansi() {
+            assert_eq!(u16::MAX, available_color_count());
+        };
+    }
+
+    #[test]
+    fn colorterm_overrides_term() {
+        skip_windows_ansi_supported!();
+        temp_env::with_vars(
+            [
+                ("COLORTERM", Some("truecolor")),
+                ("TERM", Some("xterm-256color")),
+            ],
+            || {
+                assert_eq!(u16::MAX, available_color_count());
+            },
+        );
+    }
+
+    #[test]
+    fn term_24bits() {
+        skip_windows_ansi_supported!();
+        temp_env::with_vars(
+            [("COLORTERM", None), ("TERM", Some("xterm-24bits"))],
+            || {
+                assert_eq!(u16::MAX, available_color_count());
+            },
+        );
+    }
+
+    #[test]
+    fn term_256color() {
+        skip_windows_ansi_supported!();
+        temp_env::with_vars(
+            [("COLORTERM", None), ("TERM", Some("xterm-256color"))],
+            || {
+                assert_eq!(256u16, available_color_count());
+            },
+        );
+    }
+
+    #[test]
+    fn default_color_count() {
+        skip_windows_ansi_supported!();
+        temp_env::with_vars([("COLORTERM", None::<&str>), ("TERM", None)], || {
+            assert_eq!(8, available_color_count());
+        });
+    }
+
+    #[test]
+    fn unsupported_term_colorterm_values() {
+        skip_windows_ansi_supported!();
+        temp_env::with_vars(
+            [
+                ("COLORTERM", Some("gibberish")),
+                ("TERM", Some("gibberish")),
+            ],
+            || {
+                assert_eq!(8u16, available_color_count());
+            },
+        );
+    }
+}