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

tree_sitter_cli/
highlight.rs
