Add support for macro args type annotations

Author: GuillaumeGomez

askama_derive/src/generator/expr.rs

@@ -509,14 +509,17 @@ impl<'a> Generator<'a, '_> {
         buf: &mut Buffer,
         generics: &WithSpan<Vec<WithSpan<TyGenerics<'a>>>>,
     ) {
-        let mut tmp = Buffer::new();
+        if generics.is_empty() {
+            return;
+        }
+        let generics_span = ctx.span_for_node(generics.span());
+        buf.write_token(Token![<], generics_span);
         for generic in &**generics {
             let span = ctx.span_for_node(generic.span());
-            self.visit_ty_generic(ctx, &mut tmp, generic, span);
-            tmp.write_token(Token![,], span);
+            self.visit_ty_generic(ctx, buf, generic, span);
+            buf.write_token(Token![,], span);
         }
-        let tmp = tmp.into_token_stream();
-        quote_into!(buf, ctx.span_for_node(generics.span()), { <#tmp> });
+        buf.write_token(Token![>], generics_span);
     }
 
     pub(super) fn visit_ty_generic(
@@ -530,7 +533,7 @@ impl<'a> Generator<'a, '_> {
         for _ in 0..refs {
             buf.write_token(Token![&], span);
         }
-        match kind {
+        match &**kind {
             TyGenericsKind::Path { path, args } => {
                 self.visit_macro_path(buf, path, span);
                 if let Some(generics) = args.as_ref() {

askama_derive/src/generator/helpers/macro_invocation.rs

@@ -2,9 +2,11 @@ use core::fmt;
 use std::borrow::Cow;
 use std::fmt::Write;
 use std::mem;
+use std::str::FromStr;
 
-use parser::node::{Call, Macro, Ws};
+use parser::node::{Call, Macro, MacroArg, Ws};
 use parser::{Expr, Span, WithSpan};
+use proc_macro2::TokenStream;
 use quote::quote_spanned;
 
 use crate::generator::node::AstLevel;
@@ -95,19 +97,37 @@ impl<'a, 'b> MacroInvocation<'a, 'b> {
     fn handle_macro_arg<'h>(
         &self,
         expr: &WithSpan<Box<Expr<'a>>>,
-        arg_name: &WithSpan<&'a str>,
+        arg: &MacroArg<'a>,
         buf: &mut Buffer,
         generator: &mut Generator<'a, 'h>,
     ) -> Result<(), CompileError> {
+        let mut ty_buf = Buffer::new();
+        let span = self.callsite_ctx.span_for_node(arg.name.span());
+        if let Some(ref ty) = arg.ty {
+            ty_buf.write_token(syn::Token![:], span);
+            // To prevent moving the value/variable, we take a reference of it.
+            ty_buf.write_token(syn::Token![&], span);
+            generator.visit_ty_generic(self.callsite_ctx, &mut ty_buf, ty, span);
+        }
+
         match &***expr {
             // If `expr` is already a form of variable then
             // don't reintroduce a new variable. This is
             // to avoid moving non-copyable values.
             &Expr::Var(name) if name != "self" => {
                 let var = generator.locals.resolve_or_self(name);
-                generator
-                    .locals
-                    .insert(Cow::Borrowed(**arg_name), LocalMeta::var_with_ref(var));
+                if arg.ty.is_none() {
+                    generator
+                        .locals
+                        .insert(Cow::Borrowed(*arg.name), LocalMeta::var_with_ref(var));
+                } else {
+                    let id = field_new(&arg.name, span);
+                    let var = TokenStream::from_str(&var).expect("invalid variable name");
+                    buf.write_tokens(quote_spanned! { span => let #id #ty_buf = &#var; });
+                    generator
+                        .locals
+                        .insert_with_default(Cow::Borrowed(*arg.name));
+                }
             }
             Expr::AssociatedItem(obj, associated_item) => {
                 let mut associated_item_buf = Buffer::new();
@@ -125,9 +145,18 @@ impl<'a, 'b> MacroInvocation<'a, 'b> {
                     .locals
                     .resolve(&associated_item)
                     .unwrap_or(associated_item);
-                generator
-                    .locals
-                    .insert(Cow::Borrowed(**arg_name), LocalMeta::var_with_ref(var));
+                if arg.ty.is_none() {
+                    generator
+                        .locals
+                        .insert(Cow::Borrowed(*arg.name), LocalMeta::var_with_ref(var));
+                } else {
+                    let id = field_new(&arg.name, span);
+                    let var = TokenStream::from_str(&var).expect("invalid variable name");
+                    buf.write_tokens(quote_spanned! { span => let #id #ty_buf = &#var; });
+                    generator
+                        .locals
+                        .insert_with_default(Cow::Borrowed(*arg.name));
+                }
             }
             // Everything else still needs to become variables,
             // to avoid having the same logic be executed
@@ -136,17 +165,16 @@ impl<'a, 'b> MacroInvocation<'a, 'b> {
             _ => {
                 let mut value = Buffer::new();
                 value.write_tokens(generator.visit_expr_root(self.callsite_ctx, expr)?);
-                let span = self.callsite_ctx.span_for_node(arg_name.span());
-                let id = field_new(arg_name, span);
-                buf.write_tokens(if !is_copyable(expr) {
-                    quote_spanned! { span => let #id = &(#value); }
+                let id = field_new(&arg.name, span);
+                buf.write_tokens(if !is_copyable(expr) || arg.ty.is_some() {
+                    quote_spanned! { span => let #id #ty_buf = &(#value); }
                 } else {
-                    quote_spanned! { span => let #id = #value; }
+                    quote_spanned! { span => let #id #ty_buf = #value; }
                 });
 
                 generator
                     .locals
-                    .insert_with_default(Cow::Borrowed(**arg_name));
+                    .insert_with_default(Cow::Borrowed(*arg.name));
             }
         }
 
@@ -216,7 +244,7 @@ impl<'a, 'b> MacroInvocation<'a, 'b> {
                     }
                 }
             };
-            self.handle_macro_arg(expr, &arg.name, buf, generator)?;
+            self.handle_macro_arg(expr, arg, buf, generator)?;
         }
 
         Ok(())

askama_parser/src/expr.rs

@@ -23,6 +23,8 @@ macro_rules! expr_prec_layer {
     };
 }
 
+const MAX_REFS: usize = 20;
+
 fn expr_prec_layer<'a: 'l, 'l>(
     i: &mut InputStream<'a, 'l>,
     inner: fn(&mut InputStream<'a, 'l>) -> ParseResult<'a, WithSpan<Box<Expr<'a>>>>,
@@ -1495,19 +1497,24 @@ fn ensure_macro_name<'a>(name: &WithSpan<&'a str>) -> ParseResult<'a, ()> {
 #[derive(Clone, Debug, PartialEq)]
 pub struct TyGenerics<'a> {
     pub refs: usize,
-    pub kind: TyGenericsKind<'a>,
+    pub kind: WithSpan<TyGenericsKind<'a>>,
 }
 
 impl<'a: 'l, 'l> TyGenerics<'a> {
-    fn parse(i: &mut InputStream<'a, 'l>) -> ParseResult<'a, WithSpan<Self>> {
-        let p = ws((repeat(0.., ws('&')), TyGenericsKind::parse));
-        let ((refs, kind), span) = p.with_span().parse_next(i)?;
-        let max_refs = 20;
-        if refs > max_refs {
-            return cut_error!(format!("too many references (> {max_refs})"), span);
+    pub(crate) fn parse(i: &mut InputStream<'a, 'l>) -> ParseResult<'a, WithSpan<Self>> {
+        let p = ws((repeat(0.., ws('&')), TyGenericsKind::parse.with_span()));
+        let ((refs, (kind, kind_span)), span) = p.with_span().parse_next(i)?;
+        if refs > MAX_REFS {
+            return cut_error!(format!("too many references (> {MAX_REFS})"), span);
         }
 
-        Ok(WithSpan::new(TyGenerics { refs, kind }, span))
+        Ok(WithSpan::new(
+            TyGenerics {
+                refs,
+                kind: WithSpan::new(kind, kind_span),
+            },
+            span,
+        ))
     }
 
     fn args(

askama_parser/src/node.rs

@@ -12,8 +12,8 @@ use winnow::{ModalParser, Parser};
 use crate::expr::BinOp;
 use crate::{
     ErrorContext, Expr, Filter, HashSet, InputStream, ParseErr, ParseResult, Span, Target,
-    WithSpan, block_end, block_start, cut_error, deny_any_rust_token, expr_end, expr_start, filter,
-    identifier, is_rust_keyword, keyword, skip_ws0, str_lit_without_prefix, ws,
+    TyGenerics, WithSpan, block_end, block_start, cut_error, deny_any_rust_token, expr_end,
+    expr_start, filter, identifier, is_rust_keyword, keyword, skip_ws0, str_lit_without_prefix, ws,
 };
 
 #[derive(Debug, PartialEq)]
@@ -655,6 +655,7 @@ pub struct Macro<'a> {
 #[derive(Debug, PartialEq)]
 pub struct MacroArg<'a> {
     pub name: WithSpan<&'a str>,
+    pub ty: Option<WithSpan<TyGenerics<'a>>>,
     pub default: Option<WithSpan<Box<Expr<'a>>>>,
 }
 
@@ -690,11 +691,13 @@ impl<'a: 'l, 'l> Macro<'a> {
         let macro_arg = |i: &mut _| {
             let mut p = (
                 ws(identifier.with_span()),
+                opt(preceded(':', ws(|i: &mut _| TyGenerics::parse(i)))),
                 opt(preceded('=', ws(|i: &mut _| Expr::parse(i, false)))),
             );
-            let ((name, name_span), default) = p.parse_next(i)?;
+            let ((name, name_span), ty, default) = p.parse_next(i)?;
             Ok(MacroArg {
                 name: WithSpan::new(name, name_span),
+                ty,
                 default,
             })
         };

testing/tests/macro.rs

@@ -705,3 +705,103 @@ val4: {{val4}}
         "val1: aa\nval2: x\nval3: default\nval4: c"
     );
 }
+
+// Goal of these tests is to ensure that the argument info is working as expected.
+#[test]
+fn test_macro_with_args_type_info() {
+    #[derive(Template)]
+    #[template(
+        source = r#"
+{%- macro test(title: Option<u32> = None) -%}
+-> {{title|fmt("{:?}")}}
+{% endmacro -%}
+
+{% let y = Some(12) -%}
+{% call test(y) %}{% endcall -%}
+{% call test(x) %}{% endcall -%}
+{% call test() %}{% endcall -%}
+        "#,
+        ext = "txt"
+    )]
+    struct F {
+        x: Option<u32>,
+    }
+
+    assert_eq!(
+        F { x: Some(4) }.render().unwrap(),
+        "-> Some(12)\n-> Some(4)\n-> None\n"
+    );
+}
+
+#[test]
+fn test_macro_with_args_type_info2() {
+    #[derive(Template)]
+    #[template(
+        source = r#"
+{%- macro test(entries: &[u32]) -%}
+{%- for entry in entries -%}
+-> {{entry}}
+{% endfor -%}
+{% endmacro -%}
+{{- test(x.as_slice()) -}}
+        "#,
+        ext = "txt"
+    )]
+    struct F {
+        x: Vec<u32>,
+    }
+
+    assert_eq!(F { x: vec![4, 2] }.render().unwrap(), "-> 4\n-> 2\n");
+}
+
+#[test]
+fn test_macro_with_args_type_info3() {
+    #[derive(Template)]
+    #[template(
+        source = r#"
+{%- macro test(entries: std::vec::Vec<u32>) -%}
+{%- for entry in entries -%}
+-> {{entry}}
+{% endfor -%}
+{% endmacro -%}
+{{- test(x) -}}
+        "#,
+        ext = "txt"
+    )]
+    struct F {
+        x: Vec<u32>,
+    }
+    assert_eq!(F { x: vec![4, 2] }.render().unwrap(), "-> 4\n-> 2\n");
+}
+
+#[test]
+fn test_macro_with_args_type_info4() {
+    #[derive(Template)]
+    #[template(
+        source = r#"
+{%- macro test(entries: &[Vec<u32>]) -%}
+{%- for entry in entries -%}
++>
+{%- for sub_entry in entry -%}
+-> {{sub_entry}}
+{% endfor -%}
+{% endfor -%}
+{% endmacro -%}
+
+{{- test(x.as_slice()) -}}
+        "#,
+        ext = "txt"
+    )]
+    struct F {
+        x: Vec<Vec<u32>>,
+    }
+
+    assert_eq!(
+        F {
+            x: vec![vec![4, 2], vec![5, 1]]
+        }
+        .render()
+        .unwrap(),
+        "+>-> 4\n-> 2\n+>-> 5\n-> 1\n",
+    );
+}

testing/tests/ui/macro-args-with-type.rs

@@ -0,0 +1,17 @@
+use askama::Template;
+
+#[derive(Template)]
+#[template(
+    source = r#"
+        {%- macro test(title: Option<u32> = None) -%}{% endmacro -%}
+        {%- let y = 12 -%}
+        {% call test(y) %}{% endcall -%}
+        {% call test(x) %}{% endcall -%}
+    "#,
+    ext = "txt",
+)]
+struct F {
+    x: u32,
+}
+
+fn main() {}

testing/tests/ui/macro-args-with-type.stderr

@@ -0,0 +1,19 @@
+error[E0308]: mismatched types
+ --> tests/ui/macro-args-with-type.rs:3:10
+  |
+3 | #[derive(Template)]
+  |          ^^^^^^^^ expected `&Option<u32>`, found `&{integer}`
+  |
+  = note: expected reference `&Option<u32>`
+             found reference `&{integer}`
+  = note: this error originates in the derive macro `Template` (in Nightly builds, run with -Z macro-backtrace for more info)
+
+error[E0308]: mismatched types
+ --> tests/ui/macro-args-with-type.rs:3:10
+  |
+3 | #[derive(Template)]
+  |          ^^^^^^^^ expected `&Option<u32>`, found `&u32`
+  |
+  = note: expected reference `&Option<u32>`
+             found reference `&u32`
+  = note: this error originates in the derive macro `Template` (in Nightly builds, run with -Z macro-backtrace for more info)