Handle imports as part of code generation

2 files changed, 17 insertions, 45 deletions

diff --git a/askama_derive/src/generator.rs b/askama_derive/src/generator.rs
index f7e7692..c6507d4 100644
--- a/askama_derive/src/generator.rs
+++ b/askama_derive/src/generator.rs
@@ -12,10 +12,11 @@ use std::collections::{HashMap, HashSet};
 use syn;
 
 
-pub fn generate(input: &TemplateInput, nodes: &[Node], imported: &HashMap<(&str, &str), Macro>)
-                -> String {
+pub fn generate(input: &TemplateInput, nodes: &[Node]) -> String {
     let mut base = None;
     let mut blocks = Vec::new();
+    let mut imported = Vec::new();
+    let mut parsed = Vec::new();
     let mut macros = HashMap::new();
 
     for n in nodes {
@@ -32,9 +33,15 @@ pub fn generate(input: &TemplateInput, nodes: &[Node], imported: &HashMap<(&str,
             Node::Macro(name, m) => {
                 macros.insert((None, *name), m);
             },
+            Node::Import(_, import_path, scope) => {
+                let path = path::find_template_from_path(import_path, Some(&input.path));
+                let src = path::get_template_source(&path);
+                imported.push((*scope, src));
+            }
             _ => {},
         }
     }
+    parsed.extend(imported.iter().map(|(scope, src)| (*scope, parser::parse(&src))));
 
     let mut check_nested = 0;
     let mut nested_blocks = Vec::new();
@@ -52,8 +59,13 @@ pub fn generate(input: &TemplateInput, nodes: &[Node], imported: &HashMap<(&str,
         check_nested += 1;
     }
 
-    for (&(scope, name), m) in imported {
-        macros.insert((Some(scope), name), m);
+    for (scope, ast) in &parsed {
+        for n in ast {
+            match n {
+                Node::Macro(name, m) => macros.insert((Some(*scope), name), &m),
+                _ => None,
+            };
+        }
     }
 
     Generator::default().build(&State {
diff --git a/askama_derive/src/lib.rs b/askama_derive/src/lib.rs
index 1faebec..b14b16f 100644
--- a/askama_derive/src/lib.rs
+++ b/askama_derive/src/lib.rs
@@ -12,13 +12,7 @@ mod generator;
 mod parser;
 
 use input::Print;
-use parser::{Macro, Node};
 use proc_macro::TokenStream;
-use shared::path;
-
-use std::borrow::Cow;
-use std::collections::HashMap;
-use std::path::Path;
 
 #[proc_macro_derive(Template, attributes(template))]
 pub fn derive_template(input: TokenStream) -> TokenStream {
@@ -40,46 +34,12 @@ pub fn derive_template(input: TokenStream) -> TokenStream {
 fn build_template(ast: &syn::DeriveInput) -> String {
     let data = input::TemplateInput::new(ast);
     let nodes = parser::parse(data.source.as_ref());
-    let imports = Imports::new(&nodes, &data.path);
     if data.meta.print == Print::Ast || data.meta.print == Print::All {
         println!("{:?}", nodes);
     }
-    let code = generator::generate(&data, &nodes, &imports.macro_map());
+    let code = generator::generate(&data, &nodes);
     if data.meta.print == Print::Code || data.meta.print == Print::All {
         println!("{}", code);
     }
     code
 }
-
-struct Imports<'a> {
-    sources: HashMap<&'a str, Cow<'a, str>>,
-}
-
-impl<'a> Imports<'a> {
-    fn new(parent_nodes: &'a [Node], parent_path: &'a Path) -> Imports<'a> {
-        let sources = parent_nodes.iter().filter_map(|n| {
-            match *n {
-                Node::Import(_, import_path, scope) => {
-                    let path = path::find_template_from_path(import_path, Some(parent_path));
-                    let src = path::get_template_source(&path);
-                    Some((scope, Cow::Owned(src)))
-                },
-                _ => None,
-            }
-        }).collect();
-        Imports { sources }
-    }
-
-    fn macro_map(&'a self) -> HashMap<(&'a str, &'a str), Macro<'a>> {
-        let mut macro_map = HashMap::new();
-        for (scope, s) in &self.sources {
-            for n in parser::parse(s.as_ref()) {
-                match n {
-                    Node::Macro(name, m) => macro_map.insert((*scope, name), m),
-                    _ => None,
-                };
-            }
-        }
-        macro_map
-    }
-}