Implement {{loop.cycle(…)}} similar to Jinja

4 files changed, 95 insertions, 7 deletions

diff --git a/askama_shared/src/generator.rs b/askama_shared/src/generator.rs
index f6d7cca..7500cc7 100644
--- a/askama_shared/src/generator.rs
+++ b/askama_shared/src/generator.rs
@@ -1278,15 +1278,38 @@ impl<'a, S: std::hash::BuildHasher> Generator<'a, S> {
         method: &str,
         args: &[Expr<'_>],
     ) -> Result<DisplayWrap, CompileError> {
-        if let Expr::Var("self") = obj {
-            buf.write("self");
+        if matches!(obj, Expr::Var("loop")) {
+            match method {
+                "cycle" => match args {
+                    [arg] => {
+                        if matches!(arg, Expr::Array(arr) if arr.is_empty()) {
+                            panic!("loop.cycle(…) cannot use an empty array.");
+                        }
+                        buf.write("({");
+                        buf.write("let _cycle = &(");
+                        self.visit_expr(buf, arg)?;
+                        buf.writeln(");")?;
+                        buf.writeln("let _len = _cycle.len();")?;
+                        buf.writeln("if _len == 0 {")?;
+                        buf.writeln("return ::core::result::Result::Err(::askama::Error::Fmt(::core::fmt::Error));")?;
+                        buf.writeln("}")?;
+                        buf.writeln("_cycle[_loop_item.index % _len]")?;
+                        buf.writeln("})")?;
+                    }
+                    _ => return Err("loop.cycle(…) expects exactly one argument".into()),
+                },
+                s => return Err(format!("unknown loop method: {:?}", s).into()),
+            }
         } else {
-            self.visit_expr(buf, obj)?;
+            if let Expr::Var("self") = obj {
+                buf.write("self");
+            } else {
+                self.visit_expr(buf, obj)?;
+            }
+            buf.write(&format!(".{}(", normalize_identifier(method)));
+            self._visit_args(buf, args)?;
+            buf.write(")");
         }
-
-        buf.write(&format!(".{}(", normalize_identifier(method)));
-        self._visit_args(buf, args)?;
-        buf.write(")");
         Ok(DisplayWrap::Unwrapped)
     }
 
diff --git a/testing/tests/loops.rs b/testing/tests/loops.rs
index 49c37e0..f014068 100644
--- a/testing/tests/loops.rs
+++ b/testing/tests/loops.rs
@@ -372,3 +372,48 @@ fn test_loop_break_continue() {
     };
     assert_eq!(t.render().unwrap(), "x1yx2yx3yx11x4yx5y");
 }
+
+#[derive(Template)]
+#[template(
+    source = r#"{% for v in values %}{{loop.cycle(["r", "g", "b"])}}{{v}},{% endfor %}"#,
+    ext = "txt"
+)]
+struct ForCycle<'a> {
+    values: &'a [u8],
+}
+
+#[test]
+fn test_for_cycle() {
+    let t = ForCycle {
+        values: &[1, 2, 3, 4, 5, 6, 7, 8, 9],
+    };
+    assert_eq!(t.render().unwrap(), "r1,g2,b3,r4,g5,b6,r7,g8,b9,");
+}
+
+#[derive(Template)]
+#[template(
+    source = r#"{% for v in values %}{{loop.cycle(cycle)}}{{v}},{% endfor %}"#,
+    ext = "txt"
+)]
+struct ForCycleDynamic<'a> {
+    values: &'a [u8],
+    cycle: &'a [char],
+}
+
+#[test]
+fn test_for_cycle_dynamic() {
+    let t = ForCycleDynamic {
+        values: &[1, 2, 3, 4, 5, 6, 7, 8, 9],
+        cycle: &['a', 'b', 'c', 'd'],
+    };
+    assert_eq!(t.render().unwrap(), "a1,b2,c3,d4,a5,b6,c7,d8,a9,");
+}
+
+#[test]
+fn test_for_cycle_empty() {
+    let t = ForCycleDynamic {
+        values: &[1, 2, 3, 4, 5, 6, 7, 8, 9],
+        cycle: &[],
+    };
+    assert!(t.render().is_err())
+}
diff --git a/testing/tests/ui/loop_cycle_wrong_argument_count.rs b/testing/tests/ui/loop_cycle_wrong_argument_count.rs
new file mode 100644
index 0000000..e87f4f4
--- /dev/null
+++ b/testing/tests/ui/loop_cycle_wrong_argument_count.rs
@@ -0,0 +1,13 @@
+use askama::Template;
+
+#[derive(Template)]
+#[template(
+    source = r#"{% for v in values %}{{ loop.cycle("r", "g", "b") }}{{ v }},{% endfor %}"#,
+    ext = "txt"
+)]
+struct ForCycle<'a> {
+    values: &'a [u8],
+}
+
+fn main() {
+}
diff --git a/testing/tests/ui/loop_cycle_wrong_argument_count.stderr b/testing/tests/ui/loop_cycle_wrong_argument_count.stderr
new file mode 100644
index 0000000..acf3bb1
--- /dev/null
+++ b/testing/tests/ui/loop_cycle_wrong_argument_count.stderr
@@ -0,0 +1,7 @@
+error: loop.cycle(…) expects exactly one argument
+ --> $DIR/loop_cycle_wrong_argument_count.rs:3:10
+  |
+3 | #[derive(Template)]
+  |          ^^^^^^^^
+  |
+  = note: this error originates in the derive macro `Template` (in Nightly builds, run with -Z macro-backtrace for more info)