Fixed whitespace issue when generating match (#399)

* Fixed #397

* Updated parser to ignore whitespace between match and when

* Updated test cases

* Updated Python script to generate match ws tests

* Added match ws tests

* Resolved rustfmt lint

askama_shared/src/generator.rs

@@ -444,8 +444,8 @@ impl<'a, S: std::hash::BuildHasher> Generator<'a, S> {
                 Node::Cond(ref conds, ws) => {
                     self.write_cond(ctx, buf, conds, ws)?;
                 }
-                Node::Match(ws1, ref expr, inter, ref arms, ws2) => {
-                    self.write_match(ctx, buf, ws1, expr, inter, arms, ws2)?;
+                Node::Match(ws1, ref expr, ref arms, ws2) => {
+                    self.write_match(ctx, buf, ws1, expr, arms, ws2)?;
                 }
                 Node::Loop(ws1, ref var, ref iter, ref body, ws2) => {
                     self.write_loop(ctx, buf, ws1, var, iter, body, ws2)?;
@@ -561,23 +561,28 @@ impl<'a, S: std::hash::BuildHasher> Generator<'a, S> {
         buf: &mut Buffer,
         ws1: WS,
         expr: &Expr,
-        inter: Option<&'a str>,
         arms: &'a [When],
         ws2: WS,
     ) -> Result<usize, CompileError> {
         self.flush_ws(ws1);
         let flushed = self.write_buf_writable(buf)?;
         let mut arm_sizes = Vec::new();
-        if let Some(inter) = inter {
-            if !inter.is_empty() {
-                self.next_ws = Some(inter);
-            }
-        }
 
         let expr_code = self.visit_expr_root(expr)?;
         buf.writeln(&format!("match &{} {{", expr_code))?;
-        for arm in arms {
+
+        let mut arm_size = 0;
+        for (i, arm) in arms.iter().enumerate() {
             let &(ws, ref variant, ref params, ref body) = arm;
+            self.handle_ws(ws);
+
+            if i > 0 {
+                arm_sizes.push(arm_size + self.write_buf_writable(buf)?);
+
+                buf.writeln("}")?;
+                self.locals.pop();
+            }
+
             self.locals.push();
             match *variant {
                 Some(ref param) => {
@@ -624,15 +629,17 @@ impl<'a, S: std::hash::BuildHasher> Generator<'a, S> {
                 }
             }
             buf.writeln(" => {")?;
-            self.handle_ws(ws);
-            let arm_size = self.handle(ctx, body, buf, AstLevel::Nested)?;
-            arm_sizes.push(arm_size + self.write_buf_writable(buf)?);
-            buf.writeln("}")?;
-            self.locals.pop();
+
+            arm_size = self.handle(ctx, body, buf, AstLevel::Nested)?;
         }
 
-        buf.writeln("}")?;
         self.handle_ws(ws2);
+        arm_sizes.push(arm_size + self.write_buf_writable(buf)?);
+        buf.writeln("}")?;
+        self.locals.pop();
+
+        buf.writeln("}")?;
+
         Ok(flushed + median(&mut arm_sizes))
     }
 

askama_shared/src/heritage.rs

@@ -89,7 +89,7 @@ impl<'a> Context<'a> {
                     Node::Loop(_, _, _, nodes, _) => {
                         nested.push(nodes);
                     }
-                    Node::Match(_, _, _, arms, _) => {
+                    Node::Match(_, _, arms, _) => {
                         for (_, _, _, arm) in arms {
                             nested.push(arm);
                         }

askama_shared/src/parser.rs

@@ -19,7 +19,7 @@ pub enum Node<'a> {
     LetDecl(WS, Target<'a>),
     Let(WS, Target<'a>, Expr<'a>),
     Cond(Vec<(WS, Option<Expr<'a>>, Vec<Node<'a>>)>, WS),
-    Match(WS, Expr<'a>, Option<&'a str>, Vec<When<'a>>, WS),
+    Match(WS, Expr<'a>, Vec<When<'a>>, WS),
     Loop(WS, Target<'a>, Expr<'a>, Vec<Node<'a>>, WS),
     Extends(Expr<'a>),
     BlockDef(WS, &'a str, Vec<Node<'a>>, WS),
@@ -771,8 +771,8 @@ fn block_match<'a>(i: &'a [u8], s: &'a Syntax<'a>) -> IResult<&'a [u8], Node<'a>
         arms.push(arm);
     }
 
-    let inter = match inter {
-        Some(Node::Lit(lws, val, rws)) => {
+    match inter {
+        Some(Node::Lit(_, val, rws)) => {
             assert!(
                 val.is_empty(),
                 "only whitespace allowed between match and first when, found {}",
@@ -783,18 +783,16 @@ fn block_match<'a>(i: &'a [u8], s: &'a Syntax<'a>) -> IResult<&'a [u8], Node<'a>
                 "only whitespace allowed between match and first when, found {}",
                 rws
             );
-            Some(lws)
         }
-        None => None,
+        None => {}
         _ => panic!("only literals allowed between match and first when"),
-    };
+    }
 
     Ok((
         i,
         Node::Match(
             WS(pws1.is_some(), nws1.is_some()),
             expr,
-            inter,
             arms,
             WS(pws2.is_some(), nws2.is_some()),
         ),

testing/tests/gen_ws_tests.py

@@ -4,12 +4,23 @@
 from itertools import product, chain, zip_longest, tee
 
 
+# The amount of branches to generate
+BRANCHES = 2 # branches
+
 IF, ELSE_IF, ELSE, END_IF = 0, 1, 2, 3
 
 NL = "\\n"
 dash = lambda ws: [" ", "-"][ws]
 
 
+def trim(s, ws):
+    if ws[0]:
+        s = s.lstrip()
+    if ws[1]:
+        s = s.rstrip()
+    return s
+
+
 def cond_kind(i, n):
     i += 1
     if i == 1:
@@ -71,17 +82,35 @@ def write_cond(conds, active_branch):
     return code, expected
 
 
-if __name__ == "__main__":
-    # The amount of branches to generate
-    n = 2 # branches
+def write_match(contents, arms, match_ws):
+    before, expr, after = contents
+    code = before
 
+    pws, nws = match_ws[0]
+    code += f"{{%{dash(pws)} match {expr} {dash(nws)}%}}"
+
+    for (arm, expr), (pws, nws) in zip(arms, match_ws[1:-1]):
+        code += f"{{%{dash(pws)} when {arm} {dash(nws)}%}}{expr}"
+
+    pws, nws = match_ws[-1]
+    code += f"{{%{dash(pws)} endmatch {dash(nws)}%}}"
+    code += after
+
+    return code
 
-    with open("ws.rs", "w") as f:
-        f.write("""\
-// This file is auto generated by gen_ws_tests.py
 
-use askama::Template;
+def write_match_result(active_arm, contents, arms, match_ws):
+    before, expr, after = contents
+    expected = ""
 
+    expected += trim(before, (False, match_ws[0][0]))
+    expected += trim(arms[active_arm][1], (match_ws[1:][active_arm][1], match_ws[1:][active_arm+1][0]))
+    expected += trim(after, (match_ws[-1][1], False))
+    return expected
+
+
+def write_cond_tests(f):
+    f.write("""
 macro_rules! test_template {
     ($source:literal, $rendered:expr) => {{
         #[derive(Template)]
@@ -97,16 +126,61 @@ def write_cond(conds, active_branch):
 fn test_cond_ws() {
 """)
 
-        for branches in range(1, n + 1):
-            for x in product([False, True], repeat=(branches+1)*2):
-                # it = iter(x)
-                # conds = list(zip(it, it))
-                conds = list(zip(x[::2], x[1::2]))
+    for branches in range(1, BRANCHES + 1):
+        for x in product([False, True], repeat=(branches+1)*2):
+            # it = iter(x)
+            # conds = list(zip(it, it))
+            conds = list(zip(x[::2], x[1::2]))
+
+            for i in range(branches):
+                code, expected = write_cond(conds, i)
+                f.write(f'    test_template!("{code}", "{expected}");\n')
+
+        if branches != BRANCHES:
+            f.write("\n")
+    f.write("}\n")
 
-                for i in range(branches):
-                    code, expected = write_cond(conds, i)
-                    f.write(f'    test_template!("{code}", "{expected}");\n')
 
-            if branches != n:
-                f.write("\n")
-        f.write("}\n")
+def write_match_tests(f):
+    f.write("""
+#[rustfmt::skip]
+macro_rules! test_match {
+    ($source:literal, $some_rendered:expr, $none_rendered:expr) => {{
+        #[derive(Template)]
+        #[template(source = $source, ext = "txt")]
+        struct MatchWS {
+            item: Option<&'static str>,
+        }
+
+        assert_eq!(MatchWS { item: Some("foo") }.render().unwrap(), $some_rendered);
+        assert_eq!(MatchWS { item: None }.render().unwrap(), $none_rendered);
+    }};
+}
+
+#[rustfmt::skip]
+#[test]
+fn test_match_ws() {
+""")
+
+    contents = "before ", "item", "      after"
+    arms = [("Some with (item)", "  foo   "), ("None", "    bar     ")]
+
+    for x in product([False, True], repeat=len(arms)*2+1):
+        x = [False, False, *x, False]
+        arms_ws = list(zip(x[::2], x[1::2]))
+
+        code = write_match(contents, arms, arms_ws)
+        some_expected = write_match_result(0, contents, arms, arms_ws)
+        none_expected = write_match_result(1, contents, arms, arms_ws)
+
+        f.write(f'    test_match!("{code}", "{some_expected}", "{none_expected}");\n')
+
+    f.write("}\n")
+
+
+if __name__ == "__main__":
+    with open("ws.rs", "w") as f:
+        f.write("// This file is auto generated by gen_ws_tests.py\n\n")
+        f.write("use askama::Template;\n")
+        write_cond_tests(f)
+        write_match_tests(f)

testing/tests/matches.rs

@@ -17,19 +17,19 @@ struct MatchOptRefTemplate<'a> {
 #[test]
 fn test_match_option() {
     let s = MatchOptTemplate { item: Some("foo") };
-    assert_eq!(s.render().unwrap(), "\n\nFound literal foo\n");
+    assert_eq!(s.render().unwrap(), "\nFound literal foo\n");
 
     let s = MatchOptTemplate { item: Some("bar") };
-    assert_eq!(s.render().unwrap(), "\n\nFound bar\n");
+    assert_eq!(s.render().unwrap(), "\nFound bar\n");
 
     let s = MatchOptTemplate { item: None };
-    assert_eq!(s.render().unwrap(), "\n\nNot Found\n");
+    assert_eq!(s.render().unwrap(), "\nNot Found\n");
 }
 
 #[test]
 fn test_match_ref_deref() {
     let s = MatchOptRefTemplate { item: &Some("foo") };
-    assert_eq!(s.render().unwrap(), "\n\nFound literal foo\n");
+    assert_eq!(s.render().unwrap(), "\nFound literal foo\n");
 }
 
 #[derive(Template)]
@@ -41,10 +41,10 @@ struct MatchLitTemplate<'a> {
 #[test]
 fn test_match_literal() {
     let s = MatchLitTemplate { item: "bar" };
-    assert_eq!(s.render().unwrap(), "\n\nFound literal bar\n");
+    assert_eq!(s.render().unwrap(), "\nFound literal bar\n");
 
     let s = MatchLitTemplate { item: "qux" };
-    assert_eq!(s.render().unwrap(), "\n\nElse found qux\n");
+    assert_eq!(s.render().unwrap(), "\nElse found qux\n");
 }
 
 #[derive(Template)]
@@ -56,10 +56,10 @@ struct MatchLitCharTemplate {
 #[test]
 fn test_match_literal_char() {
     let s = MatchLitCharTemplate { item: 'b' };
-    assert_eq!(s.render().unwrap(), "\n\nFound literal b\n");
+    assert_eq!(s.render().unwrap(), "\nFound literal b\n");
 
     let s = MatchLitCharTemplate { item: 'c' };
-    assert_eq!(s.render().unwrap(), "\n\nElse found c\n");
+    assert_eq!(s.render().unwrap(), "\nElse found c\n");
 }
 
 #[derive(Template)]
@@ -71,10 +71,10 @@ struct MatchLitNumTemplate {
 #[test]
 fn test_match_literal_num() {
     let s = MatchLitNumTemplate { item: 42 };
-    assert_eq!(s.render().unwrap(), "\n\nFound answer to everything\n");
+    assert_eq!(s.render().unwrap(), "\nFound answer to everything\n");
 
     let s = MatchLitNumTemplate { item: 23 };
-    assert_eq!(s.render().unwrap(), "\n\nElse found 23\n");
+    assert_eq!(s.render().unwrap(), "\nElse found 23\n");
 }
 
 #[allow(dead_code)]
@@ -99,7 +99,7 @@ fn test_match_custom_enum() {
             b: 255,
         },
     };
-    assert_eq!(s.render().unwrap(), "\n\nColorful: #A000FF\n");
+    assert_eq!(s.render().unwrap(), "\nColorful: #A000FF\n");
 }
 
 #[derive(Template)]

testing/tests/whitespace.rs

@@ -37,5 +37,5 @@ fn test_extra_whitespace() {
     let mut template = AllowWhitespaces::default();
     template.nested_1.nested_2.array = &["a0", "a1", "a2", "a3"];
     template.nested_1.nested_2.hash.insert("key", "value");
-    assert_eq!(template.render().unwrap(), "\n0\n0\n0\n0\n\n\n\n0\n0\n0\n0\n0\n\na0\na1\nvalue\n\n\n\n\n\n[\n  \"a0\",\n  \"a1\",\n  \"a2\",\n  \"a3\"\n]\n[\n  \"a0\",\n  \"a1\",\n  \"a2\",\n  \"a3\"\n][\n  \"a0\",\n  \"a1\",\n  \"a2\",\n  \"a3\"\n]\n[\n  \"a1\"\n][\n  \"a1\"\n]\n[\n  \"a1\",\n  \"a2\"\n][\n  \"a1\",\n  \"a2\"\n]\n[\n  \"a1\"\n][\n  \"a1\"\n]1-1-1\n3333 3\n2222 2\n0000 0\n3333 3\n\ntruefalse\nfalsefalsefalse\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n");
+    assert_eq!(template.render().unwrap(), "\n0\n0\n0\n0\n\n\n\n0\n0\n0\n0\n0\n\na0\na1\nvalue\n\n\n\n\n\n[\n  \"a0\",\n  \"a1\",\n  \"a2\",\n  \"a3\"\n]\n[\n  \"a0\",\n  \"a1\",\n  \"a2\",\n  \"a3\"\n][\n  \"a0\",\n  \"a1\",\n  \"a2\",\n  \"a3\"\n]\n[\n  \"a1\"\n][\n  \"a1\"\n]\n[\n  \"a1\",\n  \"a2\"\n][\n  \"a1\",\n  \"a2\"\n]\n[\n  \"a1\"\n][\n  \"a1\"\n]1-1-1\n3333 3\n2222 2\n0000 0\n3333 3\n\ntruefalse\nfalsefalsefalse\n\n\n\n\n\n\n\n\n\n\n\n\n\n");
 }

testing/tests/ws.rs

@@ -161,3 +161,54 @@ fn test_cond_ws() {
     test_template!("\n1\r\n{%- if true -%}\n\n2\r\n\r\n{%- else -%}\n\n\n3\r\n\r\n\r\n{%- endif -%}\n\n\n\n4\r\n\r\n\r\n\r\n", "\n124");
     test_template!("\n1\r\n{%- if false -%}\n\n2\r\n\r\n{%- else -%}\n\n\n3\r\n\r\n\r\n{%- endif -%}\n\n\n\n4\r\n\r\n\r\n\r\n", "\n134");
 }
+
+#[rustfmt::skip]
+macro_rules! test_match {
+    ($source:literal, $some_rendered:expr, $none_rendered:expr) => {{
+        #[derive(Template)]
+        #[template(source = $source, ext = "txt")]
+        struct MatchWS {
+            item: Option<&'static str>,
+        }
+
+        assert_eq!(MatchWS { item: Some("foo") }.render().unwrap(), $some_rendered);
+        assert_eq!(MatchWS { item: None }.render().unwrap(), $none_rendered);
+    }};
+}
+
+#[rustfmt::skip]
+#[test]
+fn test_match_ws() {
+    test_match!("before {%  match item  %}{%  when Some with (item)  %}  foo   {%  when None  %}    bar     {%  endmatch  %}      after", "before   foo         after", "before     bar           after");
+    test_match!("before {%  match item  %}{%  when Some with (item)  %}  foo   {%  when None  %}    bar     {%- endmatch  %}      after", "before   foo         after", "before     bar      after");
+    test_match!("before {%  match item  %}{%  when Some with (item)  %}  foo   {%  when None -%}    bar     {%  endmatch  %}      after", "before   foo         after", "before bar           after");
+    test_match!("before {%  match item  %}{%  when Some with (item)  %}  foo   {%  when None -%}    bar     {%- endmatch  %}      after", "before   foo         after", "before bar      after");
+    test_match!("before {%  match item  %}{%  when Some with (item)  %}  foo   {%- when None  %}    bar     {%  endmatch  %}      after", "before   foo      after", "before     bar           after");
+    test_match!("before {%  match item  %}{%  when Some with (item)  %}  foo   {%- when None  %}    bar     {%- endmatch  %}      after", "before   foo      after", "before     bar      after");
+    test_match!("before {%  match item  %}{%  when Some with (item)  %}  foo   {%- when None -%}    bar     {%  endmatch  %}      after", "before   foo      after", "before bar           after");
+    test_match!("before {%  match item  %}{%  when Some with (item)  %}  foo   {%- when None -%}    bar     {%- endmatch  %}      after", "before   foo      after", "before bar      after");
+    test_match!("before {%  match item  %}{%  when Some with (item) -%}  foo   {%  when None  %}    bar     {%  endmatch  %}      after", "before foo         after", "before     bar           after");
+    test_match!("before {%  match item  %}{%  when Some with (item) -%}  foo   {%  when None  %}    bar     {%- endmatch  %}      after", "before foo         after", "before     bar      after");
+    test_match!("before {%  match item  %}{%  when Some with (item) -%}  foo   {%  when None -%}    bar     {%  endmatch  %}      after", "before foo         after", "before bar           after");
+    test_match!("before {%  match item  %}{%  when Some with (item) -%}  foo   {%  when None -%}    bar     {%- endmatch  %}      after", "before foo         after", "before bar      after");
+    test_match!("before {%  match item  %}{%  when Some with (item) -%}  foo   {%- when None  %}    bar     {%  endmatch  %}      after", "before foo      after", "before     bar           after");
+    test_match!("before {%  match item  %}{%  when Some with (item) -%}  foo   {%- when None  %}    bar     {%- endmatch  %}      after", "before foo      after", "before     bar      after");
+    test_match!("before {%  match item  %}{%  when Some with (item) -%}  foo   {%- when None -%}    bar     {%  endmatch  %}      after", "before foo      after", "before bar           after");
+    test_match!("before {%  match item  %}{%  when Some with (item) -%}  foo   {%- when None -%}    bar     {%- endmatch  %}      after", "before foo      after", "before bar      after");
+    test_match!("before {%  match item  %}{%- when Some with (item)  %}  foo   {%  when None  %}    bar     {%  endmatch  %}      after", "before   foo         after", "before     bar           after");
+    test_match!("before {%  match item  %}{%- when Some with (item)  %}  foo   {%  when None  %}    bar     {%- endmatch  %}      after", "before   foo         after", "before     bar      after");
+    test_match!("before {%  match item  %}{%- when Some with (item)  %}  foo   {%  when None -%}    bar     {%  endmatch  %}      after", "before   foo         after", "before bar           after");
+    test_match!("before {%  match item  %}{%- when Some with (item)  %}  foo   {%  when None -%}    bar     {%- endmatch  %}      after", "before   foo         after", "before bar      after");
+    test_match!("before {%  match item  %}{%- when Some with (item)  %}  foo   {%- when None  %}    bar     {%  endmatch  %}      after", "before   foo      after", "before     bar           after");
+    test_match!("before {%  match item  %}{%- when Some with (item)  %}  foo   {%- when None  %}    bar     {%- endmatch  %}      after", "before   foo      after", "before     bar      after");
+    test_match!("before {%  match item  %}{%- when Some with (item)  %}  foo   {%- when None -%}    bar     {%  endmatch  %}      after", "before   foo      after", "before bar           after");
+    test_match!("before {%  match item  %}{%- when Some with (item)  %}  foo   {%- when None -%}    bar     {%- endmatch  %}      after", "before   foo      after", "before bar      after");
+    test_match!("before {%  match item  %}{%- when Some with (item) -%}  foo   {%  when None  %}    bar     {%  endmatch  %}      after", "before foo         after", "before     bar           after");
+    test_match!("before {%  match item  %}{%- when Some with (item) -%}  foo   {%  when None  %}    bar     {%- endmatch  %}      after", "before foo         after", "before     bar      after");
+    test_match!("before {%  match item  %}{%- when Some with (item) -%}  foo   {%  when None -%}    bar     {%  endmatch  %}      after", "before foo         after", "before bar           after");
+    test_match!("before {%  match item  %}{%- when Some with (item) -%}  foo   {%  when None -%}    bar     {%- endmatch  %}      after", "before foo         after", "before bar      after");
+    test_match!("before {%  match item  %}{%- when Some with (item) -%}  foo   {%- when None  %}    bar     {%  endmatch  %}      after", "before foo      after", "before     bar           after");
+    test_match!("before {%  match item  %}{%- when Some with (item) -%}  foo   {%- when None  %}    bar     {%- endmatch  %}      after", "before foo      after", "before     bar      after");
+    test_match!("before {%  match item  %}{%- when Some with (item) -%}  foo   {%- when None -%}    bar     {%  endmatch  %}      after", "before foo      after", "before bar           after");
+    test_match!("before {%  match item  %}{%- when Some with (item) -%}  foo   {%- when None -%}    bar     {%- endmatch  %}      after", "before foo      after", "before bar      after");
+}