templates: add raw_text to directly output content

Templates can be formatted (using labels) and are usually sanitized (unless for plain text output). `raw_text(content)` bypasses both. ```toml 'hyperlink(url, text)' = ''' raw_text("\e]8;;" ++ url ++ "\e\\") ++ text ++ raw_text("\e]8;;\e\\") ''' ``` In this example, `raw_text` not only outputs the intended escape codes, it also strips away any escape codes that might otherwise be part of the `url` (from any labels attached to the `url` content). Change-Id: Id000000040ea6fd8e2d720219931485960c570dd
martinvonz · Oct 8, 2024 · beb0986 · beb0986
1 parent 68f4860
commit beb0986
Show file tree

Hide file tree

Showing 8 changed files with 144 additions and 25 deletions.
diff --git a/cli/src/formatter.rs b/cli/src/formatter.rs
@@ -578,13 +578,15 @@ impl<W: Write> Drop for ColorFormatter<W> {
 #[derive(Clone, Debug, Default)]
 pub struct FormatRecorder {
     data: Vec<u8>,
-    label_ops: Vec<(usize, LabelOp)>,
+    ops: Vec<(usize, FormatOp)>,
 }
 
 #[derive(Clone, Debug, Eq, PartialEq)]
-enum LabelOp {
+enum FormatOp {
     PushLabel(String),
     PopLabel,
+    PushRawText,
+    PopRawText,
 }
 
 impl FormatRecorder {
@@ -596,8 +598,8 @@ impl FormatRecorder {
         &self.data
     }
 
-    fn push_label_op(&mut self, op: LabelOp) {
-        self.label_ops.push((self.data.len(), op));
+    fn push_op(&mut self, op: FormatOp) {
+        self.ops.push((self.data.len(), op));
     }
 
     pub fn replay(&self, formatter: &mut dyn Formatter) -> io::Result<()> {
@@ -612,21 +614,36 @@ impl FormatRecorder {
         mut write_data: impl FnMut(&mut dyn Formatter, Range<usize>) -> io::Result<()>,
     ) -> io::Result<()> {
         let mut last_pos = 0;
-        let mut flush_data = |formatter: &mut dyn Formatter, pos| -> io::Result<()> {
-            if last_pos != pos {
+        let mut flush_data = |formatter: &mut dyn Formatter, pos, raw_text| -> io::Result<()> {
+            if pos == last_pos {
+                return Ok(());
+            }
+            if raw_text {
+                let mut plain_text_formatter = PlainTextFormatter::new(formatter.raw());
+                write_data(&mut plain_text_formatter, last_pos..pos)?;
+            } else {
                 write_data(formatter, last_pos..pos)?;
-                last_pos = pos;
             }
+            last_pos = pos;
             Ok(())
         };
-        for (pos, op) in &self.label_ops {
-            flush_data(formatter, *pos)?;
+        let mut last_op = &FormatOp::PopRawText;
+        for (pos, op) in &self.ops {
+            // Use the {Push/Pop}RawText ops recorded by RawTextRecorder
+            // (returned by FormatRecorder.raw()) to identify when to write to
+            // the underlying raw output of the formatter.
+            let raw_text = *last_op == FormatOp::PushRawText;
+            // PushRawText has to be immediately followed by PopRawText.
+            assert!(raw_text == (*op == FormatOp::PopRawText));
+            flush_data(formatter, *pos, *last_op == FormatOp::PushRawText)?;
             match op {
-                LabelOp::PushLabel(label) => formatter.push_label(label)?,
-                LabelOp::PopLabel => formatter.pop_label()?,
+                FormatOp::PushLabel(label) => formatter.push_label(label)?,
+                FormatOp::PopLabel => formatter.pop_label()?,
+                FormatOp::PushRawText | FormatOp::PopRawText => (),
             }
+            last_op = op;
         }
-        flush_data(formatter, self.data.len())
+        flush_data(formatter, self.data.len(), false)
     }
 }
 
@@ -641,18 +658,33 @@ impl Write for FormatRecorder {
     }
 }
 
+struct RawTextRecorder<'a>(&'a mut FormatRecorder);
+
+impl<'a> Write for RawTextRecorder<'a> {
+    fn write(&mut self, data: &[u8]) -> io::Result<usize> {
+        self.0.push_op(FormatOp::PushRawText);
+        let result = self.0.write(data);
+        self.0.push_op(FormatOp::PopRawText);
+        result
+    }
+
+    fn flush(&mut self) -> io::Result<()> {
+        self.0.flush()
+    }
+}
+
 impl Formatter for FormatRecorder {
     fn raw(&mut self) -> &mut dyn Write {
-        panic!("raw output isn't supported by FormatRecorder")
+        Box::leak(Box::new(RawTextRecorder(self)))
     }
 
     fn push_label(&mut self, label: &str) -> io::Result<()> {
-        self.push_label_op(LabelOp::PushLabel(label.to_owned()));
+        self.push_op(FormatOp::PushLabel(label.to_owned()));
         Ok(())
     }
 
     fn pop_label(&mut self) -> io::Result<()> {
-        self.push_label_op(LabelOp::PopLabel);
+        self.push_op(FormatOp::PopLabel);
         Ok(())
     }
 }

diff --git a/cli/src/template.pest b/cli/src/template.pest
@@ -1,11 +1,11 @@
 // Copyright 2020 The Jujutsu Authors
-// 
+//
 // Licensed under the Apache License, Version 2.0 (the "License");
 // you may not use this file except in compliance with the License.
 // You may obtain a copy of the License at
-// 
+//
 // https://www.apache.org/licenses/LICENSE-2.0
-// 
+//
 // Unless required by applicable law or agreed to in writing, software
 // distributed under the License is distributed on an "AS IS" BASIS,
 // WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
@@ -19,10 +19,11 @@
 
 whitespace = _{ " " | "\t" | "\r" | "\n" | "\x0c" }
 
+ansi_escape = @{ "\\" ~ ("e" | "x1b" | "x1B") }
 string_escape = @{ "\\" ~ ("t" | "r" | "n" | "0" | "\"" | "\\") }
 string_content_char = @{ !("\"" | "\\") ~ ANY }
 string_content = @{ string_content_char+ }
-string_literal = ${ "\"" ~ (string_content | string_escape)* ~ "\"" }
+string_literal = ${ "\"" ~ (string_content | string_escape | ansi_escape)* ~ "\"" }
 
 raw_string_content = @{ (!"'" ~ ANY)* }
 raw_string_literal = ${ "'" ~ raw_string_content ~ "'" }

diff --git a/cli/src/template_builder.rs b/cli/src/template_builder.rs
@@ -39,6 +39,7 @@ use crate::templater::ListTemplate;
 use crate::templater::Literal;
 use crate::templater::PlainTextFormattedProperty;
 use crate::templater::PropertyPlaceholder;
+use crate::templater::RawTextTemplate;
 use crate::templater::ReformatTemplate;
 use crate::templater::SeparateTemplate;
 use crate::templater::SizeHint;
@@ -1116,6 +1117,11 @@ fn builtin_functions<'a, L: TemplateLanguage<'a> + ?Sized>() -> TemplateBuildFun
             content, labels,
         ))))
     });
+    map.insert("raw_text", |language, diagnostics, build_ctx, function| {
+        let [content_node] = function.expect_exact_arguments()?;
+        let content = expect_plain_text_expression(language, diagnostics, build_ctx, content_node)?;
+        Ok(L::wrap_template(Box::new(RawTextTemplate(content))))
+    });
     map.insert("if", |language, diagnostics, build_ctx, function| {
         let ([condition_node, true_node], [false_node]) = function.expect_arguments()?;
         let condition =
@@ -2009,6 +2015,39 @@ mod tests {
         insta::assert_snapshot!(env.render_ok(r#""abcdef".substr(-2, -4)"#), @"");
     }
 
+    #[test]
+    fn test_string_ansi_escape() {
+        let env = TestTemplateEnv::new();
+
+        insta::assert_snapshot!(env.render_ok(r#""\e""#), @"␛");
+        insta::assert_snapshot!(env.render_ok(r#""\x1b""#), @"␛");
+        insta::assert_snapshot!(env.render_ok(r#""\x1B""#), @"␛");
+        insta::assert_snapshot!(
+            env.render_ok(r#""]8;;"
+                ++ "http://example.com"
+                ++ "\e\\"
+                ++ "Example"
+                ++ "\x1b]8;;\x1B\\""#),
+            @r#"␛]8;;http://example.com␛\Example␛]8;;␛\"#);
+    }
+
+    #[test]
+    fn test_raw_text_function_ansi_escape() {
+        // Note: same as test_string_ansi_escape above, wrapped with raw_text.
+        let env = TestTemplateEnv::new();
+
+        insta::assert_snapshot!(env.render_ok(r#"raw_text("\e")"#), @"");
+        insta::assert_snapshot!(env.render_ok(r#"raw_text("\x1b")"#), @"");
+        insta::assert_snapshot!(env.render_ok(r#"raw_text("\x1B")"#), @"");
+        insta::assert_snapshot!(
+            env.render_ok(r#"raw_text("]8;;"
+                ++ "http://example.com"
+                ++ "\e\\"
+                ++ "Example"
+                ++ "\x1b]8;;\x1B\\")"#),
+            @r#"]8;;http://example.com\Example]8;;\"#);
+    }
+
     #[test]
     fn test_signature() {
         let mut env = TestTemplateEnv::new();
@@ -2334,6 +2373,29 @@ mod tests {
             @"[38;5;1mtext[39m");
     }
 
+    #[test]
+    fn test_raw_text_function() {
+        // Note: same as test_label_function above, wrapped with raw_text.
+        let mut env = TestTemplateEnv::new();
+        env.add_keyword("empty", || L::wrap_boolean(Literal(true)));
+        env.add_color("error", crossterm::style::Color::DarkRed);
+        env.add_color("warning", crossterm::style::Color::DarkYellow);
+
+        // Literal
+        insta::assert_snapshot!(
+            env.render_ok(r#"raw_text(label("error", "text"))"#), @"text");
+
+        // Evaluated property
+        insta::assert_snapshot!(
+            env.render_ok(r#"raw_text(label("error".first_line(), "text"))"#),
+            @"text");
+
+        // Template
+        insta::assert_snapshot!(
+            env.render_ok(r#"raw_text(label(if(empty, "error", "warning"), "text"))"#),
+            @"text");
+    }
+
     #[test]
     fn test_coalesce_function() {
         let mut env = TestTemplateEnv::new();

diff --git a/cli/src/template_parser.rs b/cli/src/template_parser.rs
@@ -47,14 +47,16 @@ struct TemplateParser;
 
 const STRING_LITERAL_PARSER: StringLiteralParser<Rule> = StringLiteralParser {
     content_rule: Rule::string_content,
-    escape_rule: Rule::string_escape,
+    string_escape_rule: Rule::string_escape,
+    ansi_escape_rule: Some(Rule::ansi_escape),
 };
 
 impl Rule {
     fn to_symbol(self) -> Option<&'static str> {
         match self {
             Rule::EOI => None,
             Rule::whitespace => None,
+            Rule::ansi_escape => None,
             Rule::string_escape => None,
             Rule::string_content_char => None,
             Rule::string_content => None,

diff --git a/cli/src/templater.rs b/cli/src/templater.rs
@@ -16,6 +16,7 @@ use std::cell::RefCell;
 use std::error;
 use std::fmt;
 use std::io;
+use std::io::Write;
 use std::iter;
 use std::rc::Rc;
 
@@ -184,6 +185,17 @@ where
     }
 }
 
+pub struct RawTextTemplate<T>(pub T);
+
+impl<T: TemplateProperty<Output = String>> Template for RawTextTemplate<T> {
+    fn format(&self, formatter: &mut TemplateFormatter) -> io::Result<()> {
+        match self.0.extract() {
+            Ok(content) => formatter.raw().write_all(content.as_bytes()),
+            Err(err) => formatter.handle_error(err),
+        }
+    }
+}
+
 /// Renders contents in order, and returns the first non-empty output.
 pub struct CoalesceTemplate<T>(pub Vec<T>);
 
@@ -697,6 +709,10 @@ impl<'a> TemplateFormatter<'a> {
         move |formatter| TemplateFormatter::new(formatter, error_handler)
     }
 
+    pub fn raw(&mut self) -> &mut dyn Write {
+        self.formatter.raw()
+    }
+
     pub fn labeled<S: AsRef<str>>(
         &mut self,
         label: S,

diff --git a/lib/src/dsl_util.rs b/lib/src/dsl_util.rs
@@ -393,8 +393,10 @@ where
 pub struct StringLiteralParser<R> {
     /// String content part.
     pub content_rule: R,
-    /// Escape sequence part including backslash character.
-    pub escape_rule: R,
+    /// String escape sequence part including backslash character.
+    pub string_escape_rule: R,
+    /// ANSI escape sequence part.
+    pub ansi_escape_rule: Option<R>,
 }
 
 impl<R: RuleType> StringLiteralParser<R> {
@@ -404,7 +406,7 @@ impl<R: RuleType> StringLiteralParser<R> {
         for part in pairs {
             if part.as_rule() == self.content_rule {
                 result.push_str(part.as_str());
-            } else if part.as_rule() == self.escape_rule {
+            } else if part.as_rule() == self.string_escape_rule {
                 match &part.as_str()[1..] {
                     "\"" => result.push('"'),
                     "\\" => result.push('\\'),
@@ -414,6 +416,8 @@ impl<R: RuleType> StringLiteralParser<R> {
                     "0" => result.push('\0'),
                     char => panic!("invalid escape: \\{char:?}"),
                 }
+            } else if Some(part.as_rule()) == self.ansi_escape_rule {
+                result.push('\x1b');
             } else {
                 panic!("unexpected part of string: {part:?}");
             }

diff --git a/lib/src/fileset_parser.rs b/lib/src/fileset_parser.rs
@@ -37,7 +37,8 @@ struct FilesetParser;
 
 const STRING_LITERAL_PARSER: StringLiteralParser<Rule> = StringLiteralParser {
     content_rule: Rule::string_content,
-    escape_rule: Rule::string_escape,
+    string_escape_rule: Rule::string_escape,
+    ansi_escape_rule: None,
 };
 
 impl Rule {

diff --git a/lib/src/revset_parser.rs b/lib/src/revset_parser.rs
@@ -52,7 +52,8 @@ struct RevsetParser;
 
 const STRING_LITERAL_PARSER: StringLiteralParser<Rule> = StringLiteralParser {
     content_rule: Rule::string_content,
-    escape_rule: Rule::string_escape,
+    string_escape_rule: Rule::string_escape,
+    ansi_escape_rule: None,
 };
 
 impl Rule {