Skip to main content

run/
highlight.rs

1use once_cell::sync::Lazy;
2use syntect::easy::HighlightLines;
3use syntect::highlighting::ThemeSet;
4use syntect::parsing::{SyntaxReference, SyntaxSet};
5use syntect::util::{LinesWithEndings, as_24_bit_terminal_escaped};
6
7static SYNTAX_SET: Lazy<SyntaxSet> = Lazy::new(SyntaxSet::load_defaults_newlines);
8
9static THEME_SET: Lazy<ThemeSet> = Lazy::new(ThemeSet::load_defaults);
10
11fn supports_color() -> bool {
12    std::env::var("NO_COLOR").is_err()
13        && (std::env::var("TERM").unwrap_or_default().contains("color")
14            || std::env::var("COLORTERM").is_ok())
15}
16
17fn language_to_syntax_name(language_id: &str) -> &str {
18    match language_id.to_ascii_lowercase().as_str() {
19        "python" | "py" | "python3" | "py3" => "Python",
20        "javascript" | "js" | "node" | "nodejs" => "JavaScript",
21        "typescript" | "ts" => "JavaScript",
22        "rust" | "rs" => "Rust",
23        "go" | "golang" => "Go",
24        "c" => "C",
25        "cpp" | "c++" | "cxx" => "C++",
26        "java" => "Java",
27        "csharp" | "cs" | "c#" => "C#",
28        "ruby" | "rb" => "Ruby",
29        "php" => "PHP",
30        "bash" | "sh" | "shell" | "zsh" => "Bourne Again Shell (bash)",
31        "lua" => "Lua",
32        "perl" | "pl" => "Perl",
33        "swift" => "Swift",
34        "kotlin" | "kt" => "Kotlin",
35        "r" | "rscript" => "R",
36        "haskell" | "hs" => "Haskell",
37        "julia" | "jl" => "Plain Text",
38        "elixir" | "ex" | "exs" => "Plain Text",
39        "dart" => "Dart",
40        "groovy" | "grv" => "Groovy",
41        "crystal" | "cr" => "Crystal",
42        "zig" => "Zig",
43        "nim" => "Nim",
44        _ => "Plain Text",
45    }
46}
47
48fn get_syntax_for_language(language_id: &str) -> &'static SyntaxReference {
49    let syntax_name = language_to_syntax_name(language_id);
50    SYNTAX_SET
51        .find_syntax_by_name(syntax_name)
52        .unwrap_or_else(|| SYNTAX_SET.find_syntax_plain_text())
53}
54
55pub fn highlight_code(code: &str, language_id: &str) -> String {
56    if !supports_color() {
57        return code.to_string();
58    }
59
60    let syntax = get_syntax_for_language(language_id);
61
62    let Some(theme) = THEME_SET
63        .themes
64        .get("base16-ocean.dark")
65        .or_else(|| THEME_SET.themes.values().next())
66    else {
67        return code.to_string();
68    };
69
70    let mut highlighter = HighlightLines::new(syntax, theme);
71    let mut output = String::new();
72
73    for line in LinesWithEndings::from(code) {
74        let ranges = highlighter
75            .highlight_line(line, &SYNTAX_SET)
76            .unwrap_or_default();
77        let escaped = as_24_bit_terminal_escaped(&ranges[..], false);
78        output.push_str(&escaped);
79    }
80
81    if !output.is_empty() && !output.ends_with("\x1b[0m") {
82        output.push_str("\x1b[0m");
83    }
84
85    output
86}
87
88pub fn highlight_repl_input(code: &str, language_id: &str) -> String {
89    highlight_code(code, language_id)
90}
91
92pub fn highlight_output(code: &str, language_id: &str) -> String {
93    highlight_code(code, language_id)
94}
95
96pub fn has_syntax_support(language_id: &str) -> bool {
97    let syntax_name = language_to_syntax_name(language_id);
98    SYNTAX_SET.find_syntax_by_name(syntax_name).is_some()
99}
100
101#[cfg(test)]
102mod tests {
103    use super::*;
104
105    #[test]
106    fn test_language_mapping() {
107        assert_eq!(language_to_syntax_name("python"), "Python");
108        assert_eq!(language_to_syntax_name("rust"), "Rust");
109        assert_eq!(language_to_syntax_name("javascript"), "JavaScript");
110        assert_eq!(language_to_syntax_name("typescript"), "JavaScript");
111        assert_eq!(language_to_syntax_name("go"), "Go");
112        assert_eq!(language_to_syntax_name("java"), "Java");
113        assert_eq!(language_to_syntax_name("csharp"), "C#");
114        assert_eq!(language_to_syntax_name("cpp"), "C++");
115        assert_eq!(language_to_syntax_name("ruby"), "Ruby");
116        assert_eq!(language_to_syntax_name("php"), "PHP");
117    }
118
119    #[test]
120    fn test_all_language_aliases() {
121        assert_eq!(language_to_syntax_name("py"), "Python");
122        assert_eq!(language_to_syntax_name("py3"), "Python");
123        assert_eq!(language_to_syntax_name("python3"), "Python");
124
125        assert_eq!(language_to_syntax_name("js"), "JavaScript");
126        assert_eq!(language_to_syntax_name("node"), "JavaScript");
127        assert_eq!(language_to_syntax_name("nodejs"), "JavaScript");
128
129        assert_eq!(language_to_syntax_name("ts"), "JavaScript");
130
131        assert_eq!(language_to_syntax_name("rs"), "Rust");
132
133        assert_eq!(language_to_syntax_name("golang"), "Go");
134
135        assert_eq!(language_to_syntax_name("c++"), "C++");
136        assert_eq!(language_to_syntax_name("cxx"), "C++");
137
138        assert_eq!(language_to_syntax_name("cs"), "C#");
139        assert_eq!(language_to_syntax_name("c#"), "C#");
140    }
141
142    #[test]
143    fn test_all_languages_have_syntax() {
144        let supported = vec![
145            "python",
146            "javascript",
147            "rust",
148            "go",
149            "c",
150            "cpp",
151            "java",
152            "csharp",
153            "ruby",
154            "php",
155            "bash",
156            "lua",
157            "perl",
158            "r",
159            "haskell",
160        ];
161
162        for lang in supported {
163            assert!(
164                has_syntax_support(lang),
165                "Language {} should be supported",
166                lang
167            );
168        }
169
170        let fallback = vec![
171            "swift", "kotlin", "dart", "groovy", "crystal", "zig", "nim", "julia", "elixir",
172        ];
173        for lang in fallback {
174            let _ = highlight_code("test", lang);
175        }
176    }
177
178    #[test]
179    fn test_unknown_language_fallback() {
180        let syntax_name = language_to_syntax_name("unknownlang123");
181        assert_eq!(syntax_name, "Plain Text");
182    }
183
184    #[test]
185    fn test_syntax_available() {
186        assert!(has_syntax_support("python"));
187        assert!(has_syntax_support("rust"));
188        assert!(has_syntax_support("go"));
189        assert!(has_syntax_support("javascript"));
190        assert!(has_syntax_support("typescript"));
191    }
192
193    #[test]
194    fn test_highlight_basic() {
195        let code = "fn main() { println!(\"Hello\"); }";
196        let highlighted = highlight_code(code, "rust");
197        assert!(!highlighted.is_empty());
198        assert!(highlighted.len() >= code.len());
199    }
200
201    #[test]
202    fn test_highlight_python() {
203        let code = "def hello():\n    print('world')";
204        let highlighted = highlight_code(code, "python");
205        assert!(!highlighted.is_empty());
206    }
207
208    #[test]
209    fn test_highlight_javascript() {
210        let code = "function hello() { console.log('world'); }";
211        let highlighted = highlight_code(code, "javascript");
212        assert!(!highlighted.is_empty());
213    }
214
215    #[test]
216    fn test_highlight_go() {
217        let code = "package main\nfunc main() { fmt.Println(\"hello\") }";
218        let highlighted = highlight_code(code, "go");
219        assert!(!highlighted.is_empty());
220    }
221
222    #[test]
223    fn test_empty_code() {
224        let highlighted = highlight_code("", "python");
225        assert!(highlighted.is_empty() || highlighted == "\x1b[0m");
226    }
227
228    #[test]
229    fn test_whitespace_only() {
230        let code = "   \n  \t  ";
231        let highlighted = highlight_code(code, "rust");
232        assert!(!highlighted.is_empty());
233    }
234
235    #[test]
236    fn test_multiline_code() {
237        let code = "fn main() {\n    let x = 10;\n    println!(\"{}\", x);\n}";
238        let highlighted = highlight_code(code, "rust");
239        assert!(!highlighted.is_empty());
240        assert!(highlighted.contains('\n') || highlighted.contains("\\n"));
241    }
242
243    #[test]
244    fn test_color_reset_at_end() {
245        unsafe {
246            std::env::set_var("TERM", "xterm-256color");
247        }
248
249        let code = "x = 10";
250        let highlighted = highlight_code(code, "python");
251
252        assert!(
253            highlighted.ends_with("\x1b[0m") || !highlighted.contains("\x1b["),
254            "Highlighted code should end with color reset or be plain text"
255        );
256    }
257
258    #[test]
259    fn test_no_color_environment() {
260        unsafe {
261            std::env::set_var("NO_COLOR", "1");
262        }
263
264        let code = "fn main() {}";
265        let highlighted = highlight_code(code, "rust");
266
267        assert_eq!(highlighted, code);
268
269        unsafe {
270            std::env::remove_var("NO_COLOR");
271        }
272    }
273
274    #[test]
275    fn test_special_characters() {
276        let code = "print(\"Hello\\nWorld\\t!\")";
277        let highlighted = highlight_code(code, "python");
278        assert!(!highlighted.is_empty());
279    }
280
281    #[test]
282    fn test_string_literal() {
283        let code = "message = \"Hello World\"";
284        let highlighted = highlight_code(code, "python");
285        assert!(!highlighted.is_empty());
286    }
287
288    #[test]
289    fn test_repl_input_helper() {
290        let code = "x = 42";
291        let highlighted = highlight_repl_input(code, "python");
292        assert!(!highlighted.is_empty());
293    }
294
295    #[test]
296    fn test_output_helper() {
297        let code = "console.log('test')";
298        let highlighted = highlight_output(code, "javascript");
299        assert!(!highlighted.is_empty());
300    }
301
302    #[test]
303    fn test_bash_highlighting() {
304        let code = "echo \"Hello World\"";
305        let highlighted = highlight_code(code, "bash");
306        assert!(!highlighted.is_empty());
307    }
308
309    #[test]
310    fn test_sql_like_code_in_supported_language() {
311        let code = "query = \"SELECT * FROM users\"";
312        let highlighted = highlight_code(code, "python");
313        assert!(!highlighted.is_empty());
314    }
315
316    #[test]
317    fn test_comments_highlighting() {
318        let code = "// This is a comment\nlet x = 10;";
319        let highlighted = highlight_code(code, "javascript");
320        assert!(!highlighted.is_empty());
321    }
322
323    #[test]
324    fn test_string_highlighting() {
325        let code = "\"This is a string\"";
326        let highlighted = highlight_code(code, "python");
327        assert!(!highlighted.is_empty());
328    }
329
330    #[test]
331    fn test_number_highlighting() {
332        let code = "let num = 42;";
333        let highlighted = highlight_code(code, "javascript");
334        assert!(!highlighted.is_empty());
335    }
336
337    #[test]
338    fn test_very_long_code() {
339        let mut code = String::new();
340        for i in 0..1000 {
341            code.push_str(&format!("let var{} = {};\n", i, i));
342        }
343        let highlighted = highlight_code(&code, "javascript");
344        assert!(!highlighted.is_empty());
345    }
346
347    #[test]
348    fn test_case_insensitive_language_names() {
349        assert_eq!(language_to_syntax_name("PYTHON"), "Python");
350        assert_eq!(language_to_syntax_name("Python"), "Python");
351        assert_eq!(language_to_syntax_name("RuSt"), "Rust");
352        assert_eq!(language_to_syntax_name("JavaScript"), "JavaScript");
353    }
354
355    #[test]
356    fn test_syntax_reference_caching() {
357        let syntax1 = get_syntax_for_language("python");
358        let syntax2 = get_syntax_for_language("python");
359        assert_eq!(syntax1.name, syntax2.name);
360    }
361
362    #[test]
363    fn test_all_functional_languages() {
364        let langs = vec!["haskell", "elixir", "julia"];
365        for lang in langs {
366            let code = "main = print \"hello\"";
367            let highlighted = highlight_code(code, lang);
368            assert!(!highlighted.is_empty(), "Failed for {}", lang);
369        }
370    }
371
372    #[test]
373    fn test_systems_languages() {
374        let langs = vec!["c", "cpp", "rust", "zig", "nim"];
375        for lang in langs {
376            let code = "int main() { return 0; }";
377            let highlighted = highlight_code(code, lang);
378            assert!(!highlighted.is_empty(), "Failed for {}", lang);
379        }
380    }
381
382    #[test]
383    fn test_scripting_languages() {
384        let langs = vec!["python", "ruby", "perl", "lua", "php"];
385        for lang in langs {
386            let code = "print('hello')";
387            let highlighted = highlight_code(code, lang);
388            assert!(!highlighted.is_empty(), "Failed for {}", lang);
389        }
390    }
391
392    #[test]
393    fn test_jvm_languages() {
394        let langs = vec!["java", "kotlin", "groovy"];
395        for lang in langs {
396            let code = "public class Test { }";
397            let highlighted = highlight_code(code, lang);
398            assert!(!highlighted.is_empty(), "Failed for {}", lang);
399        }
400    }
401
402    #[test]
403    fn test_ansi_codes_present_when_colors_enabled() {
404        unsafe {
405            std::env::set_var("TERM", "xterm-256color");
406            std::env::remove_var("NO_COLOR");
407        }
408
409        let code = "fn main() {}";
410        let highlighted = highlight_code(code, "rust");
411
412        let has_ansi = highlighted.contains("\x1b[") || highlighted == code;
413        assert!(has_ansi, "Should contain ANSI codes or be plain text");
414    }
415
416    #[test]
417    fn test_typescript_mapping() {
418        assert_eq!(language_to_syntax_name("typescript"), "JavaScript");
419        let code = "const x: number = 42;";
420        let highlighted = highlight_code(code, "typescript");
421        assert!(!highlighted.is_empty());
422    }
423
424    #[test]
425    fn test_swift_language() {
426        let code = "func main() { print(\"Hello\") }";
427        let highlighted = highlight_code(code, "swift");
428        assert!(!highlighted.is_empty());
429    }
430
431    #[test]
432    fn test_kotlin_language() {
433        let code = "fun main() { println(\"Hello\") }";
434        let highlighted = highlight_code(code, "kotlin");
435        assert!(!highlighted.is_empty());
436    }
437
438    #[test]
439    fn test_dart_language() {
440        let code = "void main() { print('Hello'); }";
441        let highlighted = highlight_code(code, "dart");
442        assert!(!highlighted.is_empty());
443    }
444
445    #[test]
446    fn test_r_language() {
447        let code = "x <- 10\nprint(x)";
448        let highlighted = highlight_code(code, "r");
449        assert!(!highlighted.is_empty());
450    }
451
452    #[test]
453    fn test_crystal_language() {
454        let code = "puts \"Hello\"";
455        let highlighted = highlight_code(code, "crystal");
456        assert!(!highlighted.is_empty());
457    }
458
459    #[test]
460    fn test_zig_language() {
461        let code = "pub fn main() void {}";
462        let highlighted = highlight_code(code, "zig");
463        assert!(!highlighted.is_empty());
464    }
465
466    #[test]
467    fn test_nim_language() {
468        let code = "echo \"Hello\"";
469        let highlighted = highlight_code(code, "nim");
470        assert!(!highlighted.is_empty());
471    }
472
473    #[test]
474    fn test_only_comments() {
475        let code = "// Just a comment\n// Another comment";
476        let highlighted = highlight_code(code, "rust");
477        assert!(!highlighted.is_empty());
478    }
479
480    #[test]
481    fn test_code_with_syntax_errors() {
482        let code = "fn main( { println! }";
483        let highlighted = highlight_code(code, "rust");
484        assert!(!highlighted.is_empty());
485    }
486
487    #[test]
488    fn test_mixed_quotes() {
489        let code = r#"s1 = 'single'; s2 = "double""#;
490        let highlighted = highlight_code(code, "python");
491        assert!(!highlighted.is_empty());
492    }
493
494    #[test]
495    fn test_nested_structures() {
496        let code = "let arr = [[1, 2], [3, 4], [5, 6]];";
497        let highlighted = highlight_code(code, "javascript");
498        assert!(!highlighted.is_empty());
499    }
500
501    #[test]
502    fn test_regex_patterns() {
503        let code = r"pattern = /[a-z]+/g";
504        let highlighted = highlight_code(code, "javascript");
505        assert!(!highlighted.is_empty());
506    }
507
508    #[test]
509    fn test_multiple_statements_one_line() {
510        let code = "x = 1; y = 2; z = 3;";
511        let highlighted = highlight_code(code, "javascript");
512        assert!(!highlighted.is_empty());
513    }
514
515    #[test]
516    fn test_color_reset_prevents_bleeding() {
517        unsafe {
518            std::env::set_var("TERM", "xterm-256color");
519            std::env::remove_var("NO_COLOR");
520        }
521
522        let code1 = "x = 20";
523        let highlighted1 = highlight_code(code1, "javascript");
524
525        if highlighted1.contains("\x1b[") {
526            assert!(
527                highlighted1.ends_with("\x1b[0m"),
528                "Highlighted code must end with reset code to prevent color bleeding"
529            );
530        }
531    }
532
533    #[test]
534    fn test_all_25_languages_work() {
535        let languages = vec![
536            "python",
537            "javascript",
538            "typescript",
539            "rust",
540            "go",
541            "c",
542            "cpp",
543            "java",
544            "csharp",
545            "ruby",
546            "php",
547            "bash",
548            "lua",
549            "perl",
550            "swift",
551            "kotlin",
552            "r",
553            "dart",
554            "haskell",
555            "julia",
556            "elixir",
557            "groovy",
558            "crystal",
559            "zig",
560            "nim",
561        ];
562
563        for lang in languages {
564            let code = "test code";
565            let highlighted = highlight_code(code, lang);
566            assert!(
567                !highlighted.is_empty(),
568                "Language {} failed to highlight",
569                lang
570            );
571        }
572    }
573
574    #[test]
575    fn test_bash_aliases() {
576        let aliases = vec!["bash", "sh", "shell", "zsh"];
577        for alias in aliases {
578            assert_eq!(
579                language_to_syntax_name(alias),
580                "Bourne Again Shell (bash)",
581                "Bash alias {} failed",
582                alias
583            );
584        }
585    }
586
587    #[test]
588    fn test_empty_lines_in_code() {
589        let code = "fn main() {\n\n\n    println!(\"test\");\n\n}";
590        let highlighted = highlight_code(code, "rust");
591        assert!(!highlighted.is_empty());
592
593        assert!(highlighted.matches('\n').count() >= 4);
594    }
595
596    #[test]
597    fn test_tabs_and_spaces() {
598        let code = "\tfn main() {\n\t\tprintln!(\"test\");\n\t}";
599        let highlighted = highlight_code(code, "rust");
600        assert!(!highlighted.is_empty());
601    }
602
603    #[test]
604    fn test_colorterm_environment() {
605        unsafe {
606            std::env::set_var("COLORTERM", "truecolor");
607            std::env::remove_var("NO_COLOR");
608        }
609
610        assert!(supports_color());
611
612        unsafe {
613            std::env::remove_var("COLORTERM");
614        }
615    }
616
617    #[test]
618    fn test_repl_helpers_consistency() {
619        let code = "test = 42";
620        let h1 = highlight_code(code, "python");
621        let h2 = highlight_repl_input(code, "python");
622        let h3 = highlight_output(code, "python");
623
624        assert_eq!(h1, h2);
625        assert_eq!(h2, h3);
626    }
627}