Add support for loop.last

7 files changed, 72 insertions, 7 deletions

diff --git a/askama/src/lib.rs b/askama/src/lib.rs
index d1bb639..725c3de 100644
--- a/askama/src/lib.rs
+++ b/askama/src/lib.rs
@@ -433,6 +433,7 @@ pub trait Template {
 }
 
 pub use crate::shared::filters;
+pub use crate::shared::helpers;
 pub use crate::shared::{read_config_file, Error, MarkupDisplay, Result};
 pub use askama_derive::*;
 
diff --git a/askama_derive/src/generator.rs b/askama_derive/src/generator.rs
index 67ec2c2..6c02a26 100644
--- a/askama_derive/src/generator.rs
+++ b/askama_derive/src/generator.rs
@@ -467,11 +467,17 @@ impl<'a> Generator<'a> {
         let expr_code = self.visit_expr_root(iter);
 
         self.write_buf_writable(buf);
-        buf.write("for (_loop_index, ");
+        buf.write("for (_loop_index, _loop_last, ");
         self.visit_target(buf, var);
         match iter {
-            Expr::Range(_, _, _) => buf.writeln(&format!(") in ({}).enumerate() {{", expr_code)),
-            _ => buf.writeln(&format!(") in (&{}).into_iter().enumerate() {{", expr_code)),
+            Expr::Range(_, _, _) => buf.writeln(&format!(
+                ") in ::askama::helpers::enumerate({}) {{",
+                expr_code
+            )),
+            _ => buf.writeln(&format!(
+                ") in ::askama::helpers::enumerate((&{}).into_iter()) {{",
+                expr_code
+            )),
         };
 
         self.handle(ctx, body, buf, AstLevel::Nested);
@@ -915,6 +921,9 @@ impl<'a> Generator<'a> {
                 } else if attr == "first" {
                     buf.write("(_loop_index == 0)");
                     return DisplayWrap::Unwrapped;
+                } else if attr == "last" {
+                    buf.write("_loop_last");
+                    return DisplayWrap::Unwrapped;
                 } else {
                     panic!("unknown loop variable");
                 }
diff --git a/askama_shared/src/helpers/mod.rs b/askama_shared/src/helpers/mod.rs
new file mode 100644
index 0000000..1a3fd7b
--- /dev/null
+++ b/askama_shared/src/helpers/mod.rs
@@ -0,0 +1,48 @@
+use std::iter::Peekable;
+
+pub struct Enumerate<I>
+where
+    I: Iterator,
+{
+    iter: Peekable<I>,
+    count: usize,
+}
+
+impl<I> Iterator for Enumerate<I>
+where
+    I: Iterator,
+{
+    type Item = (usize, bool, <I as Iterator>::Item);
+
+    #[inline]
+    fn next(&mut self) -> Option<(usize, bool, <I as Iterator>::Item)> {
+        self.iter.next().map(|a| {
+            let last = self.iter.peek().is_none();
+            let ret = (self.count, last, a);
+            // Possible undefined overflow.
+            self.count += 1;
+            ret
+        })
+    }
+
+    #[inline]
+    fn size_hint(&self) -> (usize, Option<usize>) {
+        self.iter.size_hint()
+    }
+
+    #[inline]
+    fn count(self) -> usize {
+        self.iter.count()
+    }
+}
+
+#[inline]
+pub fn enumerate<I>(iter: I) -> Enumerate<I>
+where
+    I: Iterator,
+{
+    Enumerate {
+        iter: iter.peekable(),
+        count: 0,
+    }
+}
diff --git a/askama_shared/src/lib.rs b/askama_shared/src/lib.rs
index 850cb8b..a0c9321 100644
--- a/askama_shared/src/lib.rs
+++ b/askama_shared/src/lib.rs
@@ -16,6 +16,7 @@ pub use askama_escape::MarkupDisplay;
 use std::collections::BTreeMap;
 
 pub mod filters;
+pub mod helpers;
 
 #[derive(Debug)]
 pub struct Config<'a> {
diff --git a/testing/templates/for-range.html b/testing/templates/for-range.html
index 9b109bb..3688478 100644
--- a/testing/templates/for-range.html
+++ b/testing/templates/for-range.html
@@ -1,5 +1,5 @@
 {% for s in 0..2 -%}
-    foo
+    foo{% if loop.first %} (first){% endif %}{% if loop.last %} (last){% endif %}
 {% endfor -%}
 
 {% for s in init..1 -%}
diff --git a/testing/templates/precedence-for.html b/testing/templates/precedence-for.html
index c8d2cdd..57c900b 100644
--- a/testing/templates/precedence-for.html
+++ b/testing/templates/precedence-for.html
@@ -1,3 +1,3 @@
 {% for s in strings %}
-    {{- loop.index0 }}. {{ s }}{{ 2 * loop.index }}{% if !loop.first %}{% else %} (first){% endif %}
+    {{- loop.index0 }}. {{ s }}{{ 2 * loop.index }}{% if !loop.first %}{% else %} (first){% endif %}{% if loop.last %} (last){% endif %}
 {% endfor %}
diff --git a/testing/tests/loops.rs b/testing/tests/loops.rs
index b5fac7c..3ce6185 100644
--- a/testing/tests/loops.rs
+++ b/testing/tests/loops.rs
@@ -46,7 +46,10 @@ fn test_precedence_for() {
     let s = PrecedenceTemplate {
         strings: vec!["A", "alfa", "1"],
     };
-    assert_eq!(s.render().unwrap(), "0. A2 (first)\n1. alfa4\n2. 16\n");
+    assert_eq!(
+        s.render().unwrap(),
+        "0. A2 (first)\n1. alfa4\n2. 16 (last)\n"
+    );
 }
 
 #[derive(Template)]
@@ -59,5 +62,8 @@ struct ForRangeTemplate {
 #[test]
 fn test_for_range() {
     let s = ForRangeTemplate { init: -1, end: 1 };
-    assert_eq!(s.render().unwrap(), "foo\nfoo\nbar\nbar\nfoo\nbar\nbar\n");
+    assert_eq!(
+        s.render().unwrap(),
+        "foo (first)\nfoo (last)\nbar\nbar\nfoo\nbar\nbar\n"
+    );
 }