Support function arguments in Python AST

1 files changed, 14 insertions, 0 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 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()) {