Support assignment statements in the Python backend

4 files changed, 47 insertions, 5 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 63fa56b..c4b1f49 100644
--- a/src/main/java/org/zwobble/couscous/backends/python/PythonCodeGenerator.java
+++ b/src/main/java/org/zwobble/couscous/backends/python/PythonCodeGenerator.java
@@ -2,7 +2,7 @@ package org.zwobble.couscous.backends.python;
 
 import java.util.stream.Collectors;
 
-import org.zwobble.couscous.ast.Assignment;
+import org.zwobble.couscous.ast.AssignmentNode;
 import org.zwobble.couscous.ast.ClassNode;
 import org.zwobble.couscous.ast.ExpressionNode;
 import org.zwobble.couscous.ast.ExpressionStatementNode;
@@ -22,6 +22,7 @@ import org.zwobble.couscous.backends.python.ast.PythonExpressionNode;
 import org.zwobble.couscous.backends.python.ast.PythonFunctionDefinitionNode;
 import org.zwobble.couscous.backends.python.ast.PythonModuleNode;
 import org.zwobble.couscous.backends.python.ast.PythonStatementNode;
+import org.zwobble.couscous.backends.python.ast.PythonVariableReferenceNode;
 import org.zwobble.couscous.values.BooleanValue;
 import org.zwobble.couscous.values.IntegerValue;
 import org.zwobble.couscous.values.PrimitiveValue;
@@ -30,6 +31,7 @@ import org.zwobble.couscous.values.StringValue;
 import org.zwobble.couscous.values.UnitValue;
 
 import static java.util.Arrays.asList;
+import static org.zwobble.couscous.backends.python.ast.PythonAssignmentNode.pythonAssignment;
 import static org.zwobble.couscous.backends.python.ast.PythonClassNode.pythonClass;
 import static org.zwobble.couscous.backends.python.ast.PythonFunctionDefinitionNode.pythonFunctionDefinition;
 import static org.zwobble.couscous.backends.python.ast.PythonIntegerLiteralNode.pythonIntegerLiteral;
@@ -100,7 +102,14 @@ public class PythonCodeGenerator {
 
         @Override
         public PythonStatementNode visit(ExpressionStatementNode expressionStatement) {
-            throw new UnsupportedOperationException();
+            if (expressionStatement.getExpression() instanceof AssignmentNode) {
+                val assignment = (AssignmentNode) expressionStatement.getExpression();
+                return pythonAssignment(
+                    EXPRESSION_GENERATOR.visit(assignment.getTarget()),
+                    generateExpression(assignment.getValue()));
+            } else {
+                throw new UnsupportedOperationException();   
+            }
         }
 
         @Override
@@ -110,9 +119,11 @@ public class PythonCodeGenerator {
     }
     
     private static PythonExpressionNode generateExpression(ExpressionNode expression) {
-        return expression.accept(new ExpressionGenerator());
+        return expression.accept(EXPRESSION_GENERATOR);
     }
     
+    private static final ExpressionGenerator EXPRESSION_GENERATOR = new ExpressionGenerator();
+    
     private static class ExpressionGenerator implements ExpressionNodeVisitor<PythonExpressionNode> {
         @Override
         public PythonExpressionNode visit(LiteralNode literal) {
@@ -120,12 +131,12 @@ public class PythonCodeGenerator {
         }
 
         @Override
-        public PythonExpressionNode visit(VariableReferenceNode variableReference) {
+        public PythonVariableReferenceNode visit(VariableReferenceNode variableReference) {
             return pythonVariableReference(variableReference.getReferent().getName());
         }
 
         @Override
-        public PythonExpressionNode visit(Assignment assignment) {
+        public PythonExpressionNode visit(AssignmentNode assignment) {
             throw new UnsupportedOperationException();
         }
 
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 4893cb6..bc5fb19 100644
--- a/src/main/java/org/zwobble/couscous/backends/python/PythonSerializer.java
+++ b/src/main/java/org/zwobble/couscous/backends/python/PythonSerializer.java
@@ -1,5 +1,6 @@
 package org.zwobble.couscous.backends.python;
 
+import org.zwobble.couscous.backends.python.ast.PythonAssignmentNode;
 import org.zwobble.couscous.backends.python.ast.PythonBlock;
 import org.zwobble.couscous.backends.python.ast.PythonBooleanLiteralNode;
 import org.zwobble.couscous.backends.python.ast.PythonClassNode;
@@ -106,6 +107,18 @@ public class PythonSerializer implements PythonNodeVisitor<Void> {
         return null;
     }
 
+    @Override
+    public Void visit(PythonAssignmentNode assignment) {
+        writer.startStatement();
+        write(assignment.getTarget());
+        writer.writeSpace();
+        writer.writeSymbol("=");
+        writer.writeSpace();
+        write(assignment.getValue());
+        writer.endStatement();
+        return null;
+    }
+
     private void writeArgumentNames(PythonFunctionDefinitionNode functionDefinition) {
         if (!functionDefinition.getArgumentNames().isEmpty()) {
             writer.writeIdentifier(functionDefinition.getArgumentNames().get(0));
diff --git a/src/main/java/org/zwobble/couscous/backends/python/ast/PythonAssignmentNode.java b/src/main/java/org/zwobble/couscous/backends/python/ast/PythonAssignmentNode.java
new file mode 100644
index 0000000..d64e101
--- /dev/null
+++ b/src/main/java/org/zwobble/couscous/backends/python/ast/PythonAssignmentNode.java
@@ -0,0 +1,16 @@
+package org.zwobble.couscous.backends.python.ast;
+
+import org.zwobble.couscous.backends.python.ast.visitors.PythonNodeVisitor;
+
+import lombok.Value;
+
+@Value(staticConstructor="pythonAssignment")
+public class PythonAssignmentNode implements PythonStatementNode {
+    PythonVariableReferenceNode target;
+    PythonExpressionNode value;
+    
+    @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 88035c8..f5cd04f 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
@@ -1,5 +1,6 @@
 package org.zwobble.couscous.backends.python.ast.visitors;
 
+import org.zwobble.couscous.backends.python.ast.PythonAssignmentNode;
 import org.zwobble.couscous.backends.python.ast.PythonBooleanLiteralNode;
 import org.zwobble.couscous.backends.python.ast.PythonClassNode;
 import org.zwobble.couscous.backends.python.ast.PythonFunctionDefinitionNode;
@@ -20,6 +21,7 @@ public interface PythonNodeVisitor<T> {
     T visit(PythonPassNode pass);
     T visit(PythonClassNode pythonClass);
     T visit(PythonFunctionDefinitionNode functionDefinition);
+    T visit(PythonAssignmentNode assignment);
 
     T visit(PythonModuleNode module);