Support passing arguments to Python runner

4 files changed, 59 insertions, 25 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 95e54a5..63fa56b 100644
--- a/src/main/java/org/zwobble/couscous/backends/python/PythonCodeGenerator.java
+++ b/src/main/java/org/zwobble/couscous/backends/python/PythonCodeGenerator.java
@@ -24,6 +24,7 @@ import org.zwobble.couscous.backends.python.ast.PythonModuleNode;
 import org.zwobble.couscous.backends.python.ast.PythonStatementNode;
 import org.zwobble.couscous.values.BooleanValue;
 import org.zwobble.couscous.values.IntegerValue;
+import org.zwobble.couscous.values.PrimitiveValue;
 import org.zwobble.couscous.values.PrimitiveValueVisitor;
 import org.zwobble.couscous.values.StringValue;
 import org.zwobble.couscous.values.UnitValue;
@@ -35,6 +36,7 @@ import static org.zwobble.couscous.backends.python.ast.PythonIntegerLiteralNode.
 import static org.zwobble.couscous.backends.python.ast.PythonModuleNode.pythonModule;
 import static org.zwobble.couscous.backends.python.ast.PythonReturnNode.pythonReturn;
 import static org.zwobble.couscous.backends.python.ast.PythonStringLiteralNode.pythonStringLiteral;
+import static org.zwobble.couscous.backends.python.ast.PythonVariableReferenceNode.pythonVariableReference;
 
 import lombok.val;
 
@@ -49,12 +51,41 @@ public class PythonCodeGenerator {
         return pythonModule(asList(pythonClass));
     }
     
+    public static PythonExpressionNode generateCode(PrimitiveValue value) {
+        return value.accept(new PrimitiveValueVisitor<PythonExpressionNode>() {
+            @Override
+            public PythonExpressionNode visit(IntegerValue value) {
+                return pythonIntegerLiteral(value.getValue());
+            }
+
+            @Override
+            public PythonExpressionNode visit(StringValue value) {
+                return pythonStringLiteral(value.getValue());
+            }
+
+            @Override
+            public PythonExpressionNode visit(BooleanValue value) {
+                throw new UnsupportedOperationException();
+            }
+
+            @Override
+            public PythonExpressionNode visit(UnitValue unitValue) {
+                throw new UnsupportedOperationException();
+            }
+        });
+    }
+    
     private static PythonFunctionDefinitionNode generateFunction(MethodNode method) {
+        val argumentNames = method.getArguments().stream()
+            .map(argument -> argument.getName())
+            .collect(Collectors.toList());
+            
         val pythonBody = method.getBody()
             .stream()
             .map(PythonCodeGenerator::generateStatement)
             .collect(Collectors.toList());
-        return pythonFunctionDefinition(method.getName(), asList(), new PythonBlock(pythonBody));
+        
+        return pythonFunctionDefinition(method.getName(), argumentNames, new PythonBlock(pythonBody));
     }
     
     private static PythonStatementNode generateStatement(StatementNode statement) {
@@ -85,33 +116,12 @@ public class PythonCodeGenerator {
     private static class ExpressionGenerator implements ExpressionNodeVisitor<PythonExpressionNode> {
         @Override
         public PythonExpressionNode visit(LiteralNode literal) {
-            return literal.getValue().accept(new PrimitiveValueVisitor<PythonExpressionNode>() {
-                @Override
-                public PythonExpressionNode visit(IntegerValue value) {
-                    return pythonIntegerLiteral(value.getValue());
-                }
-
-                @Override
-                public PythonExpressionNode visit(StringValue value) {
-                    return pythonStringLiteral(value.getValue());
-                }
-
-                @Override
-                public PythonExpressionNode visit(BooleanValue value) {
-                    throw new UnsupportedOperationException();
-                }
-
-                @Override
-                public PythonExpressionNode visit(UnitValue unitValue) {
-                    throw new UnsupportedOperationException();
-                }
-            });
+            return generateCode(literal.getValue());
         }
 
         @Override
-        public PythonExpressionNode visit(
-            VariableReferenceNode variableReference) {
-            throw new UnsupportedOperationException();
+        public PythonExpressionNode visit(VariableReferenceNode variableReference) {
+            return pythonVariableReference(variableReference.getReferent().getName());
         }
 
         @Override
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 c6cf516..4893cb6 100644
--- a/src/main/java/org/zwobble/couscous/backends/python/PythonSerializer.java
+++ b/src/main/java/org/zwobble/couscous/backends/python/PythonSerializer.java
@@ -10,6 +10,7 @@ import org.zwobble.couscous.backends.python.ast.PythonNode;
 import org.zwobble.couscous.backends.python.ast.PythonPassNode;
 import org.zwobble.couscous.backends.python.ast.PythonReturnNode;
 import org.zwobble.couscous.backends.python.ast.PythonStringLiteralNode;
+import org.zwobble.couscous.backends.python.ast.PythonVariableReferenceNode;
 import org.zwobble.couscous.backends.python.ast.PythonWriter;
 import org.zwobble.couscous.backends.python.ast.visitors.PythonNodeVisitor;
 
@@ -54,6 +55,12 @@ public class PythonSerializer implements PythonNodeVisitor<Void> {
     }
 
     @Override
+    public Void visit(PythonVariableReferenceNode reference) {
+        writer.writeIdentifier(reference.getName());
+        return null;
+    }
+
+    @Override
     public Void visit(PythonReturnNode pythonReturn) {
         writer.startStatement();
         writer.writeKeyword("return");
diff --git a/src/main/java/org/zwobble/couscous/backends/python/ast/PythonVariableReferenceNode.java b/src/main/java/org/zwobble/couscous/backends/python/ast/PythonVariableReferenceNode.java
new file mode 100644
index 0000000..81f26d2
--- /dev/null
+++ b/src/main/java/org/zwobble/couscous/backends/python/ast/PythonVariableReferenceNode.java
@@ -0,0 +1,15 @@
+package org.zwobble.couscous.backends.python.ast;
+
+import org.zwobble.couscous.backends.python.ast.visitors.PythonNodeVisitor;
+
+import lombok.Value;
+
+@Value(staticConstructor="pythonVariableReference")
+public class PythonVariableReferenceNode implements PythonExpressionNode {
+    String name;
+
+    @Override
+    public <T> T accept(PythonNodeVisitor<T> visitor) {
+        return visitor.visit(this);
+    }
+}
diff --git a/src/main/java/org/zwobble/couscous/backends/python/ast/visitors/PythonNodeVisitor.java b/src/main/java/org/zwobble/couscous/backends/python/ast/visitors/PythonNodeVisitor.java
index 42f97a0..88035c8 100644
--- a/src/main/java/org/zwobble/couscous/backends/python/ast/visitors/PythonNodeVisitor.java
+++ b/src/main/java/org/zwobble/couscous/backends/python/ast/visitors/PythonNodeVisitor.java
@@ -8,11 +8,13 @@ import org.zwobble.couscous.backends.python.ast.PythonModuleNode;
 import org.zwobble.couscous.backends.python.ast.PythonPassNode;
 import org.zwobble.couscous.backends.python.ast.PythonReturnNode;
 import org.zwobble.couscous.backends.python.ast.PythonStringLiteralNode;
+import org.zwobble.couscous.backends.python.ast.PythonVariableReferenceNode;
 
 public interface PythonNodeVisitor<T> {
     T visit(PythonIntegerLiteralNode integerLiteral);
     T visit(PythonStringLiteralNode stringLiteral);
     T visit(PythonBooleanLiteralNode booleanLiteral);
+    T visit(PythonVariableReferenceNode reference);
     
     T visit(PythonReturnNode pythonReturn);
     T visit(PythonPassNode pass);