Use relative imports for Python backend

2 files changed, 31 insertions, 3 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 9b3367f..cc2b116 100644
--- a/src/main/java/org/zwobble/couscous/backends/python/PythonCodeGenerator.java
+++ b/src/main/java/org/zwobble/couscous/backends/python/PythonCodeGenerator.java
@@ -41,6 +41,7 @@ import org.zwobble.couscous.values.PrimitiveValueVisitor;
 import org.zwobble.couscous.values.StringValue;
 import org.zwobble.couscous.values.UnitValue;
 
+import com.google.common.base.Strings;
 import com.google.common.collect.ImmutableList;
 import com.google.common.collect.ImmutableSet;
 import com.google.common.collect.Iterables;
@@ -93,10 +94,21 @@ public class PythonCodeGenerator {
         val classes = findReferencedClasses(classNode);
         return classes.stream()
             .map(name -> pythonImport(
-                name.getQualifiedName(),
+                Strings.repeat(".", packageDepth(classNode) + 1) + name.getQualifiedName(),
                 asList(pythonImportAlias(name.getSimpleName()))));
     }
     
+    private static int packageDepth(ClassNode classNode) {
+        int depth = 0;
+        val qualifiedName = classNode.getName().getQualifiedName();
+        for (int index = 0; index < qualifiedName.length(); index++) {
+            if (qualifiedName.charAt(index) == '.') {
+                depth += 1;                
+            }
+        }
+        return depth;
+    }
+    
     private static Set<TypeName> findReferencedClasses(ClassNode classNode) {
         val imports = ImmutableSet.<TypeName>builder();
         NodeVisitors.visitAll(classNode, new NodeVisitorWithEmptyDefaults() {
diff --git a/src/main/java/org/zwobble/couscous/backends/python/PythonCompiler.java b/src/main/java/org/zwobble/couscous/backends/python/PythonCompiler.java
index 1fb9149..8787d81 100644
--- a/src/main/java/org/zwobble/couscous/backends/python/PythonCompiler.java
+++ b/src/main/java/org/zwobble/couscous/backends/python/PythonCompiler.java
@@ -1,6 +1,7 @@
 package org.zwobble.couscous.backends.python;
 
 import java.io.File;
+import java.io.IOException;
 import java.nio.file.Files;
 import java.nio.file.Path;
 import java.util.List;
@@ -17,9 +18,11 @@ import lombok.val;
 
 public class PythonCompiler {
     private final Path root;
+    private final String packageName;
 
-    public PythonCompiler(Path root) {
+    public PythonCompiler(Path root, String packageName) {
         this.root = root;
+        this.packageName = packageName;
     }
     
     public void compile(List<ClassNode> classes) {
@@ -34,7 +37,9 @@ public class PythonCompiler {
     }
 
     private Path pathForClass(TypeName className) {
-        return root.resolve(className.getQualifiedName().replace(".", File.separator) + ".py");
+        return root
+            .resolve(packageName)
+            .resolve(className.getQualifiedName().replace(".", File.separator) + ".py");
     }
 
     private void compileClass(ClassNode classNode) {
@@ -45,6 +50,17 @@ public class PythonCompiler {
     private void writeClass(TypeName name, String contents) {
         val path = pathForClass(name);
         Files.createDirectories(path.getParent());
+        createPythonPackages(path.getParent());
         Files.write(path, asList(contents));
     }
+
+    private void createPythonPackages(Path packagePath) throws IOException {
+        while (packagePath.startsWith(root)) {
+            val packageFile = packagePath.resolve("__init__.py");
+            if (!packageFile.toFile().exists()) {
+                Files.createFile(packageFile);
+            }
+            packagePath = packagePath.getParent();
+        }
+    }
 }