Skip to main content

nils_common/
rate_limits_ansi.rs

1use crate::env as shared_env;
2use std::io::{self, IsTerminal};
3
4const CURRENT_PROFILE_FG: &str = "\x1b[38;2;199;146;234m";
5
6pub fn should_color() -> bool {
7    if shared_env::no_color_enabled() {
8        return false;
9    }
10    io::stdout().is_terminal()
11}
12
13pub fn format_name_cell(
14    raw: &str,
15    width: usize,
16    is_current: bool,
17    color_enabled: Option<bool>,
18) -> String {
19    let padded = format!("{:<width$}", raw, width = width);
20
21    let enabled = color_enabled.unwrap_or_else(should_color);
22    if !enabled || !is_current {
23        return padded;
24    }
25
26    format!("{CURRENT_PROFILE_FG}{padded}{}", reset())
27}
28
29pub fn format_percent_cell(raw: &str, width: usize, color_enabled: Option<bool>) -> String {
30    let mut trimmed = raw.to_string();
31    let raw_len = trimmed.chars().count();
32    if raw_len > width {
33        trimmed = trimmed.chars().take(width).collect();
34    }
35    let padded = format!("{:>width$}", trimmed, width = width);
36
37    let enabled = color_enabled.unwrap_or_else(should_color);
38    if !enabled {
39        return padded;
40    }
41
42    let percent = match extract_percent(raw) {
43        Some(value) => value,
44        None => return padded,
45    };
46    let color = match fg_for_percent(percent) {
47        Some(value) => value,
48        None => return padded,
49    };
50
51    format!("{color}{padded}{}", reset())
52}
53
54pub fn format_percent_token(raw: &str, color_enabled: Option<bool>) -> String {
55    let width = raw.chars().count();
56    if width == 0 {
57        return String::new();
58    }
59    format_percent_cell(raw, width, color_enabled)
60}
61
62fn extract_percent(raw: &str) -> Option<i32> {
63    let mut part = raw.rsplit(':').next()?.to_string();
64    part = part
65        .trim()
66        .trim_end_matches('%')
67        .replace(char::is_whitespace, "");
68    part.parse::<i32>().ok()
69}
70
71fn fg_for_percent(percent: i32) -> Option<String> {
72    if percent <= 0 {
73        Some(fg_truecolor(99, 119, 119))
74    } else if percent >= 80 {
75        Some(fg_truecolor(127, 219, 202))
76    } else if percent >= 60 {
77        Some(fg_truecolor(173, 219, 103))
78    } else if percent >= 40 {
79        Some(fg_truecolor(236, 196, 141))
80    } else if percent >= 20 {
81        Some(fg_truecolor(247, 140, 108))
82    } else {
83        Some(fg_truecolor(240, 113, 120))
84    }
85}
86
87fn fg_truecolor(r: u8, g: u8, b: u8) -> String {
88    format!("\x1b[38;2;{r};{g};{b}m")
89}
90
91fn reset() -> &'static str {
92    "\x1b[0m"
93}
94
95#[cfg(test)]
96mod tests {
97    use super::*;
98    use nils_test_support::{EnvGuard, GlobalStateLock};
99
100    #[test]
101    fn should_color_respects_no_color_env() {
102        let lock = GlobalStateLock::new();
103        let _no_color = EnvGuard::set(&lock, "NO_COLOR", "1");
104        assert!(!should_color());
105    }
106
107    #[test]
108    fn format_percent_cell_trims_and_keeps_non_percent_values_plain() {
109        assert_eq!(format_percent_cell("5h:94%", 8, Some(false)), "  5h:94%");
110        assert_eq!(format_percent_cell("too_long", 3, Some(false)), "too");
111        assert_eq!(format_percent_cell("oops", 4, Some(true)), "oops");
112    }
113
114    #[test]
115    fn format_percent_cell_applies_color_bands() {
116        for raw in ["x:0%", "x:80%", "x:60%", "x:40%", "x:20%", "x:19%"] {
117            let rendered = format_percent_cell(raw, raw.chars().count(), Some(true));
118            assert!(rendered.starts_with("\x1b["));
119            assert!(rendered.ends_with("\x1b[0m"));
120            assert!(rendered.contains(raw));
121        }
122    }
123
124    #[test]
125    fn format_percent_token_handles_empty_input() {
126        assert_eq!(format_percent_token("", Some(true)), "");
127    }
128
129    #[test]
130    fn format_name_cell_colors_only_current_profile() {
131        assert_eq!(
132            format_name_cell("work", 15, true, Some(false)),
133            "work           "
134        );
135        assert_eq!(
136            format_name_cell("work", 15, false, Some(true)),
137            "work           "
138        );
139
140        let rendered = format_name_cell("work", 15, true, Some(true));
141        assert!(rendered.starts_with("\x1b[38;2;199;146;234m"));
142        assert!(rendered.ends_with("\x1b[0m"));
143        assert!(rendered.contains("work"));
144    }
145}