Skip to content

Commit 12ae599

Browse files
float3float3
committed
check terminal for ansi code support
1 parent e3a7a14 commit 12ae599

File tree

5 files changed

+169
-14
lines changed

5 files changed

+169
-14
lines changed

‎Cargo.lock‎

Lines changed: 50 additions & 0 deletions
Some generated files are not rendered by default. Learn more about customizing how changed files appear on GitHub.

‎Cargo.toml‎

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -16,3 +16,5 @@ loaded = ["ash/loaded"]
1616

1717
[dependencies]
1818
ash = { version = "0.38.0", default-features = false, features = ["std"] }
19+
atty = "0.2.14"
20+
winapi = "0.3.9"

‎src/ascii_art.rs‎

Lines changed: 17 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -4,19 +4,19 @@ pub(crate) const CHARS: [char; LUT_SIZE] = ['#', '$', '%', '&', '@'];
44

55
pub(crate) const BLACK: &str = "\x1B[30m";
66
pub(crate) const RED: &str = "\x1B[31m";
7-
// pub(crate) const BRIGHT_RED: &str = "\x1B[91m";
7+
pub(crate) const BRIGHT_RED: &str = "\x1B[91m";
88
pub(crate) const GREEN: &str = "\x1B[32m";
9-
// pub(crate) const BRIGHT_GREEN: &str = "\x1B[92m";
9+
pub(crate) const BRIGHT_GREEN: &str = "\x1B[92m";
1010
pub(crate) const YELLOW: &str = "\x1B[33m";
11-
// pub(crate) const BRIGHT_YELLOW: &str = "\x1B[93m";
11+
pub(crate) const BRIGHT_YELLOW: &str = "\x1B[93m";
1212
pub(crate) const BLUE: &str = "\x1B[34m";
13-
// pub(crate) const BRIGHT_BLUE: &str = "\x1B[94m";
13+
pub(crate) const BRIGHT_BLUE: &str = "\x1B[94m";
1414
// pub(crate) const MAGENTA: &str = "\x1B[35m";
1515
// pub(crate) const BRIGHT_MAGENTA: &str = "\x1B[95m";
1616
pub(crate) const CYAN: &str = "\x1B[36m";
17-
// pub(crate) const BRIGHT_CYAN: &str = "\x1B[96m";
17+
pub(crate) const BRIGHT_CYAN: &str = "\x1B[96m";
1818
pub(crate) const WHITE: &str = "\x1B[37m";
19-
// pub(crate) const BRIGHT_WHITE: &str = "\x1B[97m";
19+
pub(crate) const BRIGHT_WHITE: &str = "\x1B[97m";
2020

2121
pub(crate) const VULKAN: &[&str] = &[
2222
r#" ################ "#,
@@ -35,6 +35,7 @@ pub(crate) const VULKAN: &[&str] = &[
3535
];
3636

3737
pub(crate) const VULKAN_STYLE: [&str; LUT_SIZE] = [RED, "", "", "", BLACK];
38+
pub(crate) const VULKAN_STYLE_ALT: [&str; LUT_SIZE] = [BRIGHT_RED, "", "", "", BLACK];
3839

3940
pub(crate) const APPLE: &[&str] = &[
4041
r#" ### "#,
@@ -55,6 +56,7 @@ pub(crate) const APPLE: &[&str] = &[
5556
];
5657

5758
pub(crate) const APPLE_STYLE: [&str; LUT_SIZE] = [WHITE, "", "", "", BLACK];
59+
pub(crate) const APPLE_STYLE_ALT: [&str; LUT_SIZE] = [BRIGHT_WHITE, "", "", "", BLACK];
5860

5961
pub(crate) const GOOGLE: &[&str] = &[
6062
r#" ######### "#,
@@ -77,6 +79,8 @@ pub(crate) const GOOGLE: &[&str] = &[
7779
];
7880

7981
pub(crate) const GOOGLE_STYLE: [&str; LUT_SIZE] = [RED, BLUE, GREEN, YELLOW, BLACK];
82+
pub(crate) const GOOGLE_STYLE_ALT: [&str; LUT_SIZE] =
83+
[BRIGHT_RED, BRIGHT_BLUE, BRIGHT_GREEN, BRIGHT_YELLOW, BLACK];
8084

8185
pub(crate) const INTEL: &[&str] = &[
8286
r#" $$$ ### "#,
@@ -91,6 +95,7 @@ pub(crate) const INTEL: &[&str] = &[
9195
];
9296

9397
pub(crate) const INTEL_STYLE: [&str; LUT_SIZE] = [WHITE, CYAN, "", "", BLACK];
98+
pub(crate) const INTEL_STYLE_ALT: [&str; LUT_SIZE] = [BRIGHT_WHITE, BRIGHT_CYAN, "", "", BLACK];
9499

95100
pub(crate) const NVIDIA: &[&str] = &[
96101
r#" #########################"#,
@@ -108,6 +113,7 @@ pub(crate) const NVIDIA: &[&str] = &[
108113
];
109114

110115
pub(crate) const NVIDIA_STYLE: [&str; LUT_SIZE] = [GREEN, "", "", "", BLACK];
116+
pub(crate) const NVIDIA_STYLE_ALT: [&str; LUT_SIZE] = [BRIGHT_GREEN, "", "", "", BLACK];
111117

112118
pub(crate) const AMD: &[&str] = &[
113119
r#" ### ### ### ######### $$$$$$$$$"#,
@@ -119,6 +125,7 @@ pub(crate) const AMD: &[&str] = &[
119125
];
120126

121127
pub(crate) const AMD_STYLE: [&str; LUT_SIZE] = [WHITE, GREEN, "", "", BLACK];
128+
pub(crate) const AMD_STYLE_ALT: [&str; LUT_SIZE] = [BRIGHT_WHITE, BRIGHT_GREEN, "", "", BLACK];
122129

123130
pub(crate) const ARM: &[&str] = &[
124131
r#" ### ######## ### ### "#,
@@ -130,6 +137,7 @@ pub(crate) const ARM: &[&str] = &[
130137
];
131138

132139
pub(crate) const ARM_STYLE: [&str; LUT_SIZE] = [RED, "", "", "", BLACK];
140+
pub(crate) const ARM_STYLE_ALT: [&str; LUT_SIZE] = [BRIGHT_RED, "", "", "", BLACK];
133141

134142
// pub(crate) const MESA: &[&str] = &[
135143
// r#" ### ### $$$$$$$$ %%%%%%%% &&& "#,
@@ -160,6 +168,8 @@ pub(crate) const MICROSOFT: &[&str] = &[
160168
];
161169

162170
pub(crate) const MICROSOFT_STYLE: [&str; LUT_SIZE] = [RED, GREEN, BLUE, YELLOW, ""];
171+
pub(crate) const MICROSOFT_STYLE_ALT: [&str; LUT_SIZE] =
172+
[BRIGHT_RED, BRIGHT_GREEN, BRIGHT_BLUE, BRIGHT_YELLOW, ""];
163173

164174
pub(crate) const QUALCOMM: &[&str] = &[
165175
r#" ######## "#,
@@ -180,3 +190,4 @@ pub(crate) const QUALCOMM: &[&str] = &[
180190
];
181191

182192
pub(crate) const QUALCOMM_STYLE: [&str; LUT_SIZE] = [BLUE, "", "", "", ""];
193+
pub(crate) const QUALCOMM_STYLE_ALT: [&str; LUT_SIZE] = [BRIGHT_BLUE, "", "", "", ""];

‎src/lib.rs‎

Lines changed: 68 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -6,6 +6,7 @@ use ash::{self, vk, Entry, Instance};
66
use device::Device;
77
use std::error::Error;
88
use vendor::Vendor;
9+
use vt::enable_virtual_terminal_processing;
910

1011
const BOLD: &str = "\x1B[1m";
1112
const RESET: &str = "\x1B[0m";
@@ -28,7 +29,9 @@ pub fn fetch_device(
2829
let art = vendor.get_ascii_art();
2930

3031
let device = Device::new(instance, device_handle);
31-
let info = get_device_info(&device, vendor.get_styles()[0]);
32+
let info = get_device_info(&device, vendor.get_style()[0]);
33+
34+
let _ = enable_virtual_terminal_processing();
3235

3336
for i in 0..art.len().max(info.len()) {
3437
let art_line = art
@@ -278,6 +281,70 @@ pub fn iterate_devices() -> Result<(), Box<dyn Error>> {
278281
Ok(())
279282
}
280283

284+
#[cfg(windows)]
285+
mod vt {
286+
use std::io::{Error, Result};
287+
use winapi::um::consoleapi::{GetConsoleMode, SetConsoleMode};
288+
use winapi::um::handleapi::INVALID_HANDLE_VALUE;
289+
use winapi::um::processenv::GetStdHandle;
290+
use winapi::um::winbase::STD_OUTPUT_HANDLE;
291+
use winapi::um::wincon::ENABLE_VIRTUAL_TERMINAL_PROCESSING;
292+
293+
/// Enables Virtual Terminal Processing on Windows.
294+
pub fn enable_virtual_terminal_processing() -> Result<()> {
295+
unsafe {
296+
let std_out = GetStdHandle(STD_OUTPUT_HANDLE);
297+
if std_out == INVALID_HANDLE_VALUE {
298+
return Err(Error::last_os_error());
299+
}
300+
let mut mode = 0;
301+
if GetConsoleMode(std_out, &mut mode) == 0 {
302+
return Err(Error::last_os_error());
303+
}
304+
mode |= ENABLE_VIRTUAL_TERMINAL_PROCESSING;
305+
if SetConsoleMode(std_out, mode) == 0 {
306+
return Err(Error::last_os_error());
307+
}
308+
}
309+
Ok(())
310+
}
311+
312+
/// Checks if Virtual Terminal Processing is enabled.
313+
pub fn is_vt_enabled() -> bool {
314+
unsafe {
315+
let std_out = GetStdHandle(STD_OUTPUT_HANDLE);
316+
if std_out == INVALID_HANDLE_VALUE {
317+
return false;
318+
}
319+
let mut mode = 0;
320+
if GetConsoleMode(std_out, &mut mode) == 0 {
321+
return false;
322+
}
323+
(mode & ENABLE_VIRTUAL_TERMINAL_PROCESSING) != 0
324+
}
325+
}
326+
}
327+
328+
#[cfg(not(windows))]
329+
mod vt {
330+
use std::io::Result;
331+
332+
/// On non‑Windows platforms, VT processing is typically enabled by default.
333+
pub fn enable_virtual_terminal_processing() -> Result<()> {
334+
Ok(())
335+
}
336+
337+
/// Assume ANSI escape codes are supported.
338+
pub fn is_vt_enabled() -> bool {
339+
true
340+
}
341+
}
342+
343+
/// Returns `true` if stdout is a TTY and (on Windows) VT processing is enabled.
344+
fn is_ansi_supported() -> bool {
345+
atty::is(atty::Stream::Stdout) && vt::is_vt_enabled()
346+
}
347+
281348
#[cfg(test)]
282349
mod tests {
283350
use super::*;

‎src/vendor.rs‎

Lines changed: 32 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -1,4 +1,4 @@
1-
use crate::ascii_art::*;
1+
use crate::{ascii_art::*, is_ansi_supported};
22

33
/// Represents a GPU vendor.
44
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
@@ -47,11 +47,14 @@ impl Vendor {
4747
| Vendor::MobileEye => VULKAN,
4848
};
4949

50-
let styles = self.get_styles();
51-
50+
let style = if is_ansi_supported() {
51+
self.get_alternative_style()
52+
} else {
53+
self.get_style()
54+
};
5255
let mut modified_art: Vec<String> = art.iter().map(|&line| line.to_string()).collect();
5356

54-
for (symbol, style) in CHARS.iter().zip(styles.iter()) {
57+
for (symbol, style) in CHARS.iter().zip(style.iter()) {
5558
if !style.is_empty() {
5659
modified_art = modified_art
5760
.iter()
@@ -64,7 +67,7 @@ impl Vendor {
6467
}
6568

6669
/// Returns an array of style strings associated with the vendor.
67-
pub const fn get_styles(&self) -> [&str; LUT_SIZE] {
70+
pub const fn get_style(&self) -> [&str; LUT_SIZE] {
6871
match self {
6972
Vendor::AMD => AMD_STYLE,
7073
Vendor::Apple => APPLE_STYLE,
@@ -132,6 +135,28 @@ impl Vendor {
132135
_ => None,
133136
}
134137
}
138+
139+
fn get_alternative_style(&self) -> [&str; LUT_SIZE] {
140+
match self {
141+
Vendor::AMD => AMD_STYLE_ALT,
142+
Vendor::Apple => APPLE_STYLE_ALT,
143+
Vendor::ARM => ARM_STYLE_ALT,
144+
Vendor::Google => GOOGLE_STYLE_ALT,
145+
Vendor::Intel => INTEL_STYLE_ALT,
146+
Vendor::Nvidia => NVIDIA_STYLE_ALT,
147+
Vendor::Microsoft => MICROSOFT_STYLE_ALT,
148+
Vendor::Qualcomm => QUALCOMM_STYLE_ALT,
149+
Vendor::Mesa => VULKAN_STYLE_ALT,
150+
Vendor::Unknown
151+
| Vendor::ImgTec
152+
| Vendor::VIV
153+
| Vendor::VSI
154+
| Vendor::Kazan
155+
| Vendor::Codeplay
156+
| Vendor::Pocl
157+
| Vendor::MobileEye => VULKAN_STYLE_ALT,
158+
}
159+
}
135160
}
136161

137162
impl From<Vendor> for u32 {
@@ -177,7 +202,7 @@ mod tests {
177202
#[test]
178203
fn test_get_styles_length() {
179204
let vendor = Vendor::AMD;
180-
let styles = vendor.get_styles();
205+
let styles = vendor.get_style();
181206
assert_eq!(styles.len(), crate::ascii_art::LUT_SIZE);
182207
}
183208

@@ -187,7 +212,7 @@ mod tests {
187212
let art = vendor.get_ascii_art();
188213
assert!(!art.is_empty(), "ASCII art should not be empty");
189214

190-
let styles = vendor.get_styles();
215+
let styles = vendor.get_style();
191216
let non_empty_style = styles.iter().any(|s| !s.is_empty());
192217
if non_empty_style {
193218
let block_found = art.iter().any(|line| line.contains(BLOCK));

0 commit comments

Comments
 (0)