Fix whitespace handling for macros and imports (#55)

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

* Fix of review comments

4 files changed, 33 insertions, 11 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) {
diff --git a/testing/templates/import.html b/testing/templates/import.html
index 87d7969..51ddd2f 100644
--- a/testing/templates/import.html
+++ b/testing/templates/import.html
@@ -1,2 +1,4 @@
-{% import "macro.html" -%}
-{%- call thrice(s) %}
+
+{%- import "macro.html" -%}
+
+{% call thrice(s) %}
diff --git a/testing/templates/macro.html b/testing/templates/macro.html
index d3ab77a..30ea742 100644
--- a/testing/templates/macro.html
+++ b/testing/templates/macro.html
@@ -1,5 +1,13 @@
-{% macro thrice(param) -%}
-  {{ param }} {{ param }} {{ param }}
-{%- endmacro %}
+1
 
-{%- call thrice(s) %}
+{%- macro thrice(param) -%}
+
+{{ param }} {{ param }} {{ param }}
+
+{%- endmacro -%}
+
+2
+
+{%- call thrice(s) -%}
+
+3
diff --git a/testing/tests/macro.rs b/testing/tests/macro.rs
index 25d95d8..fb54f88 100644
--- a/testing/tests/macro.rs
+++ b/testing/tests/macro.rs
@@ -12,7 +12,7 @@ struct MacroTemplate<'a> {
 #[test]
 fn test_macro() {
     let t = MacroTemplate { s: "foo" };
-    assert_eq!(t.render().unwrap(), "foo foo foo");
+    assert_eq!(t.render().unwrap(), "12foo foo foo3");
 }
 
 #[derive(Template)]