2 files changed, 66 insertions, 31 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
new file mode 100644
index 0000000..129fb19
--- /dev/null
+++ b/src/main/java/org/zwobble/couscous/backends/python/PrimitiveMethods.java
@@ -0,0 +1,56 @@
+package org.zwobble.couscous.backends.python;
+
+import java.util.List;
+import java.util.Map;
+import java.util.Optional;
+
+import org.zwobble.couscous.backends.python.ast.PythonExpressionNode;
+import org.zwobble.couscous.values.StringValue;
+import org.zwobble.couscous.values.TypeReference;
+
+import com.google.common.collect.ImmutableMap;
+
+import static java.util.Arrays.asList;
+import static org.zwobble.couscous.backends.python.ast.PythonAttributeAccessNode.pythonAttributeAccess;
+import static org.zwobble.couscous.backends.python.ast.PythonCallNode.pythonCall;
+import static org.zwobble.couscous.backends.python.ast.PythonVariableReferenceNode.pythonVariableReference;
+
+public class PrimitiveMethods {
+    private static final Map<String, PrimitiveMethodGenerator> STRING_METHODS =
+        ImmutableMap.<String, PrimitiveMethodGenerator>builder()
+        
+            .put("length", (receiver, arguments) ->
+                pythonCall(
+                    pythonVariableReference("len"),
+                    asList(receiver)))
+            
+            .put("substring", (receiver, arguments) ->
+                pythonCall(
+                    pythonAttributeAccess(
+                        receiver,
+                        "__getitem__"),
+                    asList(pythonCall(
+                        pythonVariableReference("slice"),
+                    arguments))))
+            
+            .build();
+    
+    private static final Map<TypeReference, Map<String, PrimitiveMethodGenerator>> METHODS = 
+        ImmutableMap.<TypeReference, Map<String, PrimitiveMethodGenerator>>builder()
+            .put(StringValue.REF, STRING_METHODS)
+            .build();
+    
+    public static boolean isPrimitive(TypeReference type) {
+        return METHODS.containsKey(type);
+    }
+    
+    public static Optional<PrimitiveMethodGenerator> getPrimitiveMethod(TypeReference type, String methodName) {
+        return Optional.ofNullable(METHODS.get(type))
+            .flatMap(methodsForType -> Optional.ofNullable(methodsForType.get(methodName)));
+    }
+    
+    @FunctionalInterface
+    public interface PrimitiveMethodGenerator {
+        PythonExpressionNode generate(PythonExpressionNode receiver, List<PythonExpressionNode> arguments);
+    }
+}
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 067e826..cf48ba0 100644
--- a/src/main/java/org/zwobble/couscous/backends/python/PythonCodeGenerator.java
+++ b/src/main/java/org/zwobble/couscous/backends/python/PythonCodeGenerator.java
@@ -157,38 +157,21 @@ public class PythonCodeGenerator {
 
         @Override
         public PythonExpressionNode visit(MethodCallNode methodCall) {
+            val receiver = generateExpression(methodCall.getReceiver());
+            val arguments = generateArguments(methodCall);
             if (isPrimitive(methodCall.getReceiver())) {
-                val isStringLength = methodCall.getReceiver().getType().equals(StringValue.REF) &&
-                    methodCall.getMethodName() == "length";
-                val isSubstring = methodCall.getReceiver().getType().equals(StringValue.REF) &&
-                    methodCall.getMethodName() == "substring";
-                if (isStringLength) {
-                    return pythonCall(
-                        pythonVariableReference("len"),
-                        asList(generateExpression(methodCall.getReceiver())));
-                } else if (isSubstring) {
-                    return pythonCall(
-                        pythonAttributeAccess(
-                            generateExpression(methodCall.getReceiver()),
-                            "__getitem__"),
-                        asList(pythonCall(
-                            pythonVariableReference("slice"),
-                            generateArguments(methodCall))));
-                } else {
-                    throw new UnsupportedOperationException();
-                }
+                val primitiveMethodGenerator = PrimitiveMethods.getPrimitiveMethod(
+                    methodCall.getReceiver().getType(),
+                    methodCall.getMethodName()).get();
+                return primitiveMethodGenerator.generate(receiver, arguments);
             } else {
-                val arguments = generateArguments(methodCall);
                 return pythonCall(
-                    pythonAttributeAccess(
-                        generateExpression(methodCall.getReceiver()),
-                        methodCall.getMethodName()),
+                    pythonAttributeAccess(receiver, methodCall.getMethodName()),
                     arguments);
             }
         }
 
-        private List<PythonExpressionNode> generateArguments(
-            MethodCallNode methodCall) {
+        private List<PythonExpressionNode> generateArguments(MethodCallNode methodCall) {
             return methodCall.getArguments().stream()
                 .map(PythonCodeGenerator::generateExpression)
                 .collect(Collectors.toList());
@@ -199,12 +182,8 @@ public class PythonCodeGenerator {
             throw new UnsupportedOperationException();
         }
     }
-
+    
     private static boolean isPrimitive(ExpressionNode value) {
-        val type = value.getType();
-        return type.equals(IntegerValue.REF) || 
-            type.equals(StringValue.REF) ||
-            type.equals(BooleanValue.REF) ||
-            type.equals(UnitValue.REF);
+        return PrimitiveMethods.isPrimitive(value.getType());
     }
 }