Allow configuration of escape modes

3 files changed, 77 insertions, 11 deletions

diff --git a/askama_derive/src/generator.rs b/askama_derive/src/generator.rs
index f6c83ec..955d588 100644
--- a/askama_derive/src/generator.rs
+++ b/askama_derive/src/generator.rs
@@ -738,7 +738,7 @@ impl<'a> Generator<'a> {
                         Wrapped => expr_buf.buf,
                         Unwrapped => format!(
                             "::askama::MarkupDisplay::new_unsafe(&{}, {})",
-                            expr_buf.buf, self.input.escaping
+                            expr_buf.buf, self.input.escaper
                         ),
                     };
 
@@ -879,7 +879,7 @@ impl<'a> Generator<'a> {
         if name == "escape" || name == "safe" || name == "e" || name == "json" {
             buf.write(&format!(
                 "::askama::filters::{}({}, &",
-                name, self.input.escaping
+                name, self.input.escaper
             ));
         } else if filters::BUILT_IN_FILTERS.contains(&name) {
             buf.write(&format!("::askama::filters::{}(&", name));
diff --git a/askama_derive/src/input.rs b/askama_derive/src/input.rs
index b584298..0cc852c 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: &'a str,
+    pub escaper: &'a str,
     pub ext: Option<String>,
     pub parent: Option<&'a syn::Type>,
     pub path: PathBuf,
@@ -165,24 +165,33 @@ impl<'a> TemplateInput<'a> {
             },
         );
 
-        let escaping = escaping.unwrap_or_else(|| {
+        // Match extension against defined output formats
+
+        let extension = escaping.unwrap_or_else(|| {
             path.extension()
                 .map(|s| s.to_str().unwrap())
-                .unwrap_or("none")
+                .unwrap_or("")
                 .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),
-        };
+
+        let mut escaper = None;
+        for (extensions, path) in &config.escapers {
+            if extensions.contains(&extension) {
+                escaper = Some(path);
+                break;
+            }
+        }
+
+        let escaper = escaper.unwrap_or_else(|| {
+            panic!("no escaper defined for extension '{}'", extension);
+        });
 
         TemplateInput {
             ast,
             config,
             source,
             print,
-            escaping,
+            escaper,
             ext,
             parent,
             path,
diff --git a/askama_shared/src/lib.rs b/askama_shared/src/lib.rs
index a0c9321..0daad1b 100644
--- a/askama_shared/src/lib.rs
+++ b/askama_shared/src/lib.rs
@@ -4,6 +4,7 @@ extern crate serde_derive;
 
 use toml;
 
+use std::collections::HashSet;
 use std::env;
 use std::fs;
 use std::path::{Path, PathBuf};
@@ -23,6 +24,7 @@ pub struct Config<'a> {
     pub dirs: Vec<PathBuf>,
     pub syntaxes: BTreeMap<String, Syntax<'a>>,
     pub default_syntax: &'a str,
+    pub escapers: Vec<(HashSet<String>, String)>,
 }
 
 impl<'a> Config<'a> {
@@ -63,10 +65,28 @@ impl<'a> Config<'a> {
             panic!("default syntax \"{}\" not found", default_syntax)
         }
 
+        let mut escapers = Vec::new();
+        if let Some(configured) = raw.escaper {
+            for escaper in configured {
+                escapers.push((
+                    escaper
+                        .extensions
+                        .iter()
+                        .map(|ext| ext.to_string())
+                        .collect(),
+                    escaper.path.to_string(),
+                ));
+            }
+        }
+        for (extensions, path) in DEFAULT_ESCAPERS {
+            escapers.push((str_set(extensions), path.to_string()));
+        }
+
         Config {
             dirs,
             syntaxes,
             default_syntax,
+            escapers,
         }
     }
 
@@ -156,6 +176,7 @@ struct RawConfig<'d> {
     #[serde(borrow)]
     general: Option<General<'d>>,
     syntax: Option<Vec<RawSyntax<'d>>>,
+    escaper: Option<Vec<RawEscaper<'d>>>,
 }
 
 #[derive(Deserialize)]
@@ -176,6 +197,12 @@ struct RawSyntax<'a> {
     comment_end: Option<&'a str>,
 }
 
+#[derive(Deserialize)]
+struct RawEscaper<'a> {
+    path: &'a str,
+    extensions: Vec<&'a str>,
+}
+
 pub fn read_config_file() -> String {
     let root = PathBuf::from(env::var("CARGO_MANIFEST_DIR").unwrap());
     let filename = root.join(CONFIG_FILE_NAME);
@@ -187,8 +214,19 @@ pub fn read_config_file() -> String {
     }
 }
 
+fn str_set<T>(vals: &[T]) -> HashSet<String>
+where
+    T: ToString,
+{
+    vals.iter().map(|s| s.to_string()).collect()
+}
+
 static CONFIG_FILE_NAME: &str = "askama.toml";
 static DEFAULT_SYNTAX_NAME: &str = "default";
+static DEFAULT_ESCAPERS: &[(&[&str], &str)] = &[
+    (&["html", "htm", "xml"], "::askama::Html"),
+    (&["none", "txt", ""], "::askama::Text"),
+];
 
 #[cfg(test)]
 mod tests {
@@ -350,4 +388,23 @@ mod tests {
 
         let _config = Config::new(raw_config);
     }
+
+    #[test]
+    fn escape_modes() {
+        let config = Config::new(
+            r#"
+            [[escaper]]
+            path = "::askama::Js"
+            extensions = ["js"]
+        "#,
+        );
+        assert_eq!(
+            config.escapers,
+            vec![
+                (str_set(&["js"]), "::askama::Js".into()),
+                (str_set(&["html", "htm", "xml"]), "::askama::Html".into()),
+                (str_set(&["none", "txt", ""]), "::askama::Text".into()),
+            ]
+        );
+    }
 }