use std::ops::Range;
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum Language {
Json,
Rust,
Python,
Typescript,
Tsx,
}
impl Language {
#[allow(clippy::should_implement_trait)]
pub fn from_str(s: &str) -> Option<Self> {
match s.to_lowercase().as_str() {
"json" => Some(Self::Json),
"rust" | "rs" => Some(Self::Rust),
"python" | "py" => Some(Self::Python),
"typescript" | "ts" | "javascript" | "js" => Some(Self::Typescript),
"tsx" | "jsx" => Some(Self::Tsx),
_ => None,
}
}
}
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum TokenKind {
Keyword,
String,
Number,
Comment,
Punctuation,
Operator,
Identifier,
Property,
Type,
Function,
Default,
}
#[derive(Debug, Clone)]
pub struct HighlightedToken {
pub range: Range<usize>,
pub kind: TokenKind,
}
impl HighlightedToken {
pub fn new(range: Range<usize>, kind: TokenKind) -> Self {
Self { range, kind }
}
}
pub fn css_class(kind: TokenKind) -> &'static str {
match kind {
TokenKind::Keyword => "tok-keyword",
TokenKind::String => "tok-string",
TokenKind::Number => "tok-number",
TokenKind::Comment => "tok-comment",
TokenKind::Punctuation => "tok-punctuation",
TokenKind::Operator => "tok-operator",
TokenKind::Identifier => "tok-identifier",
TokenKind::Property => "tok-property",
TokenKind::Type => "tok-type",
TokenKind::Function => "tok-function",
TokenKind::Default => "tok-default",
}
}
#[cfg(feature = "syntax-highlight")]
pub fn highlight(code: &str, language: Language) -> Option<Vec<HighlightedToken>> {
match language {
Language::Json => highlight_json(code),
Language::Rust => highlight_rust(code),
Language::Python => highlight_python(code),
Language::Typescript => highlight_typescript(code, false),
Language::Tsx => highlight_typescript(code, true),
}
}
#[cfg(not(feature = "syntax-highlight"))]
pub fn highlight(_code: &str, _language: Language) -> Option<Vec<HighlightedToken>> {
None
}
#[cfg(feature = "syntax-highlight")]
fn highlight_json(code: &str) -> Option<Vec<HighlightedToken>> {
use tree_sitter::Parser;
let mut parser = Parser::new();
let language = tree_sitter_json::LANGUAGE.into();
if let Err(e) = parser.set_language(&language) {
tracing::warn!(?e, "highlight_json: failed to set language");
return None;
}
let tree = match parser.parse(code, None) {
Some(t) => t,
None => {
tracing::warn!("highlight_json: parse returned None");
return None;
}
};
let root = tree.root_node();
let mut tokens = Vec::new();
collect_json_tokens(&root, &mut tokens);
tracing::debug!(token_count = tokens.len(), "highlight_json: success");
Some(tokens)
}
#[cfg(feature = "syntax-highlight")]
fn collect_json_tokens(node: &tree_sitter::Node, tokens: &mut Vec<HighlightedToken>) {
let kind = match node.kind() {
"string" => {
if let Some(parent) = node.parent() {
if parent.kind() == "pair" {
if let Some(first_child) = parent.child(0) {
if first_child.id() == node.id() {
Some(TokenKind::Property)
} else {
Some(TokenKind::String)
}
} else {
Some(TokenKind::String)
}
} else {
Some(TokenKind::String)
}
} else {
Some(TokenKind::String)
}
}
"number" => Some(TokenKind::Number),
"true" | "false" | "null" => Some(TokenKind::Keyword),
"{" | "}" | "[" | "]" | ":" | "," => Some(TokenKind::Punctuation),
_ => None,
};
if let Some(kind) = kind {
let range = node.byte_range();
tokens.push(HighlightedToken::new(range, kind));
}
let mut cursor = node.walk();
for child in node.children(&mut cursor) {
collect_json_tokens(&child, tokens);
}
}
#[cfg(feature = "syntax-highlight")]
fn highlight_rust(code: &str) -> Option<Vec<HighlightedToken>> {
use tree_sitter::Parser;
let mut parser = Parser::new();
let language = tree_sitter_rust::LANGUAGE.into();
if let Err(e) = parser.set_language(&language) {
tracing::warn!(?e, "highlight_rust: failed to set language");
return None;
}
let tree = match parser.parse(code, None) {
Some(t) => t,
None => {
tracing::warn!("highlight_rust: parse returned None");
return None;
}
};
let root = tree.root_node();
let mut tokens = Vec::new();
collect_rust_tokens(&root, &mut tokens);
tracing::debug!(token_count = tokens.len(), "highlight_rust: success");
Some(tokens)
}
#[cfg(feature = "syntax-highlight")]
fn collect_rust_tokens(node: &tree_sitter::Node, tokens: &mut Vec<HighlightedToken>) {
let kind = match node.kind() {
"let" | "mut" | "fn" | "pub" | "struct" | "enum" | "impl" | "trait" | "use" | "mod"
| "if" | "else" | "match" | "for" | "while" | "loop" | "return" | "break" | "continue"
| "const" | "static" | "type" | "where" | "as" | "in" | "ref" | "self" | "Self"
| "super" | "crate" | "async" | "await" | "dyn" | "move" | "unsafe" | "extern" => {
Some(TokenKind::Keyword)
}
"true" | "false" => Some(TokenKind::Keyword),
"string_literal" | "raw_string_literal" | "char_literal" => Some(TokenKind::String),
"integer_literal" | "float_literal" => Some(TokenKind::Number),
"line_comment" | "block_comment" => Some(TokenKind::Comment),
"type_identifier" | "primitive_type" => Some(TokenKind::Type),
"identifier" if is_function_name(node) => Some(TokenKind::Function),
"field_identifier" => Some(TokenKind::Property),
"{" | "}" | "[" | "]" | "(" | ")" | ";" | "," | "::" | ":" | "->" | "=>" => {
Some(TokenKind::Punctuation)
}
"=" | "+" | "-" | "*" | "/" | "%" | "&" | "|" | "^" | "!" | "<" | ">" | "==" | "!="
| "<=" | ">=" | "&&" | "||" | "+=" | "-=" | "*=" | "/=" | ".." | "..=" | "?" => {
Some(TokenKind::Operator)
}
_ => None,
};
if let Some(kind) = kind {
let range = node.byte_range();
tokens.push(HighlightedToken::new(range, kind));
}
let mut cursor = node.walk();
for child in node.children(&mut cursor) {
collect_rust_tokens(&child, tokens);
}
}
#[cfg(feature = "syntax-highlight")]
fn is_function_name(node: &tree_sitter::Node) -> bool {
if let Some(parent) = node.parent() {
matches!(
parent.kind(),
"function_item" | "call_expression" | "method_call_expression"
)
} else {
false
}
}
#[cfg(feature = "syntax-highlight")]
fn highlight_python(code: &str) -> Option<Vec<HighlightedToken>> {
use tree_sitter::Parser;
let mut parser = Parser::new();
let language = tree_sitter_python::LANGUAGE.into();
parser.set_language(&language).ok()?;
let tree = parser.parse(code, None)?;
let root = tree.root_node();
let mut tokens = Vec::new();
collect_python_tokens(&root, &mut tokens);
Some(tokens)
}
#[cfg(feature = "syntax-highlight")]
fn collect_python_tokens(node: &tree_sitter::Node, tokens: &mut Vec<HighlightedToken>) {
let kind = match node.kind() {
"def" | "class" | "if" | "elif" | "else" | "for" | "while" | "try" | "except"
| "finally" | "with" | "as" | "import" | "from" | "return" | "yield" | "raise"
| "break" | "continue" | "pass" | "lambda" | "and" | "or" | "not" | "in" | "is"
| "global" | "nonlocal" | "assert" | "del" | "async" | "await" => Some(TokenKind::Keyword),
"true" | "false" | "none" | "True" | "False" | "None" => Some(TokenKind::Keyword),
"string" | "string_start" | "string_content" | "string_end" => Some(TokenKind::String),
"integer" | "float" => Some(TokenKind::Number),
"comment" => Some(TokenKind::Comment),
"identifier" if is_python_function_name(node) => Some(TokenKind::Function),
"attribute" => Some(TokenKind::Property),
"(" | ")" | "[" | "]" | "{" | "}" | ":" | "," | "." | "->" => Some(TokenKind::Punctuation),
"=" | "+" | "-" | "*" | "/" | "//" | "%" | "**" | "@" | "&" | "|" | "^" | "~" | "<"
| ">" | "<=" | ">=" | "==" | "!=" | "+=" | "-=" | "*=" | "/=" | "//=" | "%=" | "**="
| "&=" | "|=" | "^=" => Some(TokenKind::Operator),
_ => None,
};
if let Some(kind) = kind {
let range = node.byte_range();
tokens.push(HighlightedToken::new(range, kind));
}
let mut cursor = node.walk();
for child in node.children(&mut cursor) {
collect_python_tokens(&child, tokens);
}
}
#[cfg(feature = "syntax-highlight")]
fn is_python_function_name(node: &tree_sitter::Node) -> bool {
if let Some(parent) = node.parent() {
matches!(parent.kind(), "function_definition" | "call")
} else {
false
}
}
#[cfg(feature = "syntax-highlight")]
fn highlight_typescript(code: &str, tsx: bool) -> Option<Vec<HighlightedToken>> {
use tree_sitter::Parser;
let mut parser = Parser::new();
let language = if tsx {
tree_sitter_typescript::LANGUAGE_TSX.into()
} else {
tree_sitter_typescript::LANGUAGE_TYPESCRIPT.into()
};
if let Err(e) = parser.set_language(&language) {
tracing::warn!(?e, tsx, "highlight_typescript: failed to set language");
return None;
}
let tree = match parser.parse(code, None) {
Some(t) => t,
None => {
tracing::warn!(tsx, "highlight_typescript: parse returned None");
return None;
}
};
let root = tree.root_node();
let mut tokens = Vec::new();
collect_ts_tokens(&root, &mut tokens);
tracing::debug!(
token_count = tokens.len(),
tsx,
"highlight_typescript: success"
);
Some(tokens)
}
#[cfg(feature = "syntax-highlight")]
fn collect_ts_tokens(node: &tree_sitter::Node, tokens: &mut Vec<HighlightedToken>) {
let kind = match node.kind() {
"const" | "let" | "var" | "function" | "return" | "if" | "else" | "for" | "while"
| "do" | "switch" | "case" | "default" | "break" | "continue" | "class" | "interface"
| "type" | "enum" | "namespace" | "module" | "import" | "export" | "from" | "as"
| "extends" | "implements" | "new" | "delete" | "typeof" | "instanceof" | "in" | "of"
| "void" | "async" | "await" | "yield" | "throw" | "try" | "catch" | "finally"
| "public" | "private" | "protected" | "readonly" | "static" | "abstract" | "declare"
| "get" | "set" | "keyof" | "infer" | "satisfies" | "is" => Some(TokenKind::Keyword),
"true" | "false" | "null" | "undefined" => Some(TokenKind::Keyword),
"string" | "template_string" | "string_fragment" | "regex" => Some(TokenKind::String),
"number" => Some(TokenKind::Number),
"comment" => Some(TokenKind::Comment),
"type_identifier" | "predefined_type" => Some(TokenKind::Type),
"identifier" if is_ts_function_name(node) => Some(TokenKind::Function),
"identifier" if is_jsx_tag_name(node) => Some(TokenKind::Type),
"property_identifier" | "shorthand_property_identifier" => Some(TokenKind::Property),
"{" | "}" | "[" | "]" | "(" | ")" | ";" | "," | "." | ":" | "?." | "=>" | "<" | ">"
| "</" | "/>" => Some(TokenKind::Punctuation),
"=" | "+" | "-" | "*" | "/" | "%" | "**" | "&" | "|" | "^" | "~" | "!" | "==" | "==="
| "!=" | "!==" | "<=" | ">=" | "&&" | "||" | "??" | "+=" | "-=" | "*=" | "/=" | "%="
| "?" | "..." => Some(TokenKind::Operator),
_ => None,
};
if let Some(kind) = kind {
let range = node.byte_range();
tokens.push(HighlightedToken::new(range, kind));
}
let mut cursor = node.walk();
for child in node.children(&mut cursor) {
collect_ts_tokens(&child, tokens);
}
}
#[cfg(feature = "syntax-highlight")]
fn is_ts_function_name(node: &tree_sitter::Node) -> bool {
if let Some(parent) = node.parent() {
matches!(
parent.kind(),
"function_declaration"
| "function_expression"
| "generator_function_declaration"
| "call_expression"
| "method_definition"
| "function_signature"
)
} else {
false
}
}
#[cfg(feature = "syntax-highlight")]
fn is_jsx_tag_name(node: &tree_sitter::Node) -> bool {
if let Some(parent) = node.parent() {
matches!(
parent.kind(),
"jsx_opening_element" | "jsx_closing_element" | "jsx_self_closing_element"
)
} else {
false
}
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn test_language_from_str() {
assert_eq!(Language::from_str("json"), Some(Language::Json));
assert_eq!(Language::from_str("JSON"), Some(Language::Json));
assert_eq!(Language::from_str("rust"), Some(Language::Rust));
assert_eq!(Language::from_str("rs"), Some(Language::Rust));
assert_eq!(Language::from_str("python"), Some(Language::Python));
assert_eq!(Language::from_str("py"), Some(Language::Python));
assert_eq!(Language::from_str("typescript"), Some(Language::Typescript));
assert_eq!(Language::from_str("ts"), Some(Language::Typescript));
assert_eq!(Language::from_str("js"), Some(Language::Typescript));
assert_eq!(Language::from_str("tsx"), Some(Language::Tsx));
assert_eq!(Language::from_str("jsx"), Some(Language::Tsx));
assert_eq!(Language::from_str("unknown"), None);
}
#[test]
fn test_css_class_distinct() {
assert_eq!(css_class(TokenKind::Keyword), "tok-keyword");
assert_ne!(css_class(TokenKind::Keyword), css_class(TokenKind::String));
}
#[cfg(feature = "syntax-highlight")]
#[test]
fn test_highlight_json() {
let code = r#"{"key": "value", "num": 42, "flag": true}"#;
let tokens = highlight(code, Language::Json).expect("should highlight JSON");
assert!(!tokens.is_empty());
let property_tokens: Vec<_> = tokens
.iter()
.filter(|t| t.kind == TokenKind::Property)
.collect();
assert!(!property_tokens.is_empty(), "should have property tokens");
let number_tokens: Vec<_> = tokens
.iter()
.filter(|t| t.kind == TokenKind::Number)
.collect();
assert_eq!(number_tokens.len(), 1, "should have one number token");
let keyword_tokens: Vec<_> = tokens
.iter()
.filter(|t| t.kind == TokenKind::Keyword)
.collect();
assert_eq!(
keyword_tokens.len(),
1,
"should have one keyword token (true)"
);
}
#[cfg(feature = "syntax-highlight")]
#[test]
fn test_highlight_rust() {
let code = r#"fn main() { let x = 42; }"#;
let tokens = highlight(code, Language::Rust).expect("should highlight Rust");
assert!(!tokens.is_empty());
let keyword_tokens: Vec<_> = tokens
.iter()
.filter(|t| t.kind == TokenKind::Keyword)
.collect();
assert!(
keyword_tokens.len() >= 2,
"should have at least fn and let keywords"
);
}
#[cfg(feature = "syntax-highlight")]
#[test]
fn test_highlight_typescript() {
let code = r#"const greeting: string = "hello"; function add(a: number) { return a; }"#;
let tokens = highlight(code, Language::Typescript).expect("should highlight TS");
assert!(!tokens.is_empty());
let has_keyword = tokens.iter().any(|t| t.kind == TokenKind::Keyword);
let has_string = tokens.iter().any(|t| t.kind == TokenKind::String);
let has_type = tokens.iter().any(|t| t.kind == TokenKind::Type);
assert!(
has_keyword,
"should classify const/function/return as keywords"
);
assert!(has_string, "should classify the string literal");
assert!(has_type, "should classify the `string`/`number` types");
}
#[cfg(feature = "syntax-highlight")]
#[test]
fn test_highlight_tsx() {
let code = r#"const App = () => <div className="x">{label}</div>;"#;
let tokens = highlight(code, Language::Tsx).expect("should highlight TSX");
assert!(!tokens.is_empty());
let has_type = tokens.iter().any(|t| t.kind == TokenKind::Type);
let has_property = tokens.iter().any(|t| t.kind == TokenKind::Property);
assert!(has_type, "JSX element name should be a Type token");
assert!(
has_property,
"JSX attribute name should be a Property token"
);
}
#[cfg(not(feature = "syntax-highlight"))]
#[test]
fn test_highlight_returns_none_without_feature() {
assert!(highlight("{}", Language::Json).is_none());
}
}