Skip to main content

mdpdf_core/
highlight.rs

1use std::collections::HashMap;
2use std::io::{self, Write};
3
4use comrak::adapters::SyntaxHighlighterAdapter;
5use syntect::highlighting::ThemeSet;
6use syntect::html::{styled_line_to_highlighted_html, IncludeBackground};
7use syntect::parsing::SyntaxSet;
8use syntect::easy::HighlightLines;
9use syntect::util::LinesWithEndings;
10
11/// Syntax highlighter backed by syntect, implementing comrak's adapter trait.
12pub struct SyntectHighlighter {
13    syntax_set: SyntaxSet,
14    theme_name: String,
15    theme_set: ThemeSet,
16}
17
18impl SyntectHighlighter {
19    pub fn new(theme_name: &str) -> Self {
20        Self {
21            syntax_set: SyntaxSet::load_defaults_newlines(),
22            theme_name: theme_name.to_owned(),
23            theme_set: ThemeSet::load_defaults(),
24        }
25    }
26
27    fn find_syntax(&self, lang: &str) -> Option<&syntect::parsing::SyntaxReference> {
28        // Try exact match first
29        self.syntax_set.find_syntax_by_token(lang)
30            // Common aliases
31            .or_else(|| match lang {
32                "tsx" | "typescriptreact" => self.syntax_set.find_syntax_by_token("typescript"),
33                "jsx" => self.syntax_set.find_syntax_by_token("javascript"),
34                "sh" | "shell" | "zsh" => self.syntax_set.find_syntax_by_token("bash"),
35                "yml" => self.syntax_set.find_syntax_by_token("yaml"),
36                "wit" => self.syntax_set.find_syntax_by_token("rust"),
37                _ => None,
38            })
39    }
40}
41
42impl SyntaxHighlighterAdapter for SyntectHighlighter {
43    fn write_highlighted(
44        &self,
45        output: &mut dyn Write,
46        lang: Option<&str>,
47        code: &str,
48    ) -> io::Result<()> {
49        let syntax = lang
50            .and_then(|l| self.find_syntax(l))
51            .unwrap_or_else(|| self.syntax_set.find_syntax_plain_text());
52
53        let theme = match self.theme_set.themes.get(&self.theme_name) {
54            Some(t) => t,
55            None => {
56                // Fallback: write plain text
57                return output.write_all(code.as_bytes());
58            }
59        };
60
61        let mut highlighter = HighlightLines::new(syntax, theme);
62
63        for line in LinesWithEndings::from(code) {
64            let ranges = highlighter
65                .highlight_line(line, &self.syntax_set)
66                .map_err(|e| io::Error::new(io::ErrorKind::Other, e))?;
67            let html = styled_line_to_highlighted_html(&ranges[..], IncludeBackground::No)
68                .map_err(|e| io::Error::new(io::ErrorKind::Other, e))?;
69            output.write_all(html.as_bytes())?;
70        }
71
72        Ok(())
73    }
74
75    fn write_pre_tag(
76        &self,
77        output: &mut dyn Write,
78        attributes: HashMap<String, String>,
79    ) -> io::Result<()> {
80        // Emit extra attributes (e.g., data-lang) if present
81        let mut attrs = String::new();
82        for (k, v) in &attributes {
83            attrs.push_str(&format!(" {}=\"{}\"", k, v));
84        }
85        write!(output, "<pre class=\"code-block\"{attrs}>")
86    }
87
88    fn write_code_tag(
89        &self,
90        output: &mut dyn Write,
91        _attributes: HashMap<String, String>,
92    ) -> io::Result<()> {
93        write!(output, "<code>")
94    }
95}