parser: add type for `Node::Let`

3 files changed, 53 insertions, 57 deletions

diff --git a/askama_derive/src/generator.rs b/askama_derive/src/generator.rs
index 9781cda..643fb73 100644
--- a/askama_derive/src/generator.rs
+++ b/askama_derive/src/generator.rs
@@ -2,7 +2,9 @@ use crate::config::{get_template_source, read_config_file, Config, WhitespaceHan
 use crate::heritage::{Context, Heritage};
 use crate::input::{Print, Source, TemplateInput};
 use crate::CompileError;
-use parser::{Call, Cond, CondTest, Expr, Lit, Loop, Match, Node, Parsed, Target, Whitespace, Ws};
+use parser::{
+    Call, Cond, CondTest, Expr, Let, Lit, Loop, Match, Node, Parsed, Target, Whitespace, Ws,
+};
 
 use proc_macro::TokenStream;
 use quote::{quote, ToTokens};
@@ -635,11 +637,8 @@ impl<'a> Generator<'a> {
                 Node::Expr(ws, ref val) => {
                     self.write_expr(ws, val);
                 }
-                Node::LetDecl(ws, ref var) => {
-                    self.write_let_decl(buf, ws, var)?;
-                }
-                Node::Let(ws, ref var, ref val) => {
-                    self.write_let(buf, ws, var, val)?;
+                Node::Let(ref l) => {
+                    self.write_let(buf, l)?;
                 }
                 Node::Cond(ref conds, ws) => {
                     size_hint += self.write_cond(ctx, buf, conds, ws)?;
@@ -1055,19 +1054,6 @@ impl<'a> Generator<'a> {
         Ok(size_hint)
     }
 
-    fn write_let_decl(
-        &mut self,
-        buf: &mut Buffer,
-        ws: Ws,
-        var: &'a Target<'_>,
-    ) -> Result<(), CompileError> {
-        self.handle_ws(ws);
-        self.write_buf_writable(buf)?;
-        buf.write("let ");
-        self.visit_target(buf, false, true, var);
-        buf.writeln(";")
-    }
-
     fn is_shadowing_variable(&self, var: &Target<'a>) -> Result<bool, CompileError> {
         match var {
             Target::Name(name) => {
@@ -1103,31 +1089,33 @@ impl<'a> Generator<'a> {
         }
     }
 
-    fn write_let(
-        &mut self,
-        buf: &mut Buffer,
-        ws: Ws,
-        var: &'a Target<'_>,
-        val: &Expr<'_>,
-    ) -> Result<(), CompileError> {
-        self.handle_ws(ws);
+    fn write_let(&mut self, buf: &mut Buffer, l: &'a Let<'_>) -> Result<(), CompileError> {
+        self.handle_ws(l.ws);
+
+        let Some(val) = &l.val else {
+            self.write_buf_writable(buf)?;
+            buf.write("let ");
+            self.visit_target(buf, false, true, &l.var);
+            return buf.writeln(";");
+        };
+
         let mut expr_buf = Buffer::new(0);
         self.visit_expr(&mut expr_buf, val)?;
 
-        let shadowed = self.is_shadowing_variable(var)?;
+        let shadowed = self.is_shadowing_variable(&l.var)?;
         if shadowed {
             // Need to flush the buffer if the variable is being shadowed,
             // to ensure the old variable is used.
             self.write_buf_writable(buf)?;
         }
         if shadowed
-            || !matches!(var, &Target::Name(_))
-            || matches!(var, Target::Name(name) if self.locals.get(name).is_none())
+            || !matches!(l.var, Target::Name(_))
+            || matches!(&l.var, Target::Name(name) if self.locals.get(name).is_none())
         {
             buf.write("let ");
         }
 
-        self.visit_target(buf, true, true, var);
+        self.visit_target(buf, true, true, &l.var);
         buf.writeln(&format!(" = {};", &expr_buf.buf))
     }
 
diff --git a/askama_parser/src/lib.rs b/askama_parser/src/lib.rs
index c0650ee..560507a 100644
--- a/askama_parser/src/lib.rs
+++ b/askama_parser/src/lib.rs
@@ -16,7 +16,7 @@ use nom::{error_position, AsChar, IResult, InputTakeAtPosition};
 
 pub use self::expr::Expr;
 pub use self::node::{
-    BlockDef, Call, Cond, CondTest, Import, Lit, Loop, Macro, Match, Node, Raw, Target, When,
+    BlockDef, Call, Cond, CondTest, Import, Let, Lit, Loop, Macro, Match, Node, Raw, Target, When,
     Whitespace, Ws,
 };
 
diff --git a/askama_parser/src/node.rs b/askama_parser/src/node.rs
index ebab746..9055d99 100644
--- a/askama_parser/src/node.rs
+++ b/askama_parser/src/node.rs
@@ -20,8 +20,7 @@ pub enum Node<'a> {
     Comment(Ws),
     Expr(Ws, Expr<'a>),
     Call(Call<'a>),
-    LetDecl(Ws, Target<'a>),
-    Let(Ws, Target<'a>, Expr<'a>),
+    Let(Let<'a>),
     Cond(Vec<Cond<'a>>, Ws),
     Match(Match<'a>),
     Loop(Loop<'a>),
@@ -50,7 +49,7 @@ impl<'a> Node<'a> {
             |i| s.tag_block_start(i),
             alt((
                 map(Call::parse, Self::Call),
-                Self::r#let,
+                map(Let::parse, Self::Let),
                 |i| Self::r#if(i, s),
                 |i| Self::r#for(i, s),
                 map(|i| Match::parse(i, s), Self::Match),
@@ -69,28 +68,6 @@ impl<'a> Node<'a> {
         Ok((i, contents))
     }
 
-    fn r#let(i: &'a str) -> IResult<&'a str, Self> {
-        let mut p = tuple((
-            opt(Whitespace::parse),
-            ws(alt((keyword("let"), keyword("set")))),
-            cut(tuple((
-                ws(Target::parse),
-                opt(tuple((ws(char('=')), ws(Expr::parse)))),
-                opt(Whitespace::parse),
-            ))),
-        ));
-        let (i, (pws, _, (var, val, nws))) = p(i)?;
-
-        Ok((
-            i,
-            if let Some((_, val)) = val {
-                Self::Let(Ws(pws, nws), var, val)
-            } else {
-                Self::LetDecl(Ws(pws, nws), var)
-            },
-        ))
-    }
-
     fn r#if(i: &'a str, s: &State<'_>) -> IResult<&'a str, Self> {
         let mut p = tuple((
             opt(Whitespace::parse),
@@ -821,6 +798,37 @@ impl<'a> Raw<'a> {
     }
 }
 
+#[derive(Debug, PartialEq)]
+pub struct Let<'a> {
+    pub ws: Ws,
+    pub var: Target<'a>,
+    pub val: Option<Expr<'a>>,
+}
+
+impl<'a> Let<'a> {
+    fn parse(i: &'a str) -> IResult<&'a str, Self> {
+        let mut p = tuple((
+            opt(Whitespace::parse),
+            ws(alt((keyword("let"), keyword("set")))),
+            cut(tuple((
+                ws(Target::parse),
+                opt(preceded(ws(char('=')), ws(Expr::parse))),
+                opt(Whitespace::parse),
+            ))),
+        ));
+        let (i, (pws, _, (var, val, nws))) = p(i)?;
+
+        Ok((
+            i,
+            Let {
+                ws: Ws(pws, nws),
+                var,
+                val,
+            },
+        ))
+    }
+}
+
 /// First field is "minus/plus sign was used on the left part of the item".
 ///
 /// Second field is "minus/plus sign was used on the right part of the item".