1 files changed, 92 insertions, 18 deletions

diff --git a/testing/tests/gen_ws_tests.py b/testing/tests/gen_ws_tests.py
index d8c809c..c2e24cb 100644
--- a/testing/tests/gen_ws_tests.py
+++ b/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 @@ macro_rules! test_template {
 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)