Support assignment statements in the Python backend

11 files changed, 89 insertions, 36 deletions

diff --git a/src/main/java/org/zwobble/couscous/ast/Assignment.java b/src/main/java/org/zwobble/couscous/ast/AssignmentNode.java
index d8bdaaf..861e847 100644
--- a/src/main/java/org/zwobble/couscous/ast/Assignment.java
+++ b/src/main/java/org/zwobble/couscous/ast/AssignmentNode.java
@@ -7,9 +7,9 @@ import static org.zwobble.couscous.ast.VariableReferenceNode.reference;
 import lombok.Value;
 
 @Value
-public class Assignment implements ExpressionNode {
-    public static Assignment assign(VariableNode target, ExpressionNode value) {
-        return new Assignment(reference(target), value);
+public class AssignmentNode implements ExpressionNode {
+    public static AssignmentNode assign(VariableNode target, ExpressionNode value) {
+        return new AssignmentNode(reference(target), value);
     }
     
     VariableReferenceNode target;
diff --git a/src/main/java/org/zwobble/couscous/ast/visitors/ExpressionNodeVisitor.java b/src/main/java/org/zwobble/couscous/ast/visitors/ExpressionNodeVisitor.java
index 5d3ea9b..c47b942 100644
--- a/src/main/java/org/zwobble/couscous/ast/visitors/ExpressionNodeVisitor.java
+++ b/src/main/java/org/zwobble/couscous/ast/visitors/ExpressionNodeVisitor.java
@@ -1,6 +1,6 @@
 package org.zwobble.couscous.ast.visitors;
 
-import org.zwobble.couscous.ast.Assignment;
+import org.zwobble.couscous.ast.AssignmentNode;
 import org.zwobble.couscous.ast.LiteralNode;
 import org.zwobble.couscous.ast.MethodCallNode;
 import org.zwobble.couscous.ast.StaticMethodCallNode;
@@ -10,7 +10,7 @@ import org.zwobble.couscous.ast.VariableReferenceNode;
 public interface ExpressionNodeVisitor<T> {
     T visit(LiteralNode literal);
     T visit(VariableReferenceNode variableReference);
-    T visit(Assignment assignment);
+    T visit(AssignmentNode assignment);
     T visit(TernaryConditionalNode ternaryConditional);
     T visit(MethodCallNode methodCall);
     T visit(StaticMethodCallNode staticMethodCall);
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);
 
diff --git a/src/main/java/org/zwobble/couscous/interpreter/Evaluator.java b/src/main/java/org/zwobble/couscous/interpreter/Evaluator.java
index 9362698..c2eeb59 100644
--- a/src/main/java/org/zwobble/couscous/interpreter/Evaluator.java
+++ b/src/main/java/org/zwobble/couscous/interpreter/Evaluator.java
@@ -3,7 +3,7 @@ package org.zwobble.couscous.interpreter;
 import java.util.List;
 import java.util.stream.Collectors;
 
-import org.zwobble.couscous.ast.Assignment;
+import org.zwobble.couscous.ast.AssignmentNode;
 import org.zwobble.couscous.ast.ExpressionNode;
 import org.zwobble.couscous.ast.LiteralNode;
 import org.zwobble.couscous.ast.MethodCallNode;
@@ -43,7 +43,7 @@ public class Evaluator implements ExpressionNodeVisitor<InterpreterValue> {
     }
 
     @Override
-    public InterpreterValue visit(Assignment assignment) {
+    public InterpreterValue visit(AssignmentNode assignment) {
         val value = eval(assignment.getValue());
         environment.put(assignment.getTarget().getReferentId(), value);
         return value;
diff --git a/src/test/java/org/zwobble/couscous/tests/BackendTests.java b/src/test/java/org/zwobble/couscous/tests/BackendTests.java
index 7e7c3bd..0f3956d 100644
--- a/src/test/java/org/zwobble/couscous/tests/BackendTests.java
+++ b/src/test/java/org/zwobble/couscous/tests/BackendTests.java
@@ -1,7 +1,10 @@
 package org.zwobble.couscous.tests;
 
 import org.junit.Test;
+import org.zwobble.couscous.ast.AssignmentNode;
 import org.zwobble.couscous.ast.ClassNode;
+import org.zwobble.couscous.ast.ExpressionStatementNode;
+import org.zwobble.couscous.ast.LiteralNode;
 import org.zwobble.couscous.ast.MethodNode;
 import org.zwobble.couscous.ast.ReturnNode;
 import org.zwobble.couscous.values.PrimitiveValue;
@@ -46,6 +49,18 @@ public abstract class BackendTests {
         
         assertEquals(value("hello, world!"), result);
     }
+    
+    @Test
+    public void canReassignValueToArgument() {
+        val arg = formalArg(var(42, "x", StringValue.REF));
+        val method = staticMethod("hello")
+            .argument(arg)
+            .statement(new ExpressionStatementNode(new AssignmentNode(reference(arg), LiteralNode.literal("[updated value]"))))
+            .statement(new ReturnNode(reference(arg)));
+        val result = runMethod(method, value("[initial value]"));
+        
+        assertEquals(value("[updated value]"), result);
+    }
 
     protected PrimitiveValue runMethod(MethodNode.MethodNodeBuilder methodBuilder, PrimitiveValue... arguments) {
         val method = methodBuilder.build();
diff --git a/src/test/java/org/zwobble/couscous/tests/backends/python/PythonSerializerTests.java b/src/test/java/org/zwobble/couscous/tests/backends/python/PythonSerializerTests.java
index 944c2af..2835034 100644
--- a/src/test/java/org/zwobble/couscous/tests/backends/python/PythonSerializerTests.java
+++ b/src/test/java/org/zwobble/couscous/tests/backends/python/PythonSerializerTests.java
@@ -7,6 +7,7 @@ import org.zwobble.couscous.backends.python.ast.PythonFunctionDefinitionNode;
 import static java.util.Arrays.asList;
 import static org.junit.Assert.assertEquals;
 import static org.zwobble.couscous.backends.python.PythonSerializer.serialize;
+import static org.zwobble.couscous.backends.python.ast.PythonAssignmentNode.pythonAssignment;
 import static org.zwobble.couscous.backends.python.ast.PythonBooleanLiteralNode.pythonBooleanLiteral;
 import static org.zwobble.couscous.backends.python.ast.PythonIntegerLiteralNode.pythonIntegerLiteral;
 import static org.zwobble.couscous.backends.python.ast.PythonModuleNode.pythonModule;
@@ -80,13 +81,6 @@ public class PythonSerializerTests {
     }
     
     @Test
-    public void moduleIsSerializedStatements() {
-        val classNode = pythonModule(asList(PASS));
-        val output = serialize(classNode);
-        assertEquals("pass\n", output);
-    }
-    
-    @Test
     public void bodiesOfBlocksAreIndented() {
         val classNode = PythonClassNode.builder("Foo")
             .statement(PythonFunctionDefinitionNode.builder("one").build())
@@ -101,4 +95,18 @@ public class PythonSerializerTests {
             "        pass\n";
         assertEquals(expectedOutput, output);
     }
+    
+    @Test
+    public void assignmentIsSerializedWithEqualsSymbolSeparatingTargetAndValue() {
+        val output = serialize(pythonAssignment(
+            pythonVariableReference("x"), pythonIntegerLiteral(42)));
+        assertEquals("x = 42\n", output);
+    }
+    
+    @Test
+    public void moduleIsSerializedStatements() {
+        val classNode = pythonModule(asList(PASS));
+        val output = serialize(classNode);
+        assertEquals("pass\n", output);
+    }
 }
diff --git a/src/test/java/org/zwobble/couscous/tests/interpreter/EvaluatorTests.java b/src/test/java/org/zwobble/couscous/tests/interpreter/EvaluatorTests.java
index 3ed1b6d..40cdf8f 100644
--- a/src/test/java/org/zwobble/couscous/tests/interpreter/EvaluatorTests.java
+++ b/src/test/java/org/zwobble/couscous/tests/interpreter/EvaluatorTests.java
@@ -22,7 +22,7 @@ import org.zwobble.couscous.values.StringValue;
 import com.google.common.collect.ImmutableMap;
 
 import static org.junit.Assert.assertEquals;
-import static org.zwobble.couscous.ast.Assignment.assign;
+import static org.zwobble.couscous.ast.AssignmentNode.assign;
 import static org.zwobble.couscous.ast.FormalArgumentNode.formalArg;
 import static org.zwobble.couscous.ast.LiteralNode.literal;
 import static org.zwobble.couscous.ast.MethodCallNode.methodCall;
diff --git a/src/test/java/org/zwobble/couscous/tests/interpreter/InterpreterTests.java b/src/test/java/org/zwobble/couscous/tests/interpreter/InterpreterTests.java
index c6e6ce2..a54f150 100644
--- a/src/test/java/org/zwobble/couscous/tests/interpreter/InterpreterTests.java
+++ b/src/test/java/org/zwobble/couscous/tests/interpreter/InterpreterTests.java
@@ -5,7 +5,7 @@ import java.util.stream.Collectors;
 
 import org.junit.Test;
 import org.zwobble.couscous.MapBackedProject;
-import org.zwobble.couscous.ast.Assignment;
+import org.zwobble.couscous.ast.AssignmentNode;
 import org.zwobble.couscous.ast.ClassNode;
 import org.zwobble.couscous.ast.ExpressionStatementNode;
 import org.zwobble.couscous.ast.LiteralNode;
@@ -49,23 +49,11 @@ public class InterpreterTests extends BackendTests {
     }
     
     @Test
-    public void canReassignValueToArgument() {
-        val arg = formalArg(var(42, "x", StringValue.REF));
-        val method = staticMethod("hello")
-            .argument(arg)
-            .statement(new ExpressionStatementNode(new Assignment(reference(arg), LiteralNode.literal("[updated value]"))))
-            .statement(new ReturnNode(reference(arg)));
-        val result = runMethod(method, value("[initial value]"));
-        
-        assertEquals(value("[updated value]"), result);
-    }
-    
-    @Test
     public void errorIfTryingToAssignValueOfWrongTypeToVariable() {
         val arg = formalArg(var(42, "x", StringValue.REF));
         val method = staticMethod("hello")
             .argument(arg)
-            .statement(new ExpressionStatementNode(new Assignment(reference(arg), literal(0))));
+            .statement(new ExpressionStatementNode(new AssignmentNode(reference(arg), literal(0))));
 
         val exception = assertThrows(UnexpectedValueType.class,
             () -> runMethod(method, value("")));
@@ -91,7 +79,7 @@ public class InterpreterTests extends BackendTests {
             42, "x", StringValue.REF, LiteralNode.literal("[initial value]"));
         val method = staticMethod("hello")
             .statement(localVariableDeclaration)
-            .statement(new ExpressionStatementNode(new Assignment(reference(localVariableDeclaration), LiteralNode.literal("[updated value]"))))
+            .statement(new ExpressionStatementNode(new AssignmentNode(reference(localVariableDeclaration), LiteralNode.literal("[updated value]"))))
             .statement(new ReturnNode(reference(localVariableDeclaration)));
         val result = runMethod(method);
         
@@ -103,7 +91,7 @@ public class InterpreterTests extends BackendTests {
         val localVariableDeclaration = localVariableDeclaration(
             42, "x", StringValue.REF, literal(""));
         val method = staticMethod("hello")
-            .statement(new ExpressionStatementNode(new Assignment(reference(localVariableDeclaration), LiteralNode.literal("[updated value]"))));
+            .statement(new ExpressionStatementNode(new AssignmentNode(reference(localVariableDeclaration), LiteralNode.literal("[updated value]"))));
 
         val exception = assertThrows(VariableNotInScope.class,
             () -> runMethod(method));