tree_sitter_cli/
highlight.rs

1use std::{
2    collections::{BTreeMap, HashSet},
3    fmt::Write,
4    fs,
5    io::{self, Write as _},
6    path::{self, Path, PathBuf},
7    str,
8    sync::{atomic::AtomicUsize, Arc},
9    time::Instant,
10};
11
12use ansi_colours::{ansi256_from_rgb, rgb_from_ansi256};
13use anstyle::{Ansi256Color, AnsiColor, Color, Effects, RgbColor};
14use anyhow::Result;
15use log::{info, warn};
16use serde::{ser::SerializeMap, Deserialize, Deserializer, Serialize, Serializer};
17use serde_json::{json, Value};
18use tree_sitter_highlight::{HighlightConfiguration, HighlightEvent, Highlighter, HtmlRenderer};
19use tree_sitter_loader::Loader;
20
21pub const HTML_HEAD_HEADER: &str = "
22<!doctype HTML>
23<head>
24  <title>Tree-sitter Highlighting</title>
25  <style>
26    body {
27      font-family: monospace
28    }
29    .line-number {
30      user-select: none;
31      text-align: right;
32      color: rgba(27,31,35,.3);
33      padding: 0 10px;
34    }
35    .line {
36      white-space: pre;
37    }
38  </style>";
39
40pub const HTML_BODY_HEADER: &str = "
41</head>
42<body>
43";
44
45pub const HTML_FOOTER: &str = "
46</body>
47";
48
49#[derive(Debug, Default)]
50pub struct Style {
51    pub ansi: anstyle::Style,
52    pub css: Option<String>,
53}
54
55#[derive(Debug)]
56pub struct Theme {
57    pub styles: Vec<Style>,
58    pub highlight_names: Vec<String>,
59}
60
61#[derive(Default, Deserialize, Serialize)]
62pub struct ThemeConfig {
63    #[serde(default)]
64    pub theme: Theme,
65}
66
67impl Theme {
68    pub fn load(path: &path::Path) -> io::Result<Self> {
69        let json = fs::read_to_string(path)?;
70        Ok(serde_json::from_str(&json).unwrap_or_default())
71    }
72
73    #[must_use]
74    pub fn default_style(&self) -> Style {
75        Style::default()
76    }
77}
78
79impl<'de> Deserialize<'de> for Theme {
80    fn deserialize<D>(deserializer: D) -> std::result::Result<Self, D::Error>
81    where
82        D: Deserializer<'de>,
83    {
84        let mut styles = Vec::new();
85        let mut highlight_names = Vec::new();
86        if let Ok(colors) = BTreeMap::<String, Value>::deserialize(deserializer) {
87            styles.reserve(colors.len());
88            highlight_names.reserve(colors.len());
89            for (name, style_value) in colors {
90                let mut style = Style::default();
91                parse_style(&mut style, style_value);
92                highlight_names.push(name);
93                styles.push(style);
94            }
95        }
96        Ok(Self {
97            styles,
98            highlight_names,
99        })
100    }
101}
102
103impl Serialize for Theme {
104    fn serialize<S>(&self, serializer: S) -> std::result::Result<S::Ok, S::Error>
105    where
106        S: Serializer,
107    {
108        let mut map = serializer.serialize_map(Some(self.styles.len()))?;
109        for (name, style) in self.highlight_names.iter().zip(&self.styles) {
110            let style = &style.ansi;
111            let color = style.get_fg_color().map(|color| match color {
112                Color::Ansi(color) => match color {
113                    AnsiColor::Black => json!("black"),
114                    AnsiColor::Blue => json!("blue"),
115                    AnsiColor::Cyan => json!("cyan"),
116                    AnsiColor::Green => json!("green"),
117                    AnsiColor::Magenta => json!("purple"),
118                    AnsiColor::Red => json!("red"),
119                    AnsiColor::White => json!("white"),
120                    AnsiColor::Yellow => json!("yellow"),
121                    _ => unreachable!(),
122                },
123                Color::Ansi256(Ansi256Color(n)) => json!(n),
124                Color::Rgb(RgbColor(r, g, b)) => json!(format!("#{r:x?}{g:x?}{b:x?}")),
125            });
126            let effects = style.get_effects();
127            if effects.contains(Effects::BOLD)
128                || effects.contains(Effects::ITALIC)
129                || effects.contains(Effects::UNDERLINE)
130            {
131                let mut style_json = BTreeMap::new();
132                if let Some(color) = color {
133                    style_json.insert("color", color);
134                }
135                if effects.contains(Effects::BOLD) {
136                    style_json.insert("bold", Value::Bool(true));
137                }
138                if effects.contains(Effects::ITALIC) {
139                    style_json.insert("italic", Value::Bool(true));
140                }
141                if effects.contains(Effects::UNDERLINE) {
142                    style_json.insert("underline", Value::Bool(true));
143                }
144                map.serialize_entry(&name, &style_json)?;
145            } else if let Some(color) = color {
146                map.serialize_entry(&name, &color)?;
147            } else {
148                map.serialize_entry(&name, &Value::Null)?;
149            }
150        }
151        map.end()
152    }
153}
154
155impl Default for Theme {
156    fn default() -> Self {
157        serde_json::from_value(json!({
158            "attribute": {"color": 124, "italic": true},
159            "comment": {"color": 245, "italic": true},
160            "constant": 94,
161            "constant.builtin": {"color": 94, "bold": true},
162            "constructor": 136,
163            "embedded": null,
164            "function": 26,
165            "function.builtin": {"color": 26, "bold": true},
166            "keyword": 56,
167            "module": 136,
168            "number": {"color": 94, "bold": true},
169            "operator": {"color": 239, "bold": true},
170            "property": 124,
171            "property.builtin": {"color": 124, "bold": true},
172            "punctuation": 239,
173            "punctuation.bracket": 239,
174            "punctuation.delimiter": 239,
175            "punctuation.special": 239,
176            "string": 28,
177            "string.special": 30,
178            "tag": 18,
179            "type": 23,
180            "type.builtin": {"color": 23, "bold": true},
181            "variable": 252,
182            "variable.builtin": {"color": 252, "bold": true},
183            "variable.parameter": {"color": 252, "underline": true}
184        }))
185        .unwrap()
186    }
187}
188
189fn parse_style(style: &mut Style, json: Value) {
190    if let Value::Object(entries) = json {
191        for (property_name, value) in entries {
192            match property_name.as_str() {
193                "bold" => {
194                    if value == Value::Bool(true) {
195                        style.ansi = style.ansi.bold();
196                    }
197                }
198                "italic" => {
199                    if value == Value::Bool(true) {
200                        style.ansi = style.ansi.italic();
201                    }
202                }
203                "underline" => {
204                    if value == Value::Bool(true) {
205                        style.ansi = style.ansi.underline();
206                    }
207                }
208                "color" => {
209                    if let Some(color) = parse_color(value) {
210                        style.ansi = style.ansi.fg_color(Some(color));
211                    }
212                }
213                _ => {}
214            }
215        }
216        style.css = Some(style_to_css(style.ansi));
217    } else if let Some(color) = parse_color(json) {
218        style.ansi = style.ansi.fg_color(Some(color));
219        style.css = Some(style_to_css(style.ansi));
220    } else {
221        style.css = None;
222    }
223
224    if let Some(Color::Rgb(RgbColor(red, green, blue))) = style.ansi.get_fg_color() {
225        if !terminal_supports_truecolor() {
226            let ansi256 = Color::Ansi256(Ansi256Color(ansi256_from_rgb((red, green, blue))));
227            style.ansi = style.ansi.fg_color(Some(ansi256));
228        }
229    }
230}
231
232fn parse_color(json: Value) -> Option<Color> {
233    match json {
234        Value::Number(n) => n.as_u64().map(|n| Color::Ansi256(Ansi256Color(n as u8))),
235        Value::String(s) => match s.to_lowercase().as_str() {
236            "black" => Some(Color::Ansi(AnsiColor::Black)),
237            "blue" => Some(Color::Ansi(AnsiColor::Blue)),
238            "cyan" => Some(Color::Ansi(AnsiColor::Cyan)),
239            "green" => Some(Color::Ansi(AnsiColor::Green)),
240            "purple" => Some(Color::Ansi(AnsiColor::Magenta)),
241            "red" => Some(Color::Ansi(AnsiColor::Red)),
242            "white" => Some(Color::Ansi(AnsiColor::White)),
243            "yellow" => Some(Color::Ansi(AnsiColor::Yellow)),
244            s => {
245                if let Some((red, green, blue)) = hex_string_to_rgb(s) {
246                    Some(Color::Rgb(RgbColor(red, green, blue)))
247                } else {
248                    None
249                }
250            }
251        },
252        _ => None,
253    }
254}
255
256fn hex_string_to_rgb(s: &str) -> Option<(u8, u8, u8)> {
257    if s.starts_with('#') && s.len() >= 7 {
258        if let (Ok(red), Ok(green), Ok(blue)) = (
259            u8::from_str_radix(&s[1..3], 16),
260            u8::from_str_radix(&s[3..5], 16),
261            u8::from_str_radix(&s[5..7], 16),
262        ) {
263            Some((red, green, blue))
264        } else {
265            None
266        }
267    } else {
268        None
269    }
270}
271
272fn style_to_css(style: anstyle::Style) -> String {
273    let mut result = String::new();
274    let effects = style.get_effects();
275    if effects.contains(Effects::UNDERLINE) {
276        write!(&mut result, "text-decoration: underline;").unwrap();
277    }
278    if effects.contains(Effects::BOLD) {
279        write!(&mut result, "font-weight: bold;").unwrap();
280    }
281    if effects.contains(Effects::ITALIC) {
282        write!(&mut result, "font-style: italic;").unwrap();
283    }
284    if let Some(color) = style.get_fg_color() {
285        write_color(&mut result, color);
286    }
287    result
288}
289
290fn write_color(buffer: &mut String, color: Color) {
291    match color {
292        Color::Ansi(color) => match color {
293            AnsiColor::Black => write!(buffer, "color: black").unwrap(),
294            AnsiColor::Red => write!(buffer, "color: red").unwrap(),
295            AnsiColor::Green => write!(buffer, "color: green").unwrap(),
296            AnsiColor::Yellow => write!(buffer, "color: yellow").unwrap(),
297            AnsiColor::Blue => write!(buffer, "color: blue").unwrap(),
298            AnsiColor::Magenta => write!(buffer, "color: purple").unwrap(),
299            AnsiColor::Cyan => write!(buffer, "color: cyan").unwrap(),
300            AnsiColor::White => write!(buffer, "color: white").unwrap(),
301            _ => unreachable!(),
302        },
303        Color::Ansi256(Ansi256Color(n)) => {
304            let (r, g, b) = rgb_from_ansi256(n);
305            write!(buffer, "color: #{r:02x}{g:02x}{b:02x}").unwrap();
306        }
307        Color::Rgb(RgbColor(r, g, b)) => write!(buffer, "color: #{r:02x}{g:02x}{b:02x}").unwrap(),
308    }
309}
310
311fn terminal_supports_truecolor() -> bool {
312    std::env::var("COLORTERM")
313        .is_ok_and(|truecolor| truecolor == "truecolor" || truecolor == "24bit")
314}
315
316pub struct HighlightOptions {
317    pub theme: Theme,
318    pub check: bool,
319    pub captures_path: Option<PathBuf>,
320    pub inline_styles: bool,
321    pub html: bool,
322    pub quiet: bool,
323    pub print_time: bool,
324    pub cancellation_flag: Arc<AtomicUsize>,
325}
326
327pub fn highlight(
328    loader: &Loader,
329    path: &Path,
330    name: &str,
331    config: &HighlightConfiguration,
332    print_name: bool,
333    opts: &HighlightOptions,
334) -> Result<()> {
335    if opts.check {
336        let names = if let Some(path) = opts.captures_path.as_deref() {
337            let file = fs::read_to_string(path)?;
338            let capture_names = file
339                .lines()
340                .filter_map(|line| {
341                    if line.trim().is_empty() || line.trim().starts_with(';') {
342                        return None;
343                    }
344                    line.split(';').next().map(|s| s.trim().trim_matches('"'))
345                })
346                .collect::<HashSet<_>>();
347            config.nonconformant_capture_names(&capture_names)
348        } else {
349            config.nonconformant_capture_names(&HashSet::new())
350        };
351        if names.is_empty() {
352            info!("All highlight captures conform to standards.");
353        } else {
354            warn!(
355                "Non-standard highlight {} detected:\n* {}",
356                if names.len() > 1 {
357                    "captures"
358                } else {
359                    "capture"
360                },
361                names.join("\n* ")
362            );
363        }
364    }
365
366    let source = fs::read(path)?;
367    let stdout = io::stdout();
368    let mut stdout = stdout.lock();
369    let time = Instant::now();
370    let mut highlighter = Highlighter::new();
371    let events =
372        highlighter.highlight(config, &source, Some(&opts.cancellation_flag), |string| {
373            loader.highlight_config_for_injection_string(string)
374        })?;
375    let theme = &opts.theme;
376
377    if !opts.quiet && print_name {
378        writeln!(&mut stdout, "{name}")?;
379    }
380
381    if opts.html {
382        if !opts.quiet {
383            writeln!(&mut stdout, "{HTML_HEAD_HEADER}")?;
384            writeln!(&mut stdout, "  <style>")?;
385            let names = theme.highlight_names.iter();
386            let styles = theme.styles.iter();
387            for (name, style) in names.zip(styles) {
388                if let Some(css) = &style.css {
389                    writeln!(&mut stdout, "    .{name} {{ {css}; }}")?;
390                }
391            }
392            writeln!(&mut stdout, "  </style>")?;
393            writeln!(&mut stdout, "{HTML_BODY_HEADER}")?;
394        }
395
396        let mut renderer = HtmlRenderer::new();
397        renderer.render(events, &source, &move |highlight, output| {
398            if opts.inline_styles {
399                output.extend(b"style='");
400                output.extend(
401                    theme.styles[highlight.0]
402                        .css
403                        .as_ref()
404                        .map_or_else(|| "".as_bytes(), |css_style| css_style.as_bytes()),
405                );
406                output.extend(b"'");
407            } else {
408                output.extend(b"class='");
409                let mut parts = theme.highlight_names[highlight.0].split('.').peekable();
410                while let Some(part) = parts.next() {
411                    output.extend(part.as_bytes());
412                    if parts.peek().is_some() {
413                        output.extend(b" ");
414                    }
415                }
416                output.extend(b"'");
417            }
418        })?;
419
420        if !opts.quiet {
421            writeln!(&mut stdout, "<table>")?;
422            for (i, line) in renderer.lines().enumerate() {
423                writeln!(
424                    &mut stdout,
425                    "<tr><td class=line-number>{}</td><td class=line>{line}</td></tr>",
426                    i + 1,
427                )?;
428            }
429            writeln!(&mut stdout, "</table>")?;
430            writeln!(&mut stdout, "{HTML_FOOTER}")?;
431        }
432    } else {
433        let mut style_stack = vec![theme.default_style().ansi];
434        for event in events {
435            match event? {
436                HighlightEvent::HighlightStart(highlight) => {
437                    style_stack.push(theme.styles[highlight.0].ansi);
438                }
439                HighlightEvent::HighlightEnd => {
440                    style_stack.pop();
441                }
442                HighlightEvent::Source { start, end } => {
443                    let style = style_stack.last().unwrap();
444                    write!(&mut stdout, "{style}").unwrap();
445                    stdout.write_all(&source[start..end])?;
446                    write!(&mut stdout, "{style:#}").unwrap();
447                }
448            }
449        }
450    }
451
452    if opts.print_time {
453        info!("Time: {}ms", time.elapsed().as_millis());
454    }
455
456    Ok(())
457}
458
459#[cfg(test)]
460mod tests {
461    use std::env;
462
463    use super::*;
464
465    const JUNGLE_GREEN: &str = "#26A69A";
466    const DARK_CYAN: &str = "#00AF87";
467
468    #[test]
469    fn test_parse_style() {
470        let original_environment_variable = env::var("COLORTERM");
471
472        let mut style = Style::default();
473        assert_eq!(style.ansi.get_fg_color(), None);
474        assert_eq!(style.css, None);
475
476        // darkcyan is an ANSI color and is preserved
477        env::set_var("COLORTERM", "");
478        parse_style(&mut style, Value::String(DARK_CYAN.to_string()));
479        assert_eq!(
480            style.ansi.get_fg_color(),
481            Some(Color::Ansi256(Ansi256Color(36)))
482        );
483        assert_eq!(style.css, Some("color: #00af87".to_string()));
484
485        // junglegreen is not an ANSI color and is preserved when the terminal supports it
486        env::set_var("COLORTERM", "truecolor");
487        parse_style(&mut style, Value::String(JUNGLE_GREEN.to_string()));
488        assert_eq!(
489            style.ansi.get_fg_color(),
490            Some(Color::Rgb(RgbColor(38, 166, 154)))
491        );
492        assert_eq!(style.css, Some("color: #26a69a".to_string()));
493
494        // junglegreen gets approximated as cadetblue when the terminal does not support it
495        env::set_var("COLORTERM", "");
496        parse_style(&mut style, Value::String(JUNGLE_GREEN.to_string()));
497        assert_eq!(
498            style.ansi.get_fg_color(),
499            Some(Color::Ansi256(Ansi256Color(72)))
500        );
501        assert_eq!(style.css, Some("color: #26a69a".to_string()));
502
503        if let Ok(environment_variable) = original_environment_variable {
504            env::set_var("COLORTERM", environment_variable);
505        } else {
506            env::remove_var("COLORTERM");
507        }
508    }
509}
tree_sitter_cli/highlight.rs

tree_sitter_cli/
highlight.rs
