Improve error handling by integrating anyhow crate

pkl-html-highlighter/Cargo.toml

@@ -8,6 +8,7 @@ authors = ["[email protected]"]
 tree-sitter-highlight = "0.20"
 tree-sitter-pkl = { git = "ssh://[email protected]/apple/tree-sitter-pkl.git", tag = "0.16.0" }
 clap = { version = "4.4.2", features = ["derive"] }
+anyhow = "1.0.40"
 
 [net]
 git-fetch-with-cli = true

pkl-html-highlighter/src/main.rs

@@ -13,12 +13,12 @@
 // See the License for the specific language governing permissions and
 // limitations under the License.
 // ===----------------------------------------------------------------------===//
-use std::process::exit;
-use std::{fs, io};
-use std::io::Read;
+use anyhow::{Context, Result};
 use clap::Parser;
+use std::fs;
+use std::io::{self, Read};
+use std::process::exit;
 use tree_sitter_highlight::{HighlightConfiguration, Highlighter, HtmlRenderer};
-use tree_sitter_pkl;
 
 #[derive(Parser)]
 struct CliArgs {
@@ -57,72 +57,111 @@ static QUERY_NAMES: [&str; 26] = [
     "title.function.invoke",
     "params",
     "comment",
-    "doctag"
+    "doctag",
 ];
 
-fn highlight_expr(contents: String, highlights_query: String) -> String {
-    let highlighted = highlight(format!("expr = {}", contents), highlights_query, false);
-    let stripped = highlighted.strip_prefix("<span class=\"hljs-property\">expr</span> <span class=\"hljs-operator\">=</span> ").unwrap();
-    return String::from(stripped);
+fn highlight_expr(contents: &str, highlights_query: &str) -> Result<String> {
+    let highlighted = highlight(&format!("expr = {}", contents), highlights_query, false)?;
+    let stripped = highlighted
+        .strip_prefix(
+            "<span class=\"hljs-property\">expr</span> <span class=\"hljs-operator\">=</span> ",
+        )
+        .context("Failed to strip prefix from highlighted output")?;
+    Ok(String::from(stripped))
 }
 
-pub fn highlight(contents: String, highlights_query: String, is_expr: bool) -> String {
+pub fn highlight(contents: &str, highlights_query: &str, is_expr: bool) -> Result<String> {
     // highlight each line as if it were an expression.
     if is_expr {
         return highlight_expr(contents, highlights_query);
     }
+
     let mut pkl_config = HighlightConfiguration::new(
         tree_sitter_pkl::language(),
-        highlights_query.as_str(),
+        highlights_query,
         tree_sitter_pkl::INJECTIONS_QUERY,
         tree_sitter_pkl::LOCALS_QUERY,
-    ).unwrap();
+    )?;
+
     let attrs = &QUERY_NAMES.map(|it| {
         let classes = it
-            .split(".")
+            .split('.')
             .enumerate()
             // The first scope is prefixed with hljs-.
             // Subscopes receive N number of underscores.
             // https://highlightjs.readthedocs.io/en/stable/css-classes-reference.html#a-note-on-scopes-with-sub-scopes
-            .map(|(idx, it)|
-                 if idx == 0 {
-                     format!("hljs-{}", it)
-                 } else {
-                     format!("{}{}", it, "_".repeat(idx))
-                 }
-            )
+            .map(|(idx, it)| {
+                if idx == 0 {
+                    format!("hljs-{}", it)
+                } else {
+                    format!("{}{}", it, "_".repeat(idx))
+                }
+            })
             .collect::<Vec<String>>()
             .join(" ");
         format!("class=\"{}\"", classes)
     });
     pkl_config.configure(&QUERY_NAMES);
 
-
     let mut highlighter = Highlighter::new();
-    let events = highlighter.highlight(
-        &pkl_config,
-        contents.as_bytes(),
-        None,
-        |_| None
-    ).unwrap();
+    let events = highlighter.highlight(&pkl_config, contents.as_bytes(), None, |_| None)?;
+
     let mut renderer = HtmlRenderer::new();
-    renderer.render(
-        events,
-        contents.as_bytes(),
-        &move |it| &attrs[it.0].as_bytes()
-    ).unwrap();
-    renderer.lines().collect::<Vec<&str>>().join("")
+    renderer.render(events, contents.as_bytes(), &move |it| {
+        attrs[it.0].as_bytes()
+    })?;
+    Ok(renderer.lines().collect::<Vec<&str>>().join(""))
 }
 
-fn main() {
+fn run() -> Result<()> {
     let args: CliArgs = CliArgs::parse();
-    if args.queries.is_none() {
-        print!("Missing required argument: --queries <query>");
-        exit(1)
-    }
-    let queries = fs::read_to_string(args.queries.unwrap().as_str()).unwrap();
+    let queries_path = args
+        .queries
+        .ok_or_else(|| anyhow::anyhow!("Missing required argument: --queries <query>"))?;
+    let queries = fs::read_to_string(&queries_path)
+        .with_context(|| format!("Failed to read queries from path: {}", queries_path))?;
+
     let mut buf = String::new();
-    io::stdin().read_to_string(&mut buf).expect("Failed to read stdin!");
-    let result = highlight(buf, queries, args.expr);
-    print!("{}", result)
+    io::stdin()
+        .read_to_string(&mut buf)
+        .context("Failed to read from stdin")?;
+    let result = highlight(&buf, &queries, args.expr)?;
+    print!("{}", result);
+    Ok(())
+}
+
+fn main() {
+    if let Err(err) = run() {
+        eprintln!("Error: {:#}", err);
+        exit(1);
+    }
+}
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+
+    static TEST_HIGHLIGHT_QUERY: &str = "(source_file) @keyword";
+
+    #[test]
+    fn test_highlight_expr() {
+        let expr = "let x = 42";
+        let result = highlight_expr(expr, TEST_HIGHLIGHT_QUERY).unwrap();
+        assert!(result.contains("hljs-operator"));
+    }
+
+    #[test]
+    fn test_highlight_document() {
+        let doc = "let x = 42";
+        let result = highlight(doc, TEST_HIGHLIGHT_QUERY, false).unwrap();
+        assert!(result.contains("hljs-operator"));
+    }
+
+    #[test]
+    fn test_highlight_fails_on_invalid_query() {
+        let doc = "let x = 42";
+        let invalid_query = "(invalid_query) @unknown";
+        let result = highlight(doc, invalid_query, false);
+        assert!(result.is_err());
+    }
 }