Specify a trait that handles the output format's escaping

2 files changed, 27 insertions, 40 deletions

diff --git a/askama_derive/src/generator.rs b/askama_derive/src/generator.rs
index 16a08bb..f6c83ec 100644
--- a/askama_derive/src/generator.rs
+++ b/askama_derive/src/generator.rs
@@ -732,14 +732,14 @@ impl<'a> Generator<'a> {
                 }
                 Writable::Expr(s) => {
                     use self::DisplayWrap::*;
-                    use super::input::EscapeMode::*;
                     let mut expr_buf = Buffer::new(0);
                     let wrapped = self.visit_expr(&mut expr_buf, s);
-                    let expression = match (wrapped, &self.input.escaping) {
-                        (Wrapped, &Html) | (Wrapped, &None) | (Unwrapped, &None) => expr_buf.buf,
-                        (Unwrapped, &Html) => {
-                            format!("::askama::MarkupDisplay::from(&{})", expr_buf.buf)
-                        }
+                    let expression = match wrapped {
+                        Wrapped => expr_buf.buf,
+                        Unwrapped => format!(
+                            "::askama::MarkupDisplay::new_unsafe(&{}, {})",
+                            expr_buf.buf, self.input.escaping
+                        ),
                     };
 
                     let id = expr_cache.entry(expression.clone()).or_insert_with(|| {
@@ -876,7 +876,12 @@ impl<'a> Generator<'a> {
             return DisplayWrap::Unwrapped;
         }
 
-        if filters::BUILT_IN_FILTERS.contains(&name) {
+        if name == "escape" || name == "safe" || name == "e" || name == "json" {
+            buf.write(&format!(
+                "::askama::filters::{}({}, &",
+                name, self.input.escaping
+            ));
+        } else if filters::BUILT_IN_FILTERS.contains(&name) {
             buf.write(&format!("::askama::filters::{}(&", name));
         } else {
             buf.write(&format!("filters::{}(&", name));
diff --git a/askama_derive/src/input.rs b/askama_derive/src/input.rs
index c23d30f..b584298 100644
--- a/askama_derive/src/input.rs
+++ b/askama_derive/src/input.rs
@@ -15,7 +15,7 @@ pub struct TemplateInput<'a> {
     pub syntax: &'a Syntax<'a>,
     pub source: Source,
     pub print: Print,
-    pub escaping: EscapeMode,
+    pub escaping: &'a str,
     pub ext: Option<String>,
     pub parent: Option<&'a syn::Type>,
     pub path: PathBuf,
@@ -91,7 +91,7 @@ impl<'a> TemplateInput<'a> {
                         }
                         "escape" => {
                             if let syn::Lit::Str(ref s) = pair.lit {
-                                escaping = Some(s.value().into());
+                                escaping = Some(s.value());
                             } else {
                                 panic!("escape value must be string literal");
                             }
@@ -165,12 +165,24 @@ impl<'a> TemplateInput<'a> {
             },
         );
 
+        let escaping = escaping.unwrap_or_else(|| {
+            path.extension()
+                .map(|s| s.to_str().unwrap())
+                .unwrap_or("none")
+                .to_string()
+        });
+        let escaping = match escaping.as_str() {
+            "html" | "htm" | "xml" => "::askama::Html",
+            "txt" | "none" => "::askama::Text",
+            val => panic!("unknown value '{}' for escape mode", val),
+        };
+
         TemplateInput {
             ast,
             config,
             source,
             print,
-            escaping: escaping.unwrap_or_else(|| EscapeMode::from_path(&path)),
+            escaping,
             ext,
             parent,
             path,
@@ -185,34 +197,6 @@ pub enum Source {
 }
 
 #[derive(PartialEq)]
-pub enum EscapeMode {
-    Html,
-    None,
-}
-
-impl From<String> for EscapeMode {
-    fn from(s: String) -> EscapeMode {
-        use self::EscapeMode::*;
-        match s.as_ref() {
-            "html" => Html,
-            "none" => None,
-            v => panic!("invalid value for escape option: {}", v),
-        }
-    }
-}
-
-impl EscapeMode {
-    fn from_path(path: &PathBuf) -> EscapeMode {
-        let extension = path.extension().map(|s| s.to_str().unwrap()).unwrap_or("");
-        if HTML_EXTENSIONS.contains(&extension) {
-            EscapeMode::Html
-        } else {
-            EscapeMode::None
-        }
-    }
-}
-
-#[derive(PartialEq)]
 pub enum Print {
     All,
     Ast,
@@ -232,5 +216,3 @@ impl From<String> for Print {
         }
     }
 }
-
-const HTML_EXTENSIONS: [&str; 3] = ["html", "htm", "xml"];