Support ternary conditionals in Python backend

4 files changed, 48 insertions, 2 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 b1947a8..0db2be6 100644
--- a/src/main/java/org/zwobble/couscous/backends/python/PythonCodeGenerator.java
+++ b/src/main/java/org/zwobble/couscous/backends/python/PythonCodeGenerator.java
@@ -32,7 +32,9 @@ 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.PythonBooleanLiteralNode.pythonBooleanLiteral;
 import static org.zwobble.couscous.backends.python.ast.PythonClassNode.pythonClass;
+import static org.zwobble.couscous.backends.python.ast.PythonConditionalExpressionNode.pythonConditionalExpression;
 import static org.zwobble.couscous.backends.python.ast.PythonFunctionDefinitionNode.pythonFunctionDefinition;
 import static org.zwobble.couscous.backends.python.ast.PythonIntegerLiteralNode.pythonIntegerLiteral;
 import static org.zwobble.couscous.backends.python.ast.PythonModuleNode.pythonModule;
@@ -67,7 +69,7 @@ public class PythonCodeGenerator {
 
             @Override
             public PythonExpressionNode visit(BooleanValue value) {
-                throw new UnsupportedOperationException();
+                return pythonBooleanLiteral(value.getValue());
             }
 
             @Override
@@ -144,7 +146,10 @@ public class PythonCodeGenerator {
 
         @Override
         public PythonExpressionNode visit(TernaryConditionalNode ternaryConditional) {
-            throw new UnsupportedOperationException();
+            return pythonConditionalExpression(
+                generateExpression(ternaryConditional.getCondition()),
+                generateExpression(ternaryConditional.getIfTrue()),
+                generateExpression(ternaryConditional.getIfFalse()));
         }
 
         @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 bc5fb19..728b71a 100644
--- a/src/main/java/org/zwobble/couscous/backends/python/PythonSerializer.java
+++ b/src/main/java/org/zwobble/couscous/backends/python/PythonSerializer.java
@@ -4,6 +4,8 @@ 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;
+import org.zwobble.couscous.backends.python.ast.PythonConditionalExpressionNode;
+import org.zwobble.couscous.backends.python.ast.PythonExpressionNode;
 import org.zwobble.couscous.backends.python.ast.PythonFunctionDefinitionNode;
 import org.zwobble.couscous.backends.python.ast.PythonIntegerLiteralNode;
 import org.zwobble.couscous.backends.python.ast.PythonModuleNode;
@@ -62,6 +64,20 @@ public class PythonSerializer implements PythonNodeVisitor<Void> {
     }
 
     @Override
+    public Void visit(PythonConditionalExpressionNode conditional) {
+        writeParenthesised(conditional.getTrueValue());
+        writer.writeSpace();
+        writer.writeKeyword("if");
+        writer.writeSpace();
+        writeParenthesised(conditional.getCondition());
+        writer.writeSpace();
+        writer.writeKeyword("else");
+        writer.writeSpace();
+        writeParenthesised(conditional.getFalseValue());
+        return null;
+    }
+
+    @Override
     public Void visit(PythonReturnNode pythonReturn) {
         writer.startStatement();
         writer.writeKeyword("return");
@@ -119,6 +135,12 @@ public class PythonSerializer implements PythonNodeVisitor<Void> {
         return null;
     }
 
+    private void writeParenthesised(PythonExpressionNode expression) {
+        writer.writeSymbol("(");
+        write(expression);
+        writer.writeSymbol(")");
+    }
+
     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/PythonConditionalExpressionNode.java b/src/main/java/org/zwobble/couscous/backends/python/ast/PythonConditionalExpressionNode.java
new file mode 100644
index 0000000..9fd406c
--- /dev/null
+++ b/src/main/java/org/zwobble/couscous/backends/python/ast/PythonConditionalExpressionNode.java
@@ -0,0 +1,17 @@
+package org.zwobble.couscous.backends.python.ast;
+
+import org.zwobble.couscous.backends.python.ast.visitors.PythonNodeVisitor;
+
+import lombok.Value;
+
+@Value(staticConstructor="pythonConditionalExpression")
+public class PythonConditionalExpressionNode implements PythonExpressionNode {
+    PythonExpressionNode condition;
+    PythonExpressionNode trueValue;
+    PythonExpressionNode falseValue;
+    
+    @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 f5cd04f..1f5b1a2 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
@@ -3,6 +3,7 @@ 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.PythonConditionalExpressionNode;
 import org.zwobble.couscous.backends.python.ast.PythonFunctionDefinitionNode;
 import org.zwobble.couscous.backends.python.ast.PythonIntegerLiteralNode;
 import org.zwobble.couscous.backends.python.ast.PythonModuleNode;
@@ -16,6 +17,7 @@ public interface PythonNodeVisitor<T> {
     T visit(PythonStringLiteralNode stringLiteral);
     T visit(PythonBooleanLiteralNode booleanLiteral);
     T visit(PythonVariableReferenceNode reference);
+    T visit(PythonConditionalExpressionNode conditional);
     
     T visit(PythonReturnNode pythonReturn);
     T visit(PythonPassNode pass);