Implement support for macro definition handling

2 files changed, 59 insertions, 8 deletions

diff --git a/askama_derive/src/generator.rs b/askama_derive/src/generator.rs
index aa3eadf..dd34bec 100644
--- a/askama_derive/src/generator.rs
+++ b/askama_derive/src/generator.rs
@@ -5,7 +5,7 @@ use quote::{Tokens, ToTokens};
 
 use std::{cmp, hash, str};
 use std::path::Path;
-use std::collections::HashSet;
+use std::collections::{HashMap, HashSet};
 
 use syn;
 
@@ -14,6 +14,7 @@ pub fn generate(ast: &syn::DeriveInput, path: &Path, mut nodes: Vec<Node>) -> St
     let mut blocks = Vec::new();
     let mut block_names = Vec::new();
     let mut content = Vec::new();
+    let mut macros = HashMap::new();
     for n in nodes.drain(..) {
         match n {
             Node::Extends(path) => {
@@ -27,11 +28,14 @@ pub fn generate(ast: &syn::DeriveInput, path: &Path, mut nodes: Vec<Node>) -> St
                 block_names.push(name);
                 content.push(Node::Block(ws1, name, ws2));
             },
+            Node::Macro(ws1, name, params, contents, ws2) => {
+                macros.insert(name, (ws1, name, params, contents, ws2));
+            },
             _ => { content.push(n); },
         }
     }
 
-    let mut gen = Generator::default();
+    let mut gen = Generator::default(&macros);
     if !blocks.is_empty() {
         let trait_name = trait_name_for_path(&base, path);
         if base.is_none() {
@@ -96,11 +100,12 @@ struct Generator<'a> {
     locals: SetChain<'a, &'a str>,
     next_ws: Option<&'a str>,
     skip_ws: bool,
+    macros: &'a MacroMap<'a>,
 }
 
 impl<'a> Generator<'a> {
 
-    fn new<'n>(locals: SetChain<'n, &'n str>, indent: u8) -> Generator<'n> {
+    fn new<'n>(macros: &'n MacroMap, locals: SetChain<'n, &'n str>, indent: u8) -> Generator<'n> {
         Generator {
             buf: String::new(),
             indent: indent,
@@ -108,16 +113,17 @@ impl<'a> Generator<'a> {
             locals: locals,
             next_ws: None,
             skip_ws: false,
+            macros: macros,
         }
     }
 
-    fn default<'n>() -> Generator<'n> {
-        Self::new(SetChain::new(), 0)
+    fn default<'n>(macros: &'n MacroMap) -> Generator<'n> {
+        Self::new(macros, SetChain::new(), 0)
     }
 
     fn child<'n>(&'n mut self) -> Generator<'n> {
         let locals = SetChain::with_parent(&self.locals);
-        Self::new(locals, self.indent)
+        Self::new(self.macros, locals, self.indent)
     }
 
     /* Helper methods for writing to internal buffer */
@@ -447,8 +453,9 @@ impl<'a> Generator<'a> {
                 Node::Include(ref ws, ref path) => {
                     self.handle_include(ws, path);
                 },
+                Node::Macro(_, _, _, _, _) |
                 Node::Extends(_) => {
-                    panic!("no extends or block definition allowed in content");
+                    panic!("no extends or macros allowed in content");
                 },
             }
         }
@@ -627,3 +634,5 @@ impl<'a, T: 'a> SetChain<'a, T> where T: cmp::Eq + hash::Hash {
         assert!(self.scopes.len() > 0);
     }
 }
+
+type MacroMap<'a> = HashMap<&'a str, (WS, &'a str, Vec<&'a str>, Vec<Node<'a>>, WS)>;
diff --git a/askama_derive/src/parser.rs b/askama_derive/src/parser.rs
index 9e8f087..9a49715 100644
--- a/askama_derive/src/parser.rs
+++ b/askama_derive/src/parser.rs
@@ -34,6 +34,7 @@ pub enum Node<'a> {
     BlockDef(WS, &'a str, Vec<Node<'a>>, WS),
     Block(WS, &'a str, WS),
     Include(WS, &'a str),
+    Macro(WS, &'a str, Vec<&'a str>, Vec<Node<'a>>, WS),
 }
 
 pub type Cond<'a> = (WS, Option<Expr<'a>>, Vec<Node<'a>>);
@@ -147,6 +148,25 @@ named!(arguments<Option<Vec<Expr>>>, opt!(
     )
 ));
 
+named!(parameters<Vec<&'a str>>, do_parse!(
+    tag_s!("(") >>
+    vals: opt!(do_parse!(
+        arg0: ws!(identifier) >>
+        args: many0!(do_parse!(
+            tag_s!(",") >>
+            argn: ws!(identifier) >>
+            (argn)
+        )) >>
+        ({
+            let mut res = vec![arg0];
+            res.extend(args);
+            res
+        })
+    )) >>
+    tag_s!(")") >>
+    (vals.unwrap_or(Vec::new()))
+));
+
 named!(expr_group<Expr>, map!(
     delimited!(char!('('), expr_any, char!(')')),
     |s| Expr::Group(Box::new(s))
@@ -347,6 +367,27 @@ named!(block_include<Node>, do_parse!(
     }))
 ));
 
+named!(block_macro<Node>, do_parse!(
+    pws1: opt!(tag_s!("-")) >>
+    ws!(tag_s!("macro")) >>
+    name: ws!(identifier) >>
+    params: ws!(parameters) >>
+    nws1: opt!(tag_s!("-")) >>
+    tag_s!("%}") >>
+    contents: parse_template >>
+    tag_s!("{%") >>
+    pws2: opt!(tag_s!("-")) >>
+    ws!(tag_s!("endmacro")) >>
+    nws2: opt!(tag_s!("-")) >>
+    (Node::Macro(
+         WS(pws1.is_some(), nws1.is_some()),
+         name,
+         params,
+         contents,
+         WS(pws2.is_some(), nws2.is_some())
+    ))
+));
+
 named!(block_node<Node>, do_parse!(
     tag_s!("{%") >>
     contents: alt!(
@@ -355,7 +396,8 @@ named!(block_node<Node>, do_parse!(
         block_for |
         block_extends |
         block_include |
-        block_block
+        block_block |
+        block_macro
     ) >>
     tag_s!("%}") >>
     (contents)