Python backend: import class into module

6 files changed, 49 insertions, 22 deletions

diff --git a/src/main/java/org/zwobble/couscous/ast/TypeName.java b/src/main/java/org/zwobble/couscous/ast/TypeName.java
index ce80f82..cbc0fe0 100644
--- a/src/main/java/org/zwobble/couscous/ast/TypeName.java
+++ b/src/main/java/org/zwobble/couscous/ast/TypeName.java
@@ -8,11 +8,12 @@ import lombok.Value;
 @AllArgsConstructor(access=AccessLevel.PRIVATE)
 public class TypeName {
     public static TypeName of(String qualifiedName) {
-        return new TypeName(
-            qualifiedName,
-            qualifiedName.substring(qualifiedName.lastIndexOf(".") + 1));
+        return new TypeName(qualifiedName);
     }
     
     String qualifiedName;
-    String simpleName;
+    
+    public String getSimpleName() {
+        return qualifiedName.substring(qualifiedName.lastIndexOf(".") + 1);
+    }
 }
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 4f54d45..8b26b10 100644
--- a/src/main/java/org/zwobble/couscous/backends/python/PythonCodeGenerator.java
+++ b/src/main/java/org/zwobble/couscous/backends/python/PythonCodeGenerator.java
@@ -2,12 +2,10 @@ package org.zwobble.couscous.backends.python;
 
 import java.util.List;
 import java.util.Set;
-import java.util.regex.Pattern;
 import java.util.stream.Collectors;
 import java.util.stream.Stream;
 
 import org.zwobble.couscous.ast.AssignmentNode;
-import org.zwobble.couscous.ast.TypeName;
 import org.zwobble.couscous.ast.ClassNode;
 import org.zwobble.couscous.ast.ExpressionNode;
 import org.zwobble.couscous.ast.ExpressionStatementNode;
@@ -19,6 +17,7 @@ import org.zwobble.couscous.ast.ReturnNode;
 import org.zwobble.couscous.ast.StatementNode;
 import org.zwobble.couscous.ast.StaticMethodCallNode;
 import org.zwobble.couscous.ast.TernaryConditionalNode;
+import org.zwobble.couscous.ast.TypeName;
 import org.zwobble.couscous.ast.VariableReferenceNode;
 import org.zwobble.couscous.ast.visitors.ExpressionNodeMapper;
 import org.zwobble.couscous.ast.visitors.NodeVisitorWithEmptyDefaults;
@@ -51,13 +50,13 @@ import static org.zwobble.couscous.backends.python.ast.PythonCallNode.pythonCall
 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.PythonImportAliasNode.pythonImportAlias;
 import static org.zwobble.couscous.backends.python.ast.PythonImportNode.pythonImport;
 import static org.zwobble.couscous.backends.python.ast.PythonIntegerLiteralNode.pythonIntegerLiteral;
 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 static org.zwobble.couscous.util.ExtraLists.foldLeft;
 
 import lombok.val;
 
@@ -80,7 +79,9 @@ public class PythonCodeGenerator {
     private static Stream<PythonImportNode> generateImports(ClassNode classNode) {
         val classes = findReferencedClasses(classNode);
         return classes.stream()
-            .map(name -> pythonImport(name.getQualifiedName()));
+            .map(name -> pythonImport(
+                name.getQualifiedName(),
+                asList(pythonImportAlias(name.getSimpleName()))));
     }
     
     private static Set<TypeName> findReferencedClasses(ClassNode classNode) {
@@ -210,16 +211,7 @@ public class PythonCodeGenerator {
         @Override
         public PythonExpressionNode visit(StaticMethodCallNode staticMethodCall) {
             val className = staticMethodCall.getClassName();
-            val moduleParts = asList(className.getQualifiedName().split(Pattern.quote(".")));
-            
-            val module = foldLeft(
-                moduleParts,
-                part -> (PythonExpressionNode)pythonVariableReference(part),
-                (parent, part) -> pythonAttributeAccess(parent, part));
-            
-            val classReference = pythonAttributeAccess(
-                module,
-                className.getSimpleName());
+            val classReference = pythonVariableReference(className.getSimpleName());
 
             val methodReference = pythonAttributeAccess(
                 classReference,
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 9c82b5d..249d9fa 100644
--- a/src/main/java/org/zwobble/couscous/backends/python/PythonSerializer.java
+++ b/src/main/java/org/zwobble/couscous/backends/python/PythonSerializer.java
@@ -164,9 +164,21 @@ public class PythonSerializer implements PythonNodeVisitor {
     @Override
     public void visit(PythonImportNode importNode) {
         writer.writeStatement(() -> {
+            writer.writeKeyword("from");
+            writer.writeSpace();
+            writer.writeIdentifier(importNode.getModuleName());
+            writer.writeSpace();
             writer.writeKeyword("import");
             writer.writeSpace();
-            writer.writeIdentifier(importNode.getName());
+            writeWithSeparator(
+                importNode.getAliases(),
+                alias -> {
+                    writer.writeIdentifier(alias.getName());
+                },
+                () -> {
+                    writer.writeSymbol(",");
+                    writer.writeSpace();
+                });
         });
     }
 
diff --git a/src/main/java/org/zwobble/couscous/backends/python/ast/PythonImportAliasNode.java b/src/main/java/org/zwobble/couscous/backends/python/ast/PythonImportAliasNode.java
new file mode 100644
index 0000000..9fb531a
--- /dev/null
+++ b/src/main/java/org/zwobble/couscous/backends/python/ast/PythonImportAliasNode.java
@@ -0,0 +1,8 @@
+package org.zwobble.couscous.backends.python.ast;
+
+import lombok.Value;
+
+@Value(staticConstructor="pythonImportAlias")
+public class PythonImportAliasNode {
+    String name;
+}
diff --git a/src/main/java/org/zwobble/couscous/backends/python/ast/PythonImportNode.java b/src/main/java/org/zwobble/couscous/backends/python/ast/PythonImportNode.java
index 2d38004..c639829 100644
--- a/src/main/java/org/zwobble/couscous/backends/python/ast/PythonImportNode.java
+++ b/src/main/java/org/zwobble/couscous/backends/python/ast/PythonImportNode.java
@@ -1,12 +1,15 @@
 package org.zwobble.couscous.backends.python.ast;
 
+import java.util.List;
+
 import org.zwobble.couscous.backends.python.ast.visitors.PythonNodeVisitor;
 
 import lombok.Value;
 
 @Value(staticConstructor="pythonImport")
 public class PythonImportNode implements PythonStatementNode {
-    String name;
+    String moduleName;
+    List<PythonImportAliasNode> aliases;
     
     @Override
     public void accept(PythonNodeVisitor visitor) {
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 ecd5d7c..46e0aa1 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
@@ -13,6 +13,7 @@ import static org.zwobble.couscous.backends.python.ast.PythonBooleanLiteralNode.
 import static org.zwobble.couscous.backends.python.ast.PythonCallNode.pythonCall;
 import static org.zwobble.couscous.backends.python.ast.PythonConditionalExpressionNode.pythonConditionalExpression;
 import static org.zwobble.couscous.backends.python.ast.PythonGetSliceNode.pythonGetSlice;
+import static org.zwobble.couscous.backends.python.ast.PythonImportAliasNode.pythonImportAlias;
 import static org.zwobble.couscous.backends.python.ast.PythonImportNode.pythonImport;
 import static org.zwobble.couscous.backends.python.ast.PythonIntegerLiteralNode.pythonIntegerLiteral;
 import static org.zwobble.couscous.backends.python.ast.PythonModuleNode.pythonModule;
@@ -147,8 +148,18 @@ public class PythonSerializerTests {
     
     @Test
     public void importIsSerializedUsingImportKeyword() {
-        val output = serialize(pythonImport("com.example.Program"));
-        assertEquals("import com.example.Program\n", output);
+        val output = serialize(pythonImport(
+            "com.example",
+            asList(pythonImportAlias("Program"))));
+        assertEquals("from com.example import Program\n", output);
+    }
+    
+    @Test
+    public void importAliasesAreSeparatedByCommas() {
+        val output = serialize(pythonImport(
+            "com.example",
+            asList(pythonImportAlias("Program"), pythonImportAlias("Runner"))));
+        assertEquals("from com.example import Program, Runner\n", output);
     }
     
     @Test