tiny_gradient/
display.rs

1use core::fmt::{Display, Formatter, Result};
2
3use crate::{gradient::Gradient as GradientGenerator, RGB};
4
5/// Methods to colorize string with a gradient.
6pub trait GradientStr {
7    /// This function takes a list of colors, which represent a gradient,
8    /// and colorizes the string.
9    fn gradient<I>(&self, colors: I) -> GradientDisplay<'_, I>
10    where
11        I: IntoIterator<Item = RGB> + Clone,
12        I::IntoIter: ExactSizeIterator + Clone;
13}
14
15impl GradientStr for str {
16    fn gradient<I>(&self, colors: I) -> GradientDisplay<'_, I>
17    where
18        I: IntoIterator<Item = RGB>,
19        I::IntoIter: ExactSizeIterator,
20    {
21        GradientDisplay::new(self, colors, ColorType::FOREGROUND)
22    }
23}
24
25impl GradientStr for &str {
26    fn gradient<I>(&self, colors: I) -> GradientDisplay<'_, I>
27    where
28        I: IntoIterator<Item = RGB>,
29        I::IntoIter: ExactSizeIterator,
30    {
31        GradientDisplay::new(self, colors, ColorType::FOREGROUND)
32    }
33}
34
35#[cfg(feature = "std")]
36impl GradientStr for String {
37    fn gradient<I>(&self, colors: I) -> GradientDisplay<'_, I>
38    where
39        I: IntoIterator<Item = RGB>,
40        I::IntoIter: ExactSizeIterator,
41    {
42        GradientDisplay::new(self, colors, ColorType::FOREGROUND)
43    }
44}
45
46/// A gradient string representation.
47#[derive(Debug, Clone, PartialEq, Eq, PartialOrd, Ord)]
48pub struct GradientDisplay<'a, I> {
49    text: &'a str,
50    colors: I,
51    color_type: ColorType,
52}
53
54impl<'a, I> GradientDisplay<'a, I> {
55    const fn new(text: &'a str, colors: I, color_type: ColorType) -> Self {
56        Self {
57            text,
58            colors,
59            color_type,
60        }
61    }
62
63    /// Colorize background.
64    /// 
65    /// Default is foreground.
66    pub const fn background(mut self) -> Self {
67        self.color_type = ColorType::BACKGROUND;
68        self
69    }
70
71    /// Colorize foreground.
72    /// 
73    /// It's a default option.
74    pub const fn foreground(mut self) -> Self {
75        self.color_type = ColorType::FOREGROUND;
76        self
77    }
78}
79
80impl<I> Display for GradientDisplay<'_, I>
81where
82    I: IntoIterator<Item = RGB> + Clone,
83    I::IntoIter: ExactSizeIterator + Clone,
84{
85    fn fmt(&self, f: &mut Formatter<'_>) -> Result {
86        display_gradient(
87            self.text,
88            self.colors.clone().into_iter(),
89            self.color_type,
90            f,
91        )
92    }
93}
94
95fn display_gradient<I>(
96    text: &str,
97    colors: I,
98    color_type: ColorType,
99    f: &mut Formatter<'_>,
100) -> Result
101where
102    I: Iterator<Item = RGB> + Clone + ExactSizeIterator,
103{
104    if colors.len() == 0 || text.is_empty() {
105        return text.fmt(f);
106    }
107
108    let line_width = text.lines().map(|l| l.chars().count()).max().unwrap_or(0);
109    if line_width == 0 {
110        return text.fmt(f);
111    }
112
113    let mut gradient_chunk = line_width;
114    if colors.len() > 2 {
115        gradient_chunk /= colors.len() - 1;
116    }
117
118    let mut gradient = None;
119    let mut next_color = RGB::new(0, 0, 0);
120    let mut _colores = colors.clone();
121    let mut i = 1;
122    for c in text.chars() {
123        if c == '\n' {
124            gradient = None;
125            i = 1;
126            c.fmt(f)?;
127            continue;
128        }
129
130        // happens only on 1st iteration
131        if gradient.is_none() {
132            _colores = colors.clone();
133            let c1: RGB = _colores.next().unwrap();
134            let c2: RGB = _colores.next().unwrap_or(c1);
135            next_color = c2;
136
137            gradient = Some(GradientGenerator::new(c1, c2, gradient_chunk).into_iter());
138        }
139
140        let color = match gradient.as_mut().unwrap().next() {
141            Some(color) => color,
142            None => {
143                i += 1;
144
145                let mut gradient_length = gradient_chunk;
146                let is_last_gradient = i + 1 == colors.len();
147                if is_last_gradient {
148                    gradient_length += line_width - gradient_chunk * i
149                };
150
151                let c1 = next_color;
152                let c2 = _colores.next().unwrap();
153                next_color = c2;
154
155                if gradient_length == 0 {
156                    gradient = Some(GradientGenerator::new(c1, c2, 0).into_iter());
157
158                    c2
159                } else {
160                    // we make +1 to be able to skip
161                    // 1st color because it was the last on prev iteration
162                    gradient =
163                        Some(GradientGenerator::new(c1, c2, gradient_length + 1).into_iter());
164                    let gradient = gradient.as_mut().unwrap();
165                    gradient.next().unwrap();
166                    gradient.next().unwrap()
167                }
168            }
169        };
170
171        colorize_char(c, color, color_type, f)?;
172    }
173
174    Ok(())
175}
176
177fn colorize_char(
178    c: char,
179    RGB { r, g, b }: RGB,
180    color_type: ColorType,
181    f: &mut Formatter<'_>,
182) -> Result {
183    f.write_fmt(format_args!(
184        "\x1b[{};2;{};{};{}m{}\x1b[0m",
185        color_type.0, r, g, b, c
186    ))
187}
188
189#[doc(hidden)]
190#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord)]
191struct ColorType(usize);
192
193impl ColorType {
194    const BACKGROUND: ColorType = ColorType(48);
195    const FOREGROUND: ColorType = ColorType(38);
196}