Fixed whitespace issue when generating if statement (#394)

* Fixed #377

askama_shared/src/generator.rs

@@ -505,8 +505,9 @@ impl<'a, S: std::hash::BuildHasher> Generator<'a, S> {
         let mut has_else = false;
         for (i, &(cws, ref cond, ref nodes)) in conds.iter().enumerate() {
             self.handle_ws(cws);
-            if arm_sizes.is_empty() {
-                flushed += self.write_buf_writable(buf)?;
+            flushed += self.write_buf_writable(buf)?;
+            if i > 0 {
+                self.locals.pop();
             }
 
             let mut arm_size = 0;
@@ -539,14 +540,14 @@ impl<'a, S: std::hash::BuildHasher> Generator<'a, S> {
             self.locals.push();
 
             arm_size += self.handle(ctx, nodes, buf, AstLevel::Nested)?;
-            arm_size += self.write_buf_writable(buf)?;
             arm_sizes.push(arm_size);
-
-            self.locals.pop();
         }
         self.handle_ws(ws);
+        flushed += self.write_buf_writable(buf)?;
         buf.writeln("}")?;
 
+        self.locals.pop();
+
         if !has_else {
             arm_sizes.push(0);
         }

testing/tests/gen_ws_tests.py

@@ -0,0 +1,112 @@
+#!/usr/bin/env python
+# -*- coding: utf-8 -*-
+
+from itertools import product, chain, zip_longest, tee
+
+
+IF, ELSE_IF, ELSE, END_IF = 0, 1, 2, 3
+
+NL = "\\n"
+dash = lambda ws: [" ", "-"][ws]
+
+
+def cond_kind(i, n):
+    i += 1
+    if i == 1:
+        return IF # if
+    elif (i == n) and (i > 1):
+        return ELSE # else
+    elif i > n:
+        return END_IF # endif
+    else:
+        return ELSE_IF # else if
+
+
+# From: https://docs.python.org/3/library/itertools.html#itertools-recipes
+def pairwise(iterable):
+    a, b = tee(iterable)
+    next(b, None)
+    return zip(a, b)
+
+
+def write_cond(conds, active_branch):
+    n = len(conds) - 1
+
+    lits = []
+    for i in range(1, n + 2 + 1):
+        ws1 = "\\n" * i
+        ws2 = "\\r\\n" * i
+        lits.append((ws1, str(i), ws2))
+
+    conds = list(conds)
+    for i, (pws, nws) in enumerate(conds):
+        kind = cond_kind(i, n)
+        b = str(i == active_branch).lower()
+        cond = [f"if {b}", f"else if {b}", "else", "endif"][kind]
+        cond = f"{{%{dash(pws)} {cond} {dash(nws)}%}}"
+        conds[i] = cond
+
+    it = map("".join, lits)
+    it = filter(None, chain.from_iterable(zip_longest(it, conds)))
+    code = "".join(it)
+
+    expected = f"{lits[0][0]}{lits[0][1]}"
+    for i, (cond, (before, after)) in enumerate(zip(conds, pairwise(lits))):
+        kind = cond_kind(i, n)
+        pws = cond.startswith("{%-")
+        nws = cond.endswith("-%}")
+
+        cond = i == active_branch
+        prev_cond = i == (active_branch + 1)
+
+        if prev_cond or (kind == IF):
+            expected += before[2] * (not pws)
+        if cond or (kind == END_IF):
+            expected += after[0] * (not nws)
+            expected += after[1]
+
+    # FIXME: Askama does not include whitespace before eof
+    # expected += lits[-1][2]
+
+    return code, expected
+
+
+if __name__ == "__main__":
+    # The amount of branches to generate
+    n = 2 # branches
+
+
+    with open("ws.rs", "w") as f:
+        f.write("""\
+// This file is auto generated by gen_ws_tests.py
+
+use askama::Template;
+
+macro_rules! test_template {
+    ($source:literal, $rendered:expr) => {{
+        #[derive(Template)]
+        #[template(source = $source, ext = "txt")]
+        struct CondWS;
+
+        assert_eq!(CondWS.render().unwrap(), $rendered);
+    }};
+}
+
+#[rustfmt::skip]
+#[test]
+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 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")