anstyle_wincon/
windows.rs

1//! Low-level wincon-styling
2
3use std::os::windows::io::AsHandle;
4use std::os::windows::io::AsRawHandle;
5
6type StdioColorResult = std::io::Result<(anstyle::AnsiColor, anstyle::AnsiColor)>;
7type StdioColorInnerResult = Result<(anstyle::AnsiColor, anstyle::AnsiColor), inner::IoError>;
8
9/// Cached [`get_colors`] call for [`std::io::stdout`]
10pub fn stdout_initial_colors() -> StdioColorResult {
11    static INITIAL: once_cell_polyfill::sync::OnceLock<StdioColorInnerResult> =
12        once_cell_polyfill::sync::OnceLock::new();
13    (*INITIAL.get_or_init(|| get_colors_(&std::io::stdout()))).map_err(Into::into)
14}
15
16/// Cached [`get_colors`] call for [`std::io::stderr`]
17pub fn stderr_initial_colors() -> StdioColorResult {
18    static INITIAL: once_cell_polyfill::sync::OnceLock<StdioColorInnerResult> =
19        once_cell_polyfill::sync::OnceLock::new();
20    (*INITIAL.get_or_init(|| get_colors_(&std::io::stderr()))).map_err(Into::into)
21}
22
23/// Apply colors to future writes
24///
25/// **Note:** Make sure any buffers are first flushed or else these colors will apply
26pub fn set_colors<S: AsHandle>(
27    stream: &mut S,
28    fg: anstyle::AnsiColor,
29    bg: anstyle::AnsiColor,
30) -> std::io::Result<()> {
31    set_colors_(stream, fg, bg).map_err(Into::into)
32}
33
34fn set_colors_<S: AsHandle>(
35    stream: &mut S,
36    fg: anstyle::AnsiColor,
37    bg: anstyle::AnsiColor,
38) -> Result<(), inner::IoError> {
39    let handle = stream.as_handle();
40    let handle = handle.as_raw_handle();
41    let attributes = inner::set_colors(fg, bg);
42    inner::set_console_text_attributes(handle, attributes)
43}
44
45/// Get the colors currently active on the console
46pub fn get_colors<S: AsHandle>(stream: &S) -> StdioColorResult {
47    get_colors_(stream).map_err(Into::into)
48}
49
50fn get_colors_<S: AsHandle>(stream: &S) -> StdioColorInnerResult {
51    let handle = stream.as_handle();
52    let handle = handle.as_raw_handle();
53    let info = inner::get_screen_buffer_info(handle)?;
54    let (fg, bg) = inner::get_colors(&info);
55    Ok((fg, bg))
56}
57
58pub(crate) fn write_colored<S: AsHandle + std::io::Write>(
59    stream: &mut S,
60    fg: Option<anstyle::AnsiColor>,
61    bg: Option<anstyle::AnsiColor>,
62    data: &[u8],
63    initial: StdioColorResult,
64) -> std::io::Result<usize> {
65    let (initial_fg, initial_bg) = initial?;
66    let non_default = fg.is_some() || bg.is_some();
67
68    if non_default {
69        let fg = fg.unwrap_or(initial_fg);
70        let bg = bg.unwrap_or(initial_bg);
71        // Ensure everything is written with the last set of colors before applying the next set
72        stream.flush()?;
73        set_colors(stream, fg, bg)?;
74    }
75    let written = stream.write(data)?;
76    if non_default {
77        // Ensure everything is written with the last set of colors before applying the next set
78        stream.flush()?;
79        set_colors(stream, initial_fg, initial_bg)?;
80    }
81    Ok(written)
82}
83
84mod inner {
85    use std::os::windows::io::RawHandle;
86
87    use windows_sys::Win32::Foundation::HANDLE;
88    use windows_sys::Win32::System::Console::CONSOLE_CHARACTER_ATTRIBUTES;
89    use windows_sys::Win32::System::Console::CONSOLE_SCREEN_BUFFER_INFO;
90    use windows_sys::Win32::System::Console::FOREGROUND_BLUE;
91    use windows_sys::Win32::System::Console::FOREGROUND_GREEN;
92    use windows_sys::Win32::System::Console::FOREGROUND_INTENSITY;
93    use windows_sys::Win32::System::Console::FOREGROUND_RED;
94
95    const FOREGROUND_CYAN: CONSOLE_CHARACTER_ATTRIBUTES = FOREGROUND_BLUE | FOREGROUND_GREEN;
96    const FOREGROUND_MAGENTA: CONSOLE_CHARACTER_ATTRIBUTES = FOREGROUND_BLUE | FOREGROUND_RED;
97    const FOREGROUND_YELLOW: CONSOLE_CHARACTER_ATTRIBUTES = FOREGROUND_GREEN | FOREGROUND_RED;
98    const FOREGROUND_WHITE: CONSOLE_CHARACTER_ATTRIBUTES =
99        FOREGROUND_BLUE | FOREGROUND_GREEN | FOREGROUND_RED;
100
101    #[derive(Copy, Clone, Debug)]
102    pub(crate) enum IoError {
103        BrokenPipe,
104        RawOs(i32),
105    }
106
107    impl From<IoError> for std::io::Error {
108        fn from(io: IoError) -> Self {
109            match io {
110                IoError::BrokenPipe => {
111                    std::io::Error::new(std::io::ErrorKind::BrokenPipe, "console is detached")
112                }
113                IoError::RawOs(code) => std::io::Error::from_raw_os_error(code),
114            }
115        }
116    }
117
118    impl IoError {
119        fn last_os_error() -> Self {
120            Self::RawOs(std::io::Error::last_os_error().raw_os_error().unwrap())
121        }
122    }
123
124    pub(crate) fn get_screen_buffer_info(
125        handle: RawHandle,
126    ) -> Result<CONSOLE_SCREEN_BUFFER_INFO, IoError> {
127        unsafe {
128            let handle: HANDLE = handle as HANDLE;
129            if handle.is_null() {
130                return Err(IoError::BrokenPipe);
131            }
132
133            let mut info: CONSOLE_SCREEN_BUFFER_INFO = std::mem::zeroed();
134            if windows_sys::Win32::System::Console::GetConsoleScreenBufferInfo(handle, &mut info)
135                != 0
136            {
137                Ok(info)
138            } else {
139                Err(IoError::last_os_error())
140            }
141        }
142    }
143
144    pub(crate) fn set_console_text_attributes(
145        handle: RawHandle,
146        attributes: CONSOLE_CHARACTER_ATTRIBUTES,
147    ) -> Result<(), IoError> {
148        unsafe {
149            let handle: HANDLE = handle as HANDLE;
150            if handle.is_null() {
151                return Err(IoError::BrokenPipe);
152            }
153
154            if windows_sys::Win32::System::Console::SetConsoleTextAttribute(handle, attributes) != 0
155            {
156                Ok(())
157            } else {
158                Err(IoError::last_os_error())
159            }
160        }
161    }
162
163    pub(crate) fn get_colors(
164        info: &CONSOLE_SCREEN_BUFFER_INFO,
165    ) -> (anstyle::AnsiColor, anstyle::AnsiColor) {
166        let attributes = info.wAttributes;
167        let bg = from_nibble(attributes >> 4);
168        let fg = from_nibble(attributes);
169        (fg, bg)
170    }
171
172    pub(crate) fn set_colors(
173        fg: anstyle::AnsiColor,
174        bg: anstyle::AnsiColor,
175    ) -> CONSOLE_CHARACTER_ATTRIBUTES {
176        to_nibble(bg) << 4 | to_nibble(fg)
177    }
178
179    fn from_nibble(color: CONSOLE_CHARACTER_ATTRIBUTES) -> anstyle::AnsiColor {
180        if color & FOREGROUND_WHITE == FOREGROUND_WHITE {
181            // 3 bits high
182            anstyle::AnsiColor::White
183        } else if color & FOREGROUND_CYAN == FOREGROUND_CYAN {
184            // 2 bits high
185            anstyle::AnsiColor::Cyan
186        } else if color & FOREGROUND_YELLOW == FOREGROUND_YELLOW {
187            // 2 bits high
188            anstyle::AnsiColor::Yellow
189        } else if color & FOREGROUND_MAGENTA == FOREGROUND_MAGENTA {
190            // 2 bits high
191            anstyle::AnsiColor::Magenta
192        } else if color & FOREGROUND_RED == FOREGROUND_RED {
193            // 1 bit high
194            anstyle::AnsiColor::Red
195        } else if color & FOREGROUND_GREEN == FOREGROUND_GREEN {
196            // 1 bit high
197            anstyle::AnsiColor::Green
198        } else if color & FOREGROUND_BLUE == FOREGROUND_BLUE {
199            // 1 bit high
200            anstyle::AnsiColor::Blue
201        } else {
202            // 0 bits high
203            anstyle::AnsiColor::Black
204        }
205        .bright(color & FOREGROUND_INTENSITY == FOREGROUND_INTENSITY)
206    }
207
208    fn to_nibble(color: anstyle::AnsiColor) -> CONSOLE_CHARACTER_ATTRIBUTES {
209        let mut attributes = 0;
210        attributes |= match color.bright(false) {
211            anstyle::AnsiColor::Black => 0,
212            anstyle::AnsiColor::Red => FOREGROUND_RED,
213            anstyle::AnsiColor::Green => FOREGROUND_GREEN,
214            anstyle::AnsiColor::Yellow => FOREGROUND_YELLOW,
215            anstyle::AnsiColor::Blue => FOREGROUND_BLUE,
216            anstyle::AnsiColor::Magenta => FOREGROUND_MAGENTA,
217            anstyle::AnsiColor::Cyan => FOREGROUND_CYAN,
218            anstyle::AnsiColor::White => FOREGROUND_WHITE,
219            anstyle::AnsiColor::BrightBlack
220            | anstyle::AnsiColor::BrightRed
221            | anstyle::AnsiColor::BrightGreen
222            | anstyle::AnsiColor::BrightYellow
223            | anstyle::AnsiColor::BrightBlue
224            | anstyle::AnsiColor::BrightMagenta
225            | anstyle::AnsiColor::BrightCyan
226            | anstyle::AnsiColor::BrightWhite => unreachable!("brights were toggled off"),
227        };
228        if color.is_bright() {
229            attributes |= FOREGROUND_INTENSITY;
230        }
231        attributes
232    }
233
234    #[test]
235    fn to_from_nibble() {
236        const COLORS: [anstyle::AnsiColor; 16] = [
237            anstyle::AnsiColor::Black,
238            anstyle::AnsiColor::Red,
239            anstyle::AnsiColor::Green,
240            anstyle::AnsiColor::Yellow,
241            anstyle::AnsiColor::Blue,
242            anstyle::AnsiColor::Magenta,
243            anstyle::AnsiColor::Cyan,
244            anstyle::AnsiColor::White,
245            anstyle::AnsiColor::BrightBlack,
246            anstyle::AnsiColor::BrightRed,
247            anstyle::AnsiColor::BrightGreen,
248            anstyle::AnsiColor::BrightYellow,
249            anstyle::AnsiColor::BrightBlue,
250            anstyle::AnsiColor::BrightMagenta,
251            anstyle::AnsiColor::BrightCyan,
252            anstyle::AnsiColor::BrightWhite,
253        ];
254        for expected in COLORS {
255            let nibble = to_nibble(expected);
256            let actual = from_nibble(nibble);
257            assert_eq!(expected, actual, "Intermediate: {nibble}");
258        }
259    }
260}