bssh/ui/
color.rs

1// Copyright 2025 Lablup Inc. and Jeongkyu Shin
2//
3// Licensed under the Apache License, Version 2.0 (the "License");
4// you may not use this file except in compliance with the License.
5// You may obtain a copy of the License at
6//
7//     http://www.apache.org/licenses/LICENSE-2.0
8//
9// Unless required by applicable law or agreed to in writing, software
10// distributed under the License is distributed on an "AS IS" BASIS,
11// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12// See the License for the specific language governing permissions and
13// limitations under the License.
14
15//! Process-wide color policy with independent stdout and stderr detection.
16
17use std::ffi::OsStr;
18use std::fmt;
19use std::io::IsTerminal;
20use std::sync::atomic::{AtomicU8, Ordering};
21
22use clap::ValueEnum;
23
24#[derive(Debug, Clone, Copy, PartialEq, Eq, ValueEnum, Default)]
25pub enum ColorMode {
26    #[default]
27    Auto,
28    Always,
29    Never,
30}
31
32#[derive(Debug, Clone, Copy, PartialEq, Eq)]
33pub enum OutputStream {
34    Stdout,
35    Stderr,
36}
37
38static COLOR_MODE: AtomicU8 = AtomicU8::new(ColorMode::Auto as u8);
39
40pub fn configure_color(mode: ColorMode) {
41    COLOR_MODE.store(mode as u8, Ordering::Relaxed);
42}
43
44fn configured_mode() -> ColorMode {
45    match COLOR_MODE.load(Ordering::Relaxed) {
46        value if value == ColorMode::Always as u8 => ColorMode::Always,
47        value if value == ColorMode::Never as u8 => ColorMode::Never,
48        _ => ColorMode::Auto,
49    }
50}
51
52fn auto_color_enabled(is_terminal: bool, no_color: Option<&OsStr>, term: Option<&OsStr>) -> bool {
53    is_terminal
54        && !no_color.is_some_and(|value| !value.is_empty())
55        && !term.is_some_and(|value| value.eq_ignore_ascii_case(OsStr::new("dumb")))
56}
57
58fn color_enabled_for(
59    mode: ColorMode,
60    is_terminal: bool,
61    no_color: Option<&OsStr>,
62    term: Option<&OsStr>,
63) -> bool {
64    match mode {
65        ColorMode::Auto => auto_color_enabled(is_terminal, no_color, term),
66        ColorMode::Always => true,
67        ColorMode::Never => false,
68    }
69}
70
71pub fn colors_enabled(stream: OutputStream) -> bool {
72    let is_terminal = match stream {
73        OutputStream::Stdout => std::io::stdout().is_terminal(),
74        OutputStream::Stderr => std::io::stderr().is_terminal(),
75    };
76    color_enabled_for(
77        configured_mode(),
78        is_terminal,
79        std::env::var_os("NO_COLOR").as_deref(),
80        std::env::var_os("TERM").as_deref(),
81    )
82}
83
84pub trait Colorize: fmt::Display {
85    fn styled_for(&self, ansi_code: &'static str, stream: OutputStream) -> StyledText {
86        StyledText {
87            text: self.to_string(),
88            ansi_code,
89            stream,
90        }
91    }
92    fn styled(&self, ansi_code: &'static str) -> StyledText {
93        self.styled_for(ansi_code, OutputStream::Stdout)
94    }
95    fn red(&self) -> StyledText {
96        self.styled("31")
97    }
98    fn red_stderr(&self) -> StyledText {
99        self.styled_for("31", OutputStream::Stderr)
100    }
101    fn green(&self) -> StyledText {
102        self.styled("32")
103    }
104    fn yellow(&self) -> StyledText {
105        self.styled("33")
106    }
107    fn blue(&self) -> StyledText {
108        self.styled("34")
109    }
110    fn cyan(&self) -> StyledText {
111        self.styled("36")
112    }
113    fn bright_blue(&self) -> StyledText {
114        self.styled("94")
115    }
116    fn bold(&self) -> StyledText {
117        self.styled("1")
118    }
119    fn dimmed(&self) -> StyledText {
120        self.styled("2")
121    }
122}
123
124impl<T: fmt::Display> Colorize for T {}
125
126#[derive(Debug, Clone, PartialEq, Eq)]
127pub struct StyledText {
128    text: String,
129    ansi_code: &'static str,
130    stream: OutputStream,
131}
132
133impl fmt::Display for StyledText {
134    fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result {
135        if colors_enabled(self.stream) {
136            write!(formatter, "\x1b[{}m{}\x1b[0m", self.ansi_code, self.text)
137        } else {
138            formatter.write_str(&self.text)
139        }
140    }
141}
142
143#[cfg(test)]
144mod tests {
145    use super::*;
146
147    #[test]
148    fn auto_requires_a_terminal() {
149        assert!(!color_enabled_for(ColorMode::Auto, false, None, None));
150        assert!(color_enabled_for(ColorMode::Auto, true, None, None));
151    }
152
153    #[test]
154    fn auto_honors_non_empty_no_color_and_dumb_term() {
155        assert!(!color_enabled_for(
156            ColorMode::Auto,
157            true,
158            Some(OsStr::new("1")),
159            None
160        ));
161        assert!(color_enabled_for(
162            ColorMode::Auto,
163            true,
164            Some(OsStr::new("")),
165            None
166        ));
167        assert!(!color_enabled_for(
168            ColorMode::Auto,
169            true,
170            None,
171            Some(OsStr::new("dumb"))
172        ));
173    }
174
175    #[test]
176    fn explicit_modes_override_environment_and_terminal_detection() {
177        assert!(color_enabled_for(
178            ColorMode::Always,
179            false,
180            Some(OsStr::new("1")),
181            Some(OsStr::new("dumb")),
182        ));
183        assert!(!color_enabled_for(ColorMode::Never, true, None, None));
184    }
185
186    #[test]
187    fn styled_text_tracks_its_destination_stream() {
188        assert_eq!("stdout".red().stream, OutputStream::Stdout);
189        assert_eq!("stderr".red_stderr().stream, OutputStream::Stderr);
190    }
191}
bssh/ui/color.rs

bssh/ui/
color.rs
