Implement `for … in … if …`

2 files changed, 33 insertions, 30 deletions

diff --git a/askama_shared/src/generator.rs b/askama_shared/src/generator.rs
index 0c2e854..5c72934 100644
--- a/askama_shared/src/generator.rs
+++ b/askama_shared/src/generator.rs
@@ -605,52 +605,50 @@ impl<'a, S: std::hash::BuildHasher> Generator<'a, S> {
 
         let flushed = self.write_buf_writable(buf)?;
         buf.writeln("{")?;
-        buf.writeln("let mut _did_not_loop = true;")?;
-        buf.write("for (");
-        self.visit_target(buf, true, true, &loop_block.var);
+        buf.writeln("let mut _did_loop = false;")?;
         match loop_block.iter {
-            Expr::Range(_, _, _) => buf.writeln(&format!(
-                ", _loop_item) in ::askama::helpers::TemplateLoop::new({}) {{",
-                expr_code
-            )),
-            Expr::Array(..) => buf.writeln(&format!(
-                ", _loop_item) in ::askama::helpers::TemplateLoop::new({}.iter()) {{",
-                expr_code
-            )),
+            Expr::Range(_, _, _) => buf.writeln(&format!("let _iter = {};", expr_code)),
+            Expr::Array(..) => buf.writeln(&format!("let _iter = {}.iter();", expr_code)),
             // If `iter` is a call then we assume it's something that returns
             // an iterator. If not then the user can explicitly add the needed
             // call without issues.
-            Expr::MethodCall(..) | Expr::PathCall(..) | Expr::Index(..) => buf.writeln(&format!(
-                ", _loop_item) in ::askama::helpers::TemplateLoop::new(({}).into_iter()) {{",
-                expr_code
-            )),
+            Expr::MethodCall(..) | Expr::PathCall(..) | Expr::Index(..) => {
+                buf.writeln(&format!("let _iter = ({}).into_iter();", expr_code))
+            }
             // If accessing `self` then it most likely needs to be
             // borrowed, to prevent an attempt of moving.
-            _ if expr_code.starts_with("self.") => buf.writeln(&format!(
-                ", _loop_item) in ::askama::helpers::TemplateLoop::new(((&{}).into_iter())) {{",
-                expr_code
-            )),
+            _ if expr_code.starts_with("self.") => {
+                buf.writeln(&format!("let _iter = (&{}).into_iter();", expr_code))
+            }
             // If accessing a field then it most likely needs to be
             // borrowed, to prevent an attempt of moving.
-            Expr::Attr(..) => buf.writeln(&format!(
-                ", _loop_item) in ::askama::helpers::TemplateLoop::new(((&{}).into_iter())) {{",
-                expr_code
-            )),
+            Expr::Attr(..) => buf.writeln(&format!("let _iter = (&{}).into_iter();", expr_code)),
             // Otherwise, we borrow `iter` assuming that it implements `IntoIterator`.
-            _ => buf.writeln(&format!(
-                ", _loop_item) in ::askama::helpers::TemplateLoop::new(({}).into_iter()) {{",
-                expr_code
-            )),
+            _ => buf.writeln(&format!("let _iter = ({}).into_iter();", expr_code)),
         }?;
+        if let Some(cond) = &loop_block.cond {
+            self.locals.push();
+            buf.write("let _iter = _iter.filter(|");
+            self.visit_target(buf, true, true, &loop_block.var);
+            buf.write("| -> bool {");
+            self.visit_expr(buf, cond)?;
+            buf.writeln("});")?;
+            self.locals.pop();
+        }
+
+        self.locals.push();
+        buf.write("for (");
+        self.visit_target(buf, true, true, &loop_block.var);
+        buf.writeln(", _loop_item) in ::askama::helpers::TemplateLoop::new(_iter) {")?;
 
-        buf.writeln("_did_not_loop = false;")?;
+        buf.writeln("_did_loop = true;")?;
         let mut size_hint1 = self.handle(ctx, &loop_block.body, buf, AstLevel::Nested)?;
         self.handle_ws(loop_block.ws2);
         size_hint1 += self.write_buf_writable(buf)?;
         self.locals.pop();
         buf.writeln("}")?;
 
-        buf.writeln("if _did_not_loop {")?;
+        buf.writeln("if !_did_loop {")?;
         self.locals.push();
         let mut size_hint2 = self.handle(ctx, &loop_block.else_block, buf, AstLevel::Nested)?;
         self.handle_ws(loop_block.ws3);
diff --git a/askama_shared/src/parser.rs b/askama_shared/src/parser.rs
index 90c2fe3..34a391d 100644
--- a/askama_shared/src/parser.rs
+++ b/askama_shared/src/parser.rs
@@ -38,6 +38,7 @@ pub struct Loop<'a> {
     pub ws1: Ws,
     pub var: Target<'a>,
     pub iter: Expr<'a>,
+    pub cond: Option<Expr<'a>>,
     pub body: Vec<Node<'a>>,
     pub ws2: Ws,
     pub else_block: Vec<Node<'a>>,
@@ -914,6 +915,7 @@ fn parse_loop_content<'a>(i: &'a [u8], s: &State<'_>) -> IResult<&'a [u8], Vec<N
 }
 
 fn block_for<'a>(i: &'a [u8], s: &State<'_>) -> IResult<&'a [u8], Node<'a>> {
+    let if_cond = preceded(ws(tag("if")), cut(ws(expr_any)));
     let else_block = |i| {
         let mut p = preceded(
             ws(tag("else")),
@@ -938,6 +940,7 @@ fn block_for<'a>(i: &'a [u8], s: &State<'_>) -> IResult<&'a [u8], Node<'a>> {
             ws(tag("in")),
             cut(tuple((
                 ws(expr_any),
+                opt(if_cond),
                 opt(char('-')),
                 |i| tag_block_end(i, s),
                 cut(tuple((
@@ -953,7 +956,8 @@ fn block_for<'a>(i: &'a [u8], s: &State<'_>) -> IResult<&'a [u8], Node<'a>> {
             ))),
         ))),
     ));
-    let (i, (pws1, _, (var, _, (iter, nws1, _, (body, (_, pws2, else_block, _, nws2)))))) = p(i)?;
+    let (i, (pws1, _, (var, _, (iter, cond, nws1, _, (body, (_, pws2, else_block, _, nws2)))))) =
+        p(i)?;
     let (nws3, else_block, pws3) = else_block.unwrap_or_default();
     Ok((
         i,
@@ -961,6 +965,7 @@ fn block_for<'a>(i: &'a [u8], s: &State<'_>) -> IResult<&'a [u8], Node<'a>> {
             ws1: Ws(pws1.is_some(), nws1.is_some()),
             var,
             iter,
+            cond,
             body,
             ws2: Ws(pws2.is_some(), nws3),
             else_block,