1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94

use std::fmt;
use std::fmt::Write as _;
use std::io::Stdout;
use std::io::Write as _;
use syntect::easy::HighlightLines;
use syntect::highlighting::ThemeSet;
use syntect::parsing::SyntaxSet;
use syntect::util::{as_24_bit_terminal_escaped, LinesWithEndings};

use crate::{Highlight, Result};

struct Line(Option<usize>);

impl Highlight for Stdout {
    fn highlight(
        &mut self,
        text: &str,
        extension: &str,
        theme: Option<&str>,
    ) -> Result<(), crate::Error> {
        if atty::is(atty::Stream::Stdout) {
            writeln!(self, "{}", highlight_text(text, extension, theme)?)?;
        } else {
            writeln!(self, "{}", text)?;
        }

        Ok(())
    }
}

pub fn highlight_text(
    text: &str,
    extension: &str,
    theme: Option<&str>,
) -> Result<String, crate::Error> {
    // Load these once at the start of your program
    let ps = SyntaxSet::load_defaults_newlines();
    let ts = ThemeSet::load_defaults();

    let syntax = if let Some(s) = ps.find_syntax_by_extension(extension) {
        s
    } else {
        ps.find_syntax_plain_text()
    };
    let mut h = HighlightLines::new(syntax, &ts.themes[theme.unwrap_or("base16-ocean.dark")]);

    let mut output = String::new();

    for line in LinesWithEndings::from(text) {
        let ranges = h.highlight_line(line, &ps).unwrap();
        let escaped = as_24_bit_terminal_escaped(&ranges[..], false);
        write!(&mut output, "{}", escaped)?;
    }

    Ok(output)
}

impl fmt::Display for Line {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        match self.0 {
            None => write!(f, "    "),
            Some(idx) => write!(f, "{:<4}", idx + 1),
        }
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use serde_json::json;

    #[test]
    fn highlight_text_should_work() {
        let v = json!({
            "foo": "bar",
            "baz": "qux"
        });
        let text = serde_json::to_string_pretty(&v).unwrap();

        insta::assert_snapshot!(highlight_text(&text, "json", None).unwrap());
    }

    #[test]
    fn stdout_highlight_should_work() {
        let v = json!({
            "foo": "bar",
            "baz": "qux"
        });
        let text = serde_json::to_string_pretty(&v).unwrap();

        let mut stdout = std::io::stdout();
        stdout.highlight(&text, "json", None).unwrap();
    }
}