Fix whitespace handling for macros and imports (#55)

Resolves https://github.com/djc/askama/issues/52.

* Fix of review comments

1 files changed, 16 insertions, 4 deletions

diff --git a/askama_shared/src/generator.rs b/askama_shared/src/generator.rs
index a37b51f..acb3f15 100644
--- a/askama_shared/src/generator.rs
+++ b/askama_shared/src/generator.rs
@@ -353,11 +353,22 @@ impl<'a> Generator<'a> {
                     self.handle_include(state, ws, path);
                 },
                 Node::Call(ref ws, name, ref args) => self.write_call(state, ws, name, args),
-                Node::Macro(_, _) |
-                Node::Import(_, _) |
+                Node::Macro(_, ref m) => {
+                    if let AstLevel::Nested = level {
+                        panic!("macro blocks only allowed at the top level");
+                    }
+                    self.flush_ws(&m.ws1);
+                    self.prepare_ws(&m.ws2);
+                },
+                Node::Import(ref ws, _) => {
+                    if let AstLevel::Nested = level {
+                        panic!("import blocks only allowed at the top level");
+                    }
+                    self.handle_ws(ws);
+                },
                 Node::Extends(_) => {
                     if let AstLevel::Nested = level {
-                        panic!("macro or extend blocks only allowed at the top level");
+                        panic!("extend blocks only allowed at the top level");
                     }
                 },
             }
@@ -436,7 +447,7 @@ impl<'a> Generator<'a> {
 
     fn write_call(&mut self, state: &'a State, ws: &WS, name: &str, args: &[Expr]) {
         let def = state.macros.get(name).expect(&format!("macro '{}' not found", name));
-        self.handle_ws(ws);
+        self.flush_ws(ws);
         self.locals.push();
         self.writeln("{");
         self.prepare_ws(&def.ws1);
@@ -451,6 +462,7 @@ impl<'a> Generator<'a> {
         self.flush_ws(&def.ws2);
         self.writeln("}");
         self.locals.pop();
+        self.prepare_ws(ws);
     }
 
     fn handle_include(&mut self, state: &'a State, ws: &WS, path: &str) {