Repository: djc/askama
Pull request #959 (2024-01-26 00:03:36)

Add support for filter block
Fixes https://github.com/djc/askama/issues/955.

This PR still lacks documentation and some more tests would be nice too (if you have ideas?). But I think it's already good enough to go through a first round of reviews.

GuillaumeGomez: Applied all suggestions (still need a new name for the `buf_writable` method).
GuillaumeGomez: Rebasing is now too difficult so I'll rework git history once PR is approved on your side. ^^'
djc: > Rebasing is now too difficult so I'll rework git history once PR is approved on your side. ^^'

The actual purpose of clean history is making review easier, so that creates a circular dependency. 😛 
GuillaumeGomez: Oh boy. XD

Ok, let me rework the history now then. ^^'
GuillaumeGomez: New clean history and all comments fixed/implemented.
GuillaumeGomez: Ping. :)
GuillaumeGomez: Any news about this? :)
GuillaumeGomez: Ping @djc. :)
GuillaumeGomez: Updated without moving the `generate_expr_with_id` for now.
GuillaumeGomez: Move the change in a commit on its own (and as first commit).
GuillaumeGomez: Ah the push was finally taken into account by github. \o/

Now just have to wait for CI. ^^'
GuillaumeGomez: :tada: 

Patch:
diff --git a/askama_derive/src/generator.rs b/askama_derive/src/generator.rs
index 6151038..ddda791 100644
--- a/askama_derive/src/generator.rs
+++ b/askama_derive/src/generator.rs
@@ -1,3 +1,4 @@
+use std::borrow::Cow;
 use std::collections::hash_map::{Entry, HashMap};
 use std::path::Path;
 use std::{cmp, hash, mem, str};
@@ -8,7 +9,8 @@ use crate::input::{Source, TemplateInput};
 use crate::{CompileError, CRATE};
 
 use parser::node::{
-    Call, Comment, CondTest, If, Include, Let, Lit, Loop, Match, Target, Whitespace, Ws,
+    Call, Comment, CondTest, FilterBlock, If, Include, Let, Lit, Loop, Match, Target, Whitespace,
+    Ws,
 };
 use parser::{Expr, Node};
 use quote::quote;
@@ -21,7 +23,7 @@ pub(crate) struct Generator<'a> {
     // The heritage contains references to blocks and their ancestry
     heritage: Option<&'a Heritage<'a>>,
     // Variables accessible directly from the current scope (not redirected to context)
-    locals: MapChain<'a, &'a str, LocalMeta>,
+    locals: MapChain<'a, Cow<'a, str>, LocalMeta>,
     // Suffix whitespace from the previous literal. Will be flushed to the
     // output buffer unless suppressed by whitespace suppression on the next
     // non-literal.
@@ -42,7 +44,7 @@ impl<'a> Generator<'a> {
         input: &'n TemplateInput<'_>,
         contexts: &'n HashMap<&'n Path, Context<'n>>,
         heritage: Option<&'n Heritage<'_>>,
-        locals: MapChain<'n, &'n str, LocalMeta>,
+        locals: MapChain<'n, Cow<'n, str>, LocalMeta>,
     ) -> Generator<'n> {
         Generator {
             input,
@@ -311,6 +313,9 @@ impl<'a> Generator<'a> {
                 Node::Call(ref call) => {
                     size_hint += self.write_call(ctx, buf, call)?;
                 }
+                Node::FilterBlock(ref filter) => {
+                    size_hint += self.write_filter_block(ctx, buf, filter)?;
+                }
                 Node::Macro(ref m) => {
                     if level != AstLevel::Top {
                         return Err("macro blocks only allowed at the top level".into());
@@ -631,7 +636,7 @@ impl<'a> Generator<'a> {
                         "no argument named `{arg_name}` in macro {name:?}"
                     )));
                 }
-                named_arguments.insert(arg_name, arg);
+                named_arguments.insert(Cow::Borrowed(arg_name), arg);
             }
         }
 
@@ -644,7 +649,7 @@ impl<'a> Generator<'a> {
         //   anything since named arguments are always last).
         let mut allow_positional = true;
         for (index, arg) in def.args.iter().enumerate() {
-            let expr = match named_arguments.get(&arg) {
+            let expr = match named_arguments.get(&Cow::Borrowed(arg)) {
                 Some(expr) => {
                     allow_positional = false;
                     expr
@@ -667,14 +672,16 @@ impl<'a> Generator<'a> {
                 // to avoid moving non-copyable values.
                 &Expr::Var(name) if name != "self" => {
                     let var = self.locals.resolve_or_self(name);
-                    self.locals.insert(arg, LocalMeta::with_ref(var));
+                    self.locals
+                        .insert(Cow::Borrowed(arg), LocalMeta::with_ref(var));
                 }
                 Expr::Attr(obj, attr) => {
                     let mut attr_buf = Buffer::new(0);
                     self.visit_attr(&mut attr_buf, obj, attr)?;
 
                     let var = self.locals.resolve(&attr_buf.buf).unwrap_or(attr_buf.buf);
-                    self.locals.insert(arg, LocalMeta::with_ref(var));
+                    self.locals
+                        .insert(Cow::Borrowed(arg), LocalMeta::with_ref(var));
                 }
                 // Everything else still needs to become variables,
                 // to avoid having the same logic be executed
@@ -692,7 +699,7 @@ impl<'a> Generator<'a> {
                     values.write("(");
                     values.write(&self.visit_expr_root(expr)?);
                     values.write(")");
-                    self.locals.insert_with_default(arg);
+                    self.locals.insert_with_default(Cow::Borrowed(arg));
                 }
             }
         }
@@ -712,6 +719,69 @@ impl<'a> Generator<'a> {
         Ok(size_hint)
     }
 
+    fn write_filter_block(
+        &mut self,
+        ctx: &'a Context<'_>,
+        buf: &mut Buffer,
+        filter: &'a FilterBlock<'_>,
+    ) -> Result<usize, CompileError> {
+        self.flush_ws(filter.ws1);
+        let mut var_name = String::new();
+        for id in 0.. {
+            var_name = format!("__filter_block{id}");
+            if self.locals.get(&Cow::Borrowed(var_name.as_str())).is_none() {
+                // No variable with this name exists, we're in the clear!
+                break;
+            }
+        }
+        let current_buf = mem::take(&mut self.buf_writable);
+
+        self.prepare_ws(filter.ws1);
+        let mut size_hint = self.handle(ctx, &filter.nodes, buf, AstLevel::Nested)?;
+        self.flush_ws(filter.ws2);
+
+        let WriteParts {
+            size_hint: write_size_hint,
+            buffers,
+        } = self.prepare_format(buf.indent + 1)?;
+        size_hint += match buffers {
+            None => return Ok(0),
+            Some(WritePartsBuffers { format, expr: None }) => {
+                buf.writeln(&format!("let {var_name} = {:#?};", &format.buf))?;
+                write_size_hint
+            }
+            Some(WritePartsBuffers {
+                format,
+                expr: Some(expr),
+            }) => {
+                buf.writeln(&format!(
+                    "let {var_name} = format!({:#?}, {});",
+                    &format.buf,
+                    expr.buf.trim(),
+                ))?;
+                write_size_hint
+            }
+        };
+
+        mem::drop(mem::replace(&mut self.buf_writable, current_buf));
+
+        let mut filter_buf = Buffer::new(buf.indent);
+        let mut args = filter.args.clone();
+        args.insert(0, Expr::Generated(var_name.clone()));
+
+        let wrap = self.visit_filter(&mut filter_buf, filter.filter_name, &args)?;
+
+        self.buf_writable
+            .push(Writable::Generated(filter_buf.buf, wrap));
+        self.prepare_ws(filter.ws2);
+
+        // We don't forget to add the created variable into the list of variables in the scope.
+        self.locals
+            .insert(Cow::Owned(var_name), LocalMeta::initialized());
+
+        Ok(size_hint)
+    }
+
     fn handle_include(
         &mut self,
         ctx: &'a Context<'_>,
@@ -775,7 +845,7 @@ impl<'a> Generator<'a> {
         match var {
             Target::Name(name) => {
                 let name = normalize_identifier(name);
-                match self.locals.get(&name) {
+                match self.locals.get(&Cow::Borrowed(name)) {
                     // declares a new variable
                     None => Ok(false),
                     // an initialized variable gets shadowed
@@ -827,7 +897,7 @@ impl<'a> Generator<'a> {
         }
         if shadowed
             || !matches!(l.var, Target::Name(_))
-            || matches!(&l.var, Target::Name(name) if self.locals.get(name).is_none())
+            || matches!(&l.var, Target::Name(name) if self.locals.get(&Cow::Borrowed(name)).is_none())
         {
             buf.write("let ");
         }
@@ -903,8 +973,38 @@ impl<'a> Generator<'a> {
 
     // Write expression buffer and empty
     fn write_buf_writable(&mut self, buf: &mut Buffer) -> Result<usize, CompileError> {
+        let WriteParts { size_hint, buffers } = self.prepare_format(buf.indent)?;
+        match buffers {
+            None => Ok(size_hint),
+            Some(WritePartsBuffers { format, expr: None }) => {
+                buf.writeln(&format!("writer.write_str({:#?})?;", &format.buf))?;
+                Ok(size_hint)
+            }
+            Some(WritePartsBuffers {
+                format,
+                expr: Some(expr),
+            }) => {
+                buf.writeln("::std::write!(")?;
+                buf.indent();
+                buf.writeln("writer,")?;
+                buf.writeln(&format!("{:#?},", &format.buf))?;
+                buf.writeln(expr.buf.trim())?;
+                buf.dedent()?;
+                buf.writeln(")?;")?;
+                Ok(size_hint)
+            }
+        }
+    }
+
+    /// This is the common code to generate an expression. It is used for filter blocks and for
+    /// expressions more generally. It stores the size it represents and the buffers. Take a look
+    /// at `WriteParts` for more details.
+    fn prepare_format(&mut self, indent: u8) -> Result<WriteParts, CompileError> {
         if self.buf_writable.is_empty() {
-            return Ok(0);
+            return Ok(WriteParts {
+                size_hint: 0,
+                buffers: None,
+            });
         }
 
         if self
@@ -918,14 +1018,21 @@ impl<'a> Generator<'a> {
                     buf_lit.write(s);
                 };
             }
-            buf.writeln(&format!("writer.write_str({:#?})?;", &buf_lit.buf))?;
-            return Ok(buf_lit.buf.len());
+            return Ok(WriteParts {
+                size_hint: buf_lit.buf.len(),
+                buffers: Some(WritePartsBuffers {
+                    format: buf_lit,
+                    expr: None,
+                }),
+            });
         }
 
+        let mut expr_cache = HashMap::with_capacity(self.buf_writable.len());
+
         let mut size_hint = 0;
         let mut buf_format = Buffer::new(0);
-        let mut buf_expr = Buffer::new(buf.indent + 1);
-        let mut expr_cache = HashMap::with_capacity(self.buf_writable.len());
+        let mut buf_expr = Buffer::new(indent + 1);
+
         for s in mem::take(&mut self.buf_writable) {
             match s {
                 Writable::Lit(s) => {
@@ -933,50 +1040,76 @@ impl<'a> Generator<'a> {
                     size_hint += s.len();
                 }
                 Writable::Expr(s) => {
-                    use self::DisplayWrap::*;
                     let mut expr_buf = Buffer::new(0);
                     let wrapped = self.visit_expr(&mut expr_buf, s)?;
-                    let expression = match wrapped {
-                        Wrapped => expr_buf.buf,
-                        Unwrapped => format!(
-                            "{CRATE}::MarkupDisplay::new_unsafe(&({}), {})",
-                            expr_buf.buf, self.input.escaper
-                        ),
-                    };
-
-                    let id = match expr_cache.entry(expression) {
-                        Entry::Occupied(e) if is_cacheable(s) => *e.get(),
-                        e => {
-                            let id = self.named;
-                            self.named += 1;
-
-                            buf_expr.write(&format!("expr{id} = "));
-                            buf_expr.write("&");
-                            buf_expr.write(e.key());
-                            buf_expr.writeln(",")?;
-
-                            if let Entry::Vacant(e) = e {
-                                e.insert(id);
-                            }
-
-                            id
-                        }
-                    };
-
-                    buf_format.write(&format!("{{expr{id}}}"));
-                    size_hint += 3;
+                    let cacheable = is_cacheable(s);
+                    size_hint += self.named_expression(
+                        &mut buf_expr,
+                        &mut buf_format,
+                        expr_buf.buf,
+                        wrapped,
+                        cacheable,
+                        &mut expr_cache,
+                    )?;
+                }
+                Writable::Generated(s, wrapped) => {
+                    size_hint += self.named_expression(
+                        &mut buf_expr,
+                        &mut buf_format,
+                        s,
+                        wrapped,
+                        false,
+                        &mut expr_cache,
+                    )?;
                 }
             }
         }
+        Ok(WriteParts {
+            size_hint,
+            buffers: Some(WritePartsBuffers {
+                format: buf_format,
+                expr: Some(buf_expr),
+            }),
+        })
+    }
 
-        buf.writeln("::std::write!(")?;
-        buf.indent();
-        buf.writeln("writer,")?;
-        buf.writeln(&format!("{:#?},", &buf_format.buf))?;
-        buf.writeln(buf_expr.buf.trim())?;
-        buf.dedent()?;
-        buf.writeln(")?;")?;
-        Ok(size_hint)
+    fn named_expression(
+        &mut self,
+        buf_expr: &mut Buffer,
+        buf_format: &mut Buffer,
+        expr: String,
+        wrapped: DisplayWrap,
+        cacheable: bool,
+        expr_cache: &mut HashMap<String, usize>,
+    ) -> Result<usize, CompileError> {
+        let expression = match wrapped {
+            DisplayWrap::Wrapped => expr,
+            DisplayWrap::Unwrapped => format!(
+                "{CRATE}::MarkupDisplay::new_unsafe(&({}), {})",
+                expr, self.input.escaper
+            ),
+        };
+        let id = match expr_cache.entry(expression) {
+            Entry::Occupied(e) if cacheable => *e.get(),
+            entry => {
+                let id = self.named;
+                self.named += 1;
+
+                buf_expr.write(&format!("expr{id} = "));
+                buf_expr.write("&");
+                buf_expr.write(entry.key());
+                buf_expr.writeln(",")?;
+
+                if let Entry::Vacant(e) = entry {
+                    e.insert(id);
+                }
+
+                id
+            }
+        };
+
+        buf_format.write(&format!("{{expr{id}}}"));
+        Ok(3)
     }
 
     fn visit_lit(&mut self, lit: &'a Lit<'_>) {
@@ -1049,6 +1182,7 @@ impl<'a> Generator<'a> {
             Expr::Try(ref expr) => self.visit_try(buf, expr.as_ref())?,
             Expr::Tuple(ref exprs) => self.visit_tuple(buf, exprs)?,
             Expr::NamedArgument(_, ref expr) => self.visit_named_argument(buf, expr)?,
+            Expr::Generated(ref s) => self.visit_generated(buf, s),
         })
     }
 
@@ -1517,6 +1651,11 @@ impl<'a> Generator<'a> {
         DisplayWrap::Unwrapped
     }
 
+    fn visit_generated(&mut self, buf: &mut Buffer, s: &str) -> DisplayWrap {
+        buf.write(s);
+        DisplayWrap::Unwrapped
+    }
+
     fn visit_bool_lit(&mut self, buf: &mut Buffer, s: &str) -> DisplayWrap {
         buf.write(s);
         DisplayWrap::Unwrapped
@@ -1551,8 +1690,10 @@ impl<'a> Generator<'a> {
             Target::Name(name) => {
                 let name = normalize_identifier(name);
                 match initialized {
-                    true => self.locals.insert(name, LocalMeta::initialized()),
-                    false => self.locals.insert_with_default(name),
+                    true => self
+                        .locals
+                        .insert(Cow::Borrowed(name), LocalMeta::initialized()),
+                    false => self.locals.insert_with_default(Cow::Borrowed(name)),
                 }
                 buf.write(name);
             }
@@ -1815,10 +1956,10 @@ where
     }
 }
 
-impl MapChain<'_, &str, LocalMeta> {
+impl MapChain<'_, Cow<'_, str>, LocalMeta> {
     fn resolve(&self, name: &str) -> Option<String> {
         let name = normalize_identifier(name);
-        self.get(&name).map(|meta| match &meta.refs {
+        self.get(&Cow::Borrowed(name)).map(|meta| match &meta.refs {
             Some(expr) => expr.clone(),
             None => name.to_string(),
         })
@@ -1906,6 +2047,7 @@ pub(crate) fn is_cacheable(expr: &Expr<'_>) -> bool {
         Expr::Call(_, _) => false,
         Expr::RustMacro(_, _) => false,
         Expr::Try(_) => false,
+        Expr::Generated(_) => true,
     }
 }
 
@@ -1925,7 +2067,7 @@ enum AstLevel {
     Nested,
 }
 
-#[derive(Clone, Copy)]
+#[derive(Clone, Copy, Debug)]
 enum DisplayWrap {
     Wrapped,
     Unwrapped,
@@ -1935,6 +2077,28 @@ enum DisplayWrap {
 enum Writable<'a> {
     Lit(&'a str),
     Expr(&'a Expr<'a>),
+    Generated(String, DisplayWrap),
+}
+
+struct WriteParts {
+    size_hint: usize,
+    buffers: Option<WritePartsBuffers>,
+}
+
+/// If "expr" is `None`, it means we can generate code like this:
+///
+/// ```ignore
+/// let var = format;
+/// ```
+///
+/// Otherwise we need to format "expr" using "format":
+///
+/// ```ignore
+/// let var = format!(format, expr);
+/// ```
+struct WritePartsBuffers {
+    format: Buffer,
+    expr: Option<Buffer>,
 }
 
 // Identifiers to be replaced with raw identifiers, so as to avoid
diff --git a/askama_derive/src/input.rs b/askama_derive/src/input.rs
index 215d060..7d2d74b 100644
--- a/askama_derive/src/input.rs
+++ b/askama_derive/src/input.rs
@@ -154,6 +154,9 @@ impl TemplateInput<'_> {
                             let import = self.config.find_template(import.path, Some(&path))?;
                             add_to_check(import)?;
                         }
+                        FilterBlock(f) => {
+                            nested.push(&f.nodes);
+                        }
                         Include(include) => {
                             let include = self.config.find_template(include.path, Some(&path))?;
                             add_to_check(include)?;
diff --git a/askama_parser/src/expr.rs b/askama_parser/src/expr.rs
index 4aa778a..be5adfd 100644
--- a/askama_parser/src/expr.rs
+++ b/askama_parser/src/expr.rs
@@ -72,6 +72,8 @@ pub enum Expr<'a> {
     Call(Box<Expr<'a>>, Vec<Expr<'a>>),
     RustMacro(Vec<&'a str>, &'a str),
     Try(Box<Expr<'a>>),
+    /// This variant should never be used directly. It is created when generating filter blocks.
+    Generated(String),
 }
 
 impl<'a> Expr<'a> {
diff --git a/askama_parser/src/node.rs b/askama_parser/src/node.rs
index 16c5382..a50c2e7 100644
--- a/askama_parser/src/node.rs
+++ b/askama_parser/src/node.rs
@@ -37,6 +37,7 @@ pub enum Node<'a> {
     Raw(Raw<'a>),
     Break(Ws),
     Continue(Ws),
+    FilterBlock(FilterBlock<'a>),
 }
 
 impl<'a> Node<'a> {
@@ -66,6 +67,7 @@ impl<'a> Node<'a> {
                 map(|i| Raw::parse(i, s), Self::Raw),
                 |i| Self::r#break(i, s),
                 |i| Self::r#continue(i, s),
+                map(|i| FilterBlock::parse(i, s), Self::FilterBlock),
             )),
             cut(|i| s.tag_block_end(i)),
         );
@@ -558,6 +560,53 @@ impl<'a> Macro<'a> {
     }
 }
 
+#[derive(Debug, PartialEq)]
+pub struct FilterBlock<'a> {
+    pub ws1: Ws,
+    pub filter_name: &'a str,
+    pub args: Vec<Expr<'a>>,
+    pub nodes: Vec<Node<'a>>,
+    pub ws2: Ws,
+}
+
+impl<'a> FilterBlock<'a> {
+    fn parse(i: &'a str, s: &State<'_>) -> ParseResult<'a, Self> {
+        let mut start = tuple((
+            opt(Whitespace::parse),
+            ws(keyword("filter")),
+            cut(tuple((
+                ws(identifier),
+                opt(|i| Expr::arguments(i, s.level.get(), false)),
+                opt(Whitespace::parse),
+                |i| s.tag_block_end(i),
+            ))),
+        ));
+        let (i, (pws1, _, (filter_name, params, nws1, _))) = start(i)?;
+
+        let mut end = cut(tuple((
+            |i| Node::many(i, s),
+            cut(tuple((
+                |i| s.tag_block_start(i),
+                opt(Whitespace::parse),
+                ws(keyword("endfilter")),
+                opt(Whitespace::parse),
+            ))),
+        )));
+        let (i, (nodes, (_, pws2, _, nws2))) = end(i)?;
+
+        Ok((
+            i,
+            Self {
+                ws1: Ws(pws1, nws1),
+                filter_name,
+                args: params.unwrap_or_default(),
+                nodes,
+                ws2: Ws(pws2, nws2),
+            },
+        ))
+    }
+}
+
 #[derive(Debug, PartialEq)]
 pub struct Import<'a> {
     pub ws: Ws,


Test patch:
diff --git a/testing/tests/filter_block.rs b/testing/tests/filter_block.rs
new file mode 100644
index 0000000..029ce13
--- /dev/null
+++ b/testing/tests/filter_block.rs
@@ -0,0 +1,130 @@
+use askama::Template;
+
+#[derive(Template)]
+#[template(
+    source = r#"{% filter lower %}
+    {{ t }} / HELLO / {{ u }}
+{% endfilter %}
+
+{{ u|lower }}
+"#,
+    ext = "html"
+)]
+struct A<T, U = u8>
+where
+    T: std::fmt::Display,
+    U: std::fmt::Display,
+{
+    t: T,
+    u: U,
+}
+
+#[test]
+fn filter_block_basic() {
+    let template = A { t: "a", u: "B" };
+    assert_eq!(template.render().unwrap(), "\n    a / hello / b\n\n\nb\n")
+}
+
+// This test ensures that we don't have variable shadowing when we have more than one
+// filter block at the same location.
+#[derive(Template)]
+#[template(
+    source = r#"{% filter lower %}
+    {{ t }} / HELLO / {{ u }}
+{% endfilter %}
+
+{% filter upper %}
+{{ u }} + TaDaM + {{ t }}
+{% endfilter %}
+
+{% filter lower %}
+    {{ t }} - CHECK - {{ t }}
+{% endfilter %}
+
+{{ u|upper }}"#,
+    ext = "html"
+)]
+struct B<T, U = u8>
+where
+    T: std::fmt::Display,
+    U: std::fmt::Display,
+{
+    t: T,
+    u: U,
+}
+
+#[test]
+fn filter_block_shadowing() {
+    let template = B { t: "a", u: "B" };
+    assert_eq!(
+        template.render().unwrap(),
+        r#"
+    a / hello / b
+
+
+
+B + TADAM + A
+
+
+
+    a - check - a
+
+
+B"#
+    );
+}
+
+// This test ensures that whitespace control is correctly handled.
+#[derive(Template)]
+#[template(
+    source = r#"{% filter lower -%}
+    {{ t }} / HELLO / {{ u }}
+{% endfilter %}
+
+{%- filter upper -%}
+{{ u }} + TaDaM + {{ t }}
+{%- endfilter -%}
+
+++b"#,
+    ext = "html"
+)]
+struct C<T, U = u8>
+where
+    T: std::fmt::Display,
+    U: std::fmt::Display,
+{
+    t: T,
+    u: U,
+}
+
+#[test]
+fn filter_block_whitespace_control() {
+    let template = C { t: "a", u: "B" };
+    assert_eq!(
+        template.render().unwrap(),
+        r#"a / hello / b
+B + TADAM + A++b"#
+    );
+}
+
+// This test ensures that HTML escape is correctly handled.
+#[derive(Template)]
+#[template(source = r#"{% filter lower %}<block>{% endfilter %}"#, ext = "html")]
+struct D;
+
+#[test]
+fn filter_block_html_escape() {
+    let template = D;
+    assert_eq!(template.render().unwrap(), r#"&lt;block&gt;"#);
+}
+
+// This test ensures that it is not escaped if it is not HTML.
+#[derive(Template)]
+#[template(source = r#"{% filter lower %}<block>{% endfilter %}"#, ext = "txt")]
+struct E;
+
+#[test]
+fn filter_block_not_html_escape() {
+    let template = E;
+    assert_eq!(template.render().unwrap(), r#"<block>"#);
+}