Skip to main content

rig_tui/
highlight.rs

1//! Code in the terminal's own colors: bat's grammars (`two-face`) and its
2//! `ansi` theme, which names colors of the terminal's palette rather than
3//! colors of its own, so highlighted code follows the terminal's theme.
4//! The grammars load once, on a thread of their own when the terminal
5//! opens, or when code is first drawn.
6//!
7//! ```
8//! use rig_tui::highlight::{Language, code_rows};
9//!
10//! let rows = code_rows("fn main() {}\n", Language::File("src/main.rs"));
11//! assert_eq!(rows.len(), 1);
12//! ```
13
14use std::path::Path;
15use std::sync::LazyLock;
16
17use ratatui::style::{Color, Style};
18use ratatui::text::{Line, Span};
19use syntect::easy::HighlightLines;
20use syntect::highlighting::FontStyle;
21use syntect::parsing::{SyntaxReference, SyntaxSet};
22use syntect::util::LinesWithEndings;
23
24use crate::{OUTPUT_INDENT, text_rows};
25
26static SYNTAXES: LazyLock<SyntaxSet> = LazyLock::new(two_face::syntax::extra_newlines);
27
28static THEME: LazyLock<syntect::highlighting::Theme> = LazyLock::new(|| {
29    let themes = two_face::theme::extra();
30    themes.get(two_face::theme::EmbeddedThemeName::Ansi).clone()
31});
32
33/// What code is in.
34#[derive(Clone, Copy, Debug)]
35pub enum Language<'a> {
36    /// A language's name or file extension, as a code block's fence says
37    /// it: `rust`, `py`, `bash`.
38    Named(&'a str),
39    /// A file's path, by its extension or its name (`Makefile`).
40    File(&'a str),
41}
42
43/// Loads the grammars and the theme on a thread of their own, so the
44/// first code drawn does not wait for them.
45pub(crate) fn warm_up() {
46    std::thread::spawn(|| {
47        LazyLock::force(&SYNTAXES);
48        LazyLock::force(&THEME);
49    });
50}
51
52/// `code`'s lines in `language`'s colors, each as its spans, tabs made
53/// spaces; `None` for a language with no grammar.
54pub fn highlight(code: &str, language: Language<'_>) -> Option<Vec<Vec<Span<'static>>>> {
55    let syntax = syntax(language)?;
56    let mut lines = HighlightLines::new(syntax, &THEME);
57    LinesWithEndings::from(code)
58        .map(|line| {
59            let pieces = lines.highlight_line(line, &SYNTAXES).ok()?;
60            let spans = pieces.into_iter().filter_map(|(style, text)| {
61                let text = text.trim_end_matches(['\n', '\r']).replace('\t', "    ");
62                (!text.is_empty()).then(|| Span::styled(text, style_of(style)))
63            });
64            Some(spans.collect())
65        })
66        .collect()
67}
68
69/// `code`'s rows as a tool's result draws them, indented by
70/// [`OUTPUT_INDENT`],
71/// in `language`'s colors, or dimmed for a language with no grammar.
72pub fn code_rows(code: &str, language: Language<'_>) -> Vec<Line<'static>> {
73    let Some(lines) = highlight(code, language) else {
74        return text_rows(code, Style::new().dim());
75    };
76    let rows = lines
77        .into_iter()
78        .map(|spans| Line::from([vec![Span::from(OUTPUT_INDENT)], spans].concat()));
79    rows.collect()
80}
81
82/// The grammar of `language`, if one is known; plain text is none.
83fn syntax(language: Language<'_>) -> Option<&'static SyntaxReference> {
84    let syntaxes: &'static SyntaxSet = &SYNTAXES;
85    let found = match language {
86        Language::Named(name) => syntaxes.find_syntax_by_token(name.trim()),
87        Language::File(path) => {
88            let path = Path::new(path);
89            let by =
90                |name: Option<&std::ffi::OsStr>| syntaxes.find_syntax_by_extension(name?.to_str()?);
91            by(path.extension()).or_else(|| by(path.file_name()))
92        }
93    };
94    found.filter(|syntax| syntax.name != "Plain Text")
95}
96
97/// A piece's style: its color and its weight, slant and underline.
98fn style_of(style: syntect::highlighting::Style) -> Style {
99    let mut drawn = Style::new();
100    if let Some(color) = color(style.foreground) {
101        drawn = drawn.fg(color);
102    }
103    let font = style.font_style;
104    if font.contains(FontStyle::BOLD) {
105        drawn = drawn.bold();
106    }
107    if font.contains(FontStyle::ITALIC) {
108        drawn = drawn.italic();
109    }
110    if font.contains(FontStyle::UNDERLINE) {
111        drawn = drawn.underlined();
112    }
113    drawn
114}
115
116/// A theme color as the `ansi` theme writes it: alpha 0 names a color of
117/// the terminal's palette by its red, alpha 1 the terminal's own text
118/// color (`None`), any other alpha a 24-bit color.
119fn color(color: syntect::highlighting::Color) -> Option<Color> {
120    match color.a {
121        0 => Some(Color::Indexed(color.r)),
122        1 => None,
123        _ => Some(Color::Rgb(color.r, color.g, color.b)),
124    }
125}
126
127#[cfg(test)]
128mod tests;