2 files changed, 46 insertions, 36 deletions

diff --git a/src/main/java/org/zwobble/couscous/backends/python/PythonSerializer.java b/src/main/java/org/zwobble/couscous/backends/python/PythonSerializer.java
index af3b542..2fa1887 100644
--- a/src/main/java/org/zwobble/couscous/backends/python/PythonSerializer.java
+++ b/src/main/java/org/zwobble/couscous/backends/python/PythonSerializer.java
@@ -78,59 +78,60 @@ public class PythonSerializer implements PythonNodeVisitor<Void> {
 
     @Override
     public Void visit(PythonReturnNode pythonReturn) {
-        writer.startStatement();
-        writer.writeKeyword("return");
-        writer.writeSpace();
-        write(pythonReturn.getValue());
-        writer.endStatement();
+        writer.writeStatement(() -> {
+            writer.writeKeyword("return");
+            writer.writeSpace();
+            write(pythonReturn.getValue());
+        });
         return null;
     }
 
     @Override
     public Void visit(PythonPassNode pass) {
-        writer.startStatement();
-        writer.writeKeyword("pass");
-        writer.endStatement();
+        writer.writeStatement(() -> {
+            writer.writeKeyword("pass");            
+        });
         return null;
     }
 
     @Override
     public Void visit(PythonClassNode pythonClass) {
-        writer.startStatement();
-        writer.writeKeyword("class");
-        writer.writeSpace();
-        writer.writeIdentifier(pythonClass.getName());
-        writer.writeSymbol("(");
-        writer.writeIdentifier("object");
-        writer.writeSymbol(")");
-        writeBlock(pythonClass.getBody());
-        writer.endStatement();
+        writer.writeStatement(() -> {
+            writer.writeKeyword("class");
+            writer.writeSpace();
+            writer.writeIdentifier(pythonClass.getName());
+            writer.writeSymbol("(");
+            writer.writeIdentifier("object");
+            writer.writeSymbol(")");
+            writeBlock(pythonClass.getBody());
+        });
         return null;
     }
 
     @Override
     public Void visit(PythonFunctionDefinitionNode functionDefinition) {
-        writer.startStatement();
-        writer.writeKeyword("def");
-        writer.writeSpace();
-        writer.writeIdentifier(functionDefinition.getName());
-        writer.writeSymbol("(");
-        writeArgumentNames(functionDefinition);
-        writer.writeSymbol(")");
-        writeBlock(functionDefinition.getBody());
-        writer.endStatement();
+        writer.writeStatement(() -> {
+            writer.writeKeyword("def");
+            writer.writeSpace();
+            writer.writeIdentifier(functionDefinition.getName());
+            writer.writeSymbol("(");
+            writeArgumentNames(functionDefinition);
+            writer.writeSymbol(")");
+            writeBlock(functionDefinition.getBody());
+        });
         return null;
     }
 
     @Override
     public Void visit(PythonAssignmentNode assignment) {
-        writer.startStatement();
-        write(assignment.getTarget());
-        writer.writeSpace();
-        writer.writeSymbol("=");
-        writer.writeSpace();
-        write(assignment.getValue());
-        writer.endStatement();
+        writer.writeStatement(() -> {
+            write(assignment.getTarget());
+            writer.writeSpace();
+            writer.writeSymbol("=");
+            writer.writeSpace();
+            write(assignment.getValue());
+            
+        });
         return null;
     }
 
diff --git a/src/main/java/org/zwobble/couscous/backends/python/PythonWriter.java b/src/main/java/org/zwobble/couscous/backends/python/PythonWriter.java
index 2551700..781db71 100644
--- a/src/main/java/org/zwobble/couscous/backends/python/PythonWriter.java
+++ b/src/main/java/org/zwobble/couscous/backends/python/PythonWriter.java
@@ -47,14 +47,23 @@ public class PythonWriter {
         depth--;
     }
 
-    public void startStatement() {
+    public void writeStatement(Action action) {
+        startStatement();
+        action.run();
+        endStatement();
+    }
+
+    private void startStatement() {
         builder.append(Strings.repeat(" ", depth * SPACES_PER_INDENT));
     }
 
-    public void endStatement() {
+    private void endStatement() {
         if (builder.charAt(builder.length() - 1) != '\n') {
             builder.append("\n");   
         }
     }
-
+    
+    interface Action {
+        void run();
+    }
 }