Add some helper methods to Generator impl

1 files changed, 40 insertions, 11 deletions

diff --git a/askama_codegen/src/generator.rs b/askama_codegen/src/generator.rs
index 13a91dd..c815676 100644
--- a/askama_codegen/src/generator.rs
+++ b/askama_codegen/src/generator.rs
@@ -3,36 +3,63 @@ use std::str;
 
 struct Generator {
     buf: String,
+    indent: u8,
+    start: bool,
 }
 
 impl Generator {
 
     fn new() -> Generator {
-        Generator { buf: String::new() }
+        Generator { buf: String::new(), indent: 0, start: true }
     }
 
     fn init(&mut self, name: &str) {
         self.write("impl askama::Template for ");
         self.write(name);
-        self.write(" {\n");
-        self.write("    fn render(&self) -> String {\n");
-        self.write("        let mut buf = String::new();\n");
+        self.writeln(" {");
+        self.indent();
+        self.writeln("fn render(&self) -> String {");
+        self.indent();
+        self.writeln("let mut buf = String::new();");
+    }
+
+    fn indent(&mut self) {
+        self.indent += 1;
+    }
+
+    fn dedent(&mut self) {
+        self.indent -= 1;
     }
 
     fn write(&mut self, s: &str) {
+        if self.start {
+            for _ in 0..(self.indent * 4) {
+                self.buf.push(' ');
+            }
+            self.start = false;
+        }
         self.buf.push_str(s);
     }
 
+    fn writeln(&mut self, s: &str) {
+        if s.is_empty() {
+            return;
+        }
+        self.write(s);
+        self.buf.push('\n');
+        self.start = true;
+    }
+
     fn visit_lit(&mut self, s: &[u8]) {
-        self.write("        buf.push_str(");
+        self.write("buf.push_str(");
         self.write(&format!("{:#?}", str::from_utf8(s).unwrap()));
-        self.write(");\n");
+        self.writeln(");");
     }
 
     fn visit_expr(&mut self, s: &[u8]) {
-        self.write("        buf.push_str(");
+        self.write("buf.push_str(");
         self.write(&format!("&self.{}", str::from_utf8(s).unwrap()));
-        self.write(");\n");
+        self.writeln(");");
     }
 
     fn handle(&mut self, tokens: &Vec<Node>) {
@@ -45,9 +72,11 @@ impl Generator {
     }
 
     fn finalize(&mut self) {
-        self.write("        buf");
-        self.write("    }\n");
-        self.write("}\n\n");
+        self.writeln("buf");
+        self.dedent();
+        self.writeln("}");
+        self.dedent();
+        self.writeln("}");
     }
 
     fn result(self) -> String {