Skip to main content

jt_consoleutils/
colorize.rs

1//! Helpers to colorize text with a left-to-right rainbow.
2//!
3//! Provides:
4//! - `colorize_text_with_width(text, Option<usize>) -> String`
5//!
6//! The functions emit 24-bit ANSI foreground escapes (`\x1b[38;2;R;G;Bm`).
7//! If `rainbow_width` is `None` the rainbow spans the widest line in `text`.
8//! The output contains newline characters and is ready to print.
9
10use crate::colors::RESET;
11
12/// Convert HSV (h in degrees 0..360, s and v in 0..1) to RGB bytes (0..255).
13fn hsv_to_rgb(h_deg: f32, s: f32, v: f32) -> (u8, u8, u8) {
14   let h = (h_deg % 360.0 + 360.0) % 360.0;
15   let c = v * s;
16   let h_prime = h / 60.0;
17   let x = c * (1.0 - ((h_prime % 2.0) - 1.0).abs());
18   let (r1, g1, b1) = if (0.0..1.0).contains(&h_prime) {
19      (c, x, 0.0)
20   } else if h_prime < 2.0 {
21      (x, c, 0.0)
22   } else if h_prime < 3.0 {
23      (0.0, c, x)
24   } else if h_prime < 4.0 {
25      (0.0, x, c)
26   } else if h_prime < 5.0 {
27      (x, 0.0, c)
28   } else {
29      (c, 0.0, x)
30   };
31   let m = v - c;
32   let r = ((r1 + m) * 255.0).round().clamp(0.0, 255.0) as u8;
33   let g = ((g1 + m) * 255.0).round().clamp(0.0, 255.0) as u8;
34   let b = ((b1 + m) * 255.0).round().clamp(0.0, 255.0) as u8;
35   (r, g, b)
36}
37
38/// Create an ANSI 24-bit foreground escape for the given RGB bytes.
39fn ansi_rgb_escape(r: u8, g: u8, b: u8) -> String {
40   format!("\x1b[38;2;{};{};{}m", r, g, b)
41}
42
43/// Colorize the provided `text` with a left-to-right rainbow. If `rainbow_width`
44/// is `None` the rainbow will span the full width of the widest line. If a
45/// width is provided the gradient is computed across that width and repeats
46/// when lines are longer.
47///
48/// The returned String contains newlines and ANSI escapes; print it directly.
49pub fn colorize_text_with_width(text: &str, rainbow_width: Option<usize>) -> String {
50   // Split into lines and characters so we can index columns.
51   let lines: Vec<String> = text.lines().map(str::to_owned).collect();
52   let char_lines: Vec<Vec<char>> = lines.iter().map(|l| l.chars().collect()).collect();
53
54   let max_width = char_lines.iter().map(|l| l.len()).max().unwrap_or(0);
55   if max_width == 0 {
56      // Nothing to color; return the original text unchanged.
57      return text.to_string();
58   }
59
60   // Determine rainbow width: span the text if None.
61   let rainbow_w = rainbow_width.unwrap_or(max_width).max(1);
62
63   // Build a palette of per-column colors across rainbow_w.
64   let mut col_colors: Vec<String> = Vec::with_capacity(rainbow_w);
65   for col in 0..rainbow_w {
66      let t = (col as f32) / (rainbow_w.max(1) as f32); // 0..1
67      // Map to hue (degrees). Phase can be tweaked if desired.
68      let hue_deg = (t * 360.0) % 360.0;
69      // Use fairly high saturation/value for vivid output; these can be tuned.
70      let (r, g, b) = hsv_to_rgb(hue_deg, 0.5_f32.min(1.0), 0.99_f32.min(1.0));
71      col_colors.push(ansi_rgb_escape(r, g, b));
72   }
73
74   let mut result = String::new();
75   for line in &char_lines {
76      let mut col = 0usize;
77      while col < line.len() {
78         result.push_str(&col_colors[col % rainbow_w]);
79         result.push(line[col]);
80         col += 1;
81      }
82      result.push('\n');
83   }
84
85   result.push_str(RESET);
86   result
87}
88
89#[cfg(test)]
90mod tests {
91   use super::*;
92
93   #[test]
94   fn empty_text_returns_unchanged() {
95      assert_eq!(colorize_text_with_width("", Some(80)), "");
96   }
97
98   #[test]
99   fn single_line_contains_ansi_escape() {
100      let output = colorize_text_with_width("hello", Some(80));
101      assert!(output.contains("\x1b[38;2;"));
102   }
103
104   #[test]
105   fn output_ends_with_reset() {
106      let output = colorize_text_with_width("hello", Some(80));
107      assert!(output.ends_with("\x1b[0m"));
108   }
109
110   #[test]
111   fn explicit_width_repeats_gradient() {
112      // A line much longer than rainbow_width should not panic.
113      let long_line = "a".repeat(200);
114      let output = colorize_text_with_width(&long_line, Some(10));
115      assert!(output.contains("\x1b[38;2;"));
116      assert!(output.ends_with("\x1b[0m"));
117   }
118
119   #[test]
120   fn none_width_uses_max_line_width() {
121      let text = "hello world";
122      let explicit = colorize_text_with_width(text, Some(text.len()));
123      let implicit = colorize_text_with_width(text, None);
124      assert_eq!(explicit, implicit);
125   }
126}