Support function arguments in Python AST

1 files changed, 26 insertions, 2 deletions

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