Implement divide and mod for Python backend

3 files changed, 33 insertions, 6 deletions

diff --git a/src/main/java/org/zwobble/couscous/backends/python/PrimitiveMethods.java b/src/main/java/org/zwobble/couscous/backends/python/PrimitiveMethods.java
index 1da18f9..5208ff1 100644
--- a/src/main/java/org/zwobble/couscous/backends/python/PrimitiveMethods.java
+++ b/src/main/java/org/zwobble/couscous/backends/python/PrimitiveMethods.java
@@ -40,8 +40,15 @@ public class PrimitiveMethods {
         addMethod(methods, "add", "__add__");
         addMethod(methods, "subtract", "__sub__");
         addMethod(methods, "multiply", "__mul__");
-        addMethod(methods, "divide", "__floordiv__");
-        addMethod(methods, "mod", "__mod__");
+
+        methods.put("divide", (receiver, arguments) ->
+            pythonCall(
+                internalReference("_div_round_to_zero"),
+                asList(receiver, arguments.get(0))));
+        methods.put("mod", (receiver, arguments) ->
+            pythonCall(
+                internalReference("_mod_round_to_zero"),
+                asList(receiver, arguments.get(0))));
 
         addMethod(methods, "greaterThan", "__gt__");
         addMethod(methods, "greaterThanOrEqual", "__ge__");
@@ -51,6 +58,10 @@ public class PrimitiveMethods {
         INT_METHODS = methods.build();
     }
     
+    private static PythonExpressionNode internalReference(String name) {
+        return pythonAttributeAccess(pythonVariableReference("_couscous"), name);
+    }
+    
     private static void addMethod(
             Builder<String, PrimitiveMethodGenerator> methods,
             String methodName,
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 a10c54d..b6f579c 100644
--- a/src/main/java/org/zwobble/couscous/backends/python/PythonCodeGenerator.java
+++ b/src/main/java/org/zwobble/couscous/backends/python/PythonCodeGenerator.java
@@ -67,17 +67,30 @@ import static org.zwobble.couscous.backends.python.ast.PythonVariableReferenceNo
 
 public class PythonCodeGenerator {
     public static PythonModuleNode generateCode(ClassNode classNode) {
-        Iterator<org.zwobble.couscous.backends.python.ast.PythonImportNode> imports = generateImports(classNode).iterator();
+        Iterator<PythonImportNode> imports = generateImports(classNode).iterator();
+        
+        PythonImportNode internalsImport = pythonImport(
+            importPathToRoot(classNode),
+            asList(pythonImportAlias("_couscous")));
+        
         PythonFunctionDefinitionNode constructor = generateConstructor(classNode.getConstructor());
         Iterable<PythonFunctionDefinitionNode> pythonMethods = Iterables.transform(classNode.getMethods(), PythonCodeGenerator::generateFunction);
         Iterable<PythonStatementNode> pythonBody = Iterables.concat(asList(constructor), pythonMethods);
         PythonClassNode pythonClass = pythonClass(classNode.getSimpleName(), ImmutableList.copyOf(pythonBody));
-        return pythonModule(ImmutableList.copyOf(Iterators.concat(imports, singletonIterator(pythonClass))));
+        
+        return pythonModule(ImmutableList.copyOf(Iterators.concat(
+            imports,
+            singletonIterator(internalsImport),
+            singletonIterator(pythonClass))));
     }
     
     private static Stream<PythonImportNode> generateImports(ClassNode classNode) {
-        final java.util.Set<org.zwobble.couscous.ast.TypeName> classes = findReferencedClasses(classNode);
-        return classes.stream().map(name -> pythonImport(Strings.repeat(".", packageDepth(classNode) + 1) + name.getQualifiedName(), asList(pythonImportAlias(name.getSimpleName()))));
+        Set<TypeName> classes = findReferencedClasses(classNode);
+        return classes.stream().map(name -> pythonImport(importPathToRoot(classNode) + name.getQualifiedName(), asList(pythonImportAlias(name.getSimpleName()))));
+    }
+
+    private static String importPathToRoot(ClassNode classNode) {
+        return Strings.repeat(".", packageDepth(classNode) + 1);
     }
     
     private static int packageDepth(ClassNode classNode) {
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 c5ea3e3..e325302 100644
--- a/src/main/java/org/zwobble/couscous/backends/python/PythonCompiler.java
+++ b/src/main/java/org/zwobble/couscous/backends/python/PythonCompiler.java
@@ -25,6 +25,9 @@ public class PythonCompiler {
             compileClass(classNode);
         }
         writeClass(TypeName.of("java.lang.Integer"), "class Integer(object):\n    def parseInt(value):\n        return int(value)");
+        writeClass(TypeName.of("_couscous"),
+            "def _div_round_to_zero(a, b): return -(-a // b) if (a < 0) ^ (b < 0) else a // b\n" +
+            "def _mod_round_to_zero(a, b): return -(-a % b) if (a < 0) ^ (b < 0) else a % b");
     }
     
     private Path pathForClass(TypeName className) {