Support function arguments in Python AST

3 files changed, 42 insertions, 3 deletions

diff --git a/src/main/java/org/zwobble/couscous/backends/python/PythonCodeGenerator.java b/src/main/java/org/zwobble/couscous/backends/python/PythonCodeGenerator.java
index eae713c..95e54a5 100644
--- a/src/main/java/org/zwobble/couscous/backends/python/PythonCodeGenerator.java
+++ b/src/main/java/org/zwobble/couscous/backends/python/PythonCodeGenerator.java
@@ -17,6 +17,7 @@ import org.zwobble.couscous.ast.TernaryConditionalNode;
 import org.zwobble.couscous.ast.VariableReferenceNode;
 import org.zwobble.couscous.ast.visitors.ExpressionNodeVisitor;
 import org.zwobble.couscous.ast.visitors.StatementNodeVisitor;
+import org.zwobble.couscous.backends.python.ast.PythonBlock;
 import org.zwobble.couscous.backends.python.ast.PythonExpressionNode;
 import org.zwobble.couscous.backends.python.ast.PythonFunctionDefinitionNode;
 import org.zwobble.couscous.backends.python.ast.PythonModuleNode;
@@ -53,7 +54,7 @@ public class PythonCodeGenerator {
             .stream()
             .map(PythonCodeGenerator::generateStatement)
             .collect(Collectors.toList());
-        return pythonFunctionDefinition(method.getName(), pythonBody);
+        return pythonFunctionDefinition(method.getName(), asList(), new PythonBlock(pythonBody));
     }
     
     private static PythonStatementNode generateStatement(StatementNode statement) {
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 b6adb80..c6cf516 100644
--- a/src/main/java/org/zwobble/couscous/backends/python/PythonSerializer.java
+++ b/src/main/java/org/zwobble/couscous/backends/python/PythonSerializer.java
@@ -13,6 +13,8 @@ import org.zwobble.couscous.backends.python.ast.PythonStringLiteralNode;
 import org.zwobble.couscous.backends.python.ast.PythonWriter;
 import org.zwobble.couscous.backends.python.ast.visitors.PythonNodeVisitor;
 
+import static com.google.common.collect.Iterables.skip;
+
 import lombok.val;
 
 public class PythonSerializer implements PythonNodeVisitor<Void> {
@@ -90,12 +92,24 @@ public class PythonSerializer implements PythonNodeVisitor<Void> {
         writer.writeSpace();
         writer.writeIdentifier(functionDefinition.getName());
         writer.writeSymbol("(");
+        writeArgumentNames(functionDefinition);
         writer.writeSymbol(")");
         writeBlock(functionDefinition.getBody());
         writer.endStatement();
         return null;
     }
 
+    private void writeArgumentNames(PythonFunctionDefinitionNode functionDefinition) {
+        if (!functionDefinition.getArgumentNames().isEmpty()) {
+            writer.writeIdentifier(functionDefinition.getArgumentNames().get(0));
+            for (val argumentName : skip(functionDefinition.getArgumentNames(), 1)) {
+                writer.writeSymbol(",");
+                writer.writeSpace();
+                writer.writeIdentifier(argumentName);
+            }
+        }
+    }
+
     @Override
     public Void visit(PythonModuleNode module) {
         for (val statement : module.getStatements()) {
diff --git a/src/main/java/org/zwobble/couscous/backends/python/ast/PythonFunctionDefinitionNode.java b/src/main/java/org/zwobble/couscous/backends/python/ast/PythonFunctionDefinitionNode.java
index 6d66df8..c67fe35 100644
--- a/src/main/java/org/zwobble/couscous/backends/python/ast/PythonFunctionDefinitionNode.java
+++ b/src/main/java/org/zwobble/couscous/backends/python/ast/PythonFunctionDefinitionNode.java
@@ -4,15 +4,39 @@ import java.util.List;
 
 import org.zwobble.couscous.backends.python.ast.visitors.PythonNodeVisitor;
 
+import com.google.common.collect.ImmutableList;
+
 import lombok.Value;
 
 @Value(staticConstructor="pythonFunctionDefinition")
 public class PythonFunctionDefinitionNode implements PythonStatementNode {
-    public static PythonFunctionDefinitionNode pythonFunctionDefinition(String name, List<PythonStatementNode> body) {
-        return pythonFunctionDefinition(name, new PythonBlock(body));
+    public static Builder builder(String name) {
+        return new Builder(name);
+    }
+    
+    public static class Builder {
+        private final String name;
+        private final ImmutableList.Builder<String> argumentNames;
+        private final ImmutableList.Builder<PythonStatementNode> body;
+        
+        public Builder(String name) {
+            this.name = name;
+            this.argumentNames = ImmutableList.builder();
+            this.body = ImmutableList.builder();
+        }
+
+        public Builder argument(String argumentName) {
+            argumentNames.add(argumentName);
+            return this;
+        }
+        
+        public PythonFunctionDefinitionNode build() {
+            return pythonFunctionDefinition(name, argumentNames.build(), new PythonBlock(body.build()));
+        }
     }
     
     String name;
+    List<String> argumentNames;
     PythonBlock body;
     
     @Override