Add tests for python interpreter/code_generator (#24)

m2cgen/exporters.py

@@ -43,8 +43,6 @@ def __init__(self, model, package_name=None, model_name="Model", indent=4):
 
 class PythonExporter(BaseExporter):
 
-    def __init__(self, model, model_name="Model", indent=4):
-        self.interpreter = interpreters.PythonInterpreter(
-            model_name=model_name,
-            indent=indent)
+    def __init__(self, model, indent=4):
+        self.interpreter = interpreters.PythonInterpreter(indent=indent)
         super(PythonExporter, self).__init__(model)

m2cgen/interpreters/interpreter.py

@@ -57,6 +57,8 @@ def interpret_feature_ref(self, expr):
             array_name=self._feature_array_name,
             index=expr.index)
 
+    # Private methods implementing visitor pattern
+
     def _do_interpret(self, expr, **kwargs):
         try:
             handler = self._select_handler(expr)

m2cgen/interpreters/python/code_generator.py

@@ -17,23 +17,13 @@ class PythonCodeGenerator(BaseCodeGenerator):
     tpl_var_declaration = CT("")
     tpl_block_termination = CT("")
 
-    def add_class_def(self, class_name):
-        class_def = "class " + class_name + "(object):"
-        self.add_code_line(class_def)
-        self.increase_indent()
-
     def add_method_def(self, name, args):
-        method_def = "def " + " " + name + "(self, "
+        method_def = "def " + " " + name + "("
         method_def += ", ".join(args)
         method_def += "):"
         self.add_code_line(method_def)
         self.increase_indent()
 
-    @contextlib.contextmanager
-    def class_definition(self, model_name):
-        self.add_class_def(model_name)
-        yield
-
     @contextlib.contextmanager
     def method_definition(self, name, args):
         self.add_method_def(name, args)

m2cgen/interpreters/python/interpreter.py

@@ -4,21 +4,19 @@
 
 class PythonInterpreter(BaseInterpreter):
 
-    def __init__(self, model_name="Model", indent=4, *args, **kwargs):
-        self.model_name = model_name
+    def __init__(self, indent=4, *args, **kwargs):
         cg = PythonCodeGenerator(indent=indent)
         super(PythonInterpreter, self).__init__(cg, *args, **kwargs)
 
     def interpret(self, expr):
         self._cg.reset_state()
 
-        with self._cg.class_definition(self.model_name):
-            with self._cg.method_definition(
-                    name="score",
-                    args=[self._feature_array_name]):
-                last_result = self._do_interpret(expr)
-                self._cg.add_return_statement(last_result)
+        with self._cg.method_definition(
+                name="score",
+                args=[self._feature_array_name]):
+            last_result = self._do_interpret(expr)
+            self._cg.add_return_statement(last_result)
 
         return [
-            (self.model_name, self._cg.code),
+            ("", self._cg.code),
         ]

tests/interpreters/test_python.py

@@ -0,0 +1,115 @@
+from m2cgen import ast
+from m2cgen import interpreters
+from tests import utils
+
+
+def test_if_expr():
+    expr = ast.IfExpr(
+        ast.CompExpr(ast.NumVal(1), ast.FeatureRef(0), ast.CompOpType.EQ),
+        ast.NumVal(2),
+        ast.NumVal(3))
+
+    interpreter = interpreters.PythonInterpreter()
+
+    expected_code = """
+def  score(input):
+    if (1) == (input[0]):
+        var0 = 2
+    else:
+        var0 = 3
+    return var0
+"""
+
+    utils.assert_code_equal(interpreter.interpret(expr)[0][1], expected_code)
+
+
+def test_bin_num_expr():
+    expr = ast.BinNumExpr(
+        ast.BinNumExpr(
+            ast.FeatureRef(0), ast.NumVal(-2), ast.BinNumOpType.DIV),
+        ast.NumVal(2),
+        ast.BinNumOpType.MUL)
+
+    interpreter = interpreters.PythonInterpreter()
+
+    expected_code = """
+def  score(input):
+    return ((input[0]) / (-2)) * (2)
+    """
+
+    utils.assert_code_equal(interpreter.interpret(expr)[0][1], expected_code)
+
+
+def test_dependable_condition():
+    left = ast.BinNumExpr(
+        ast.IfExpr(
+            ast.CompExpr(ast.NumVal(1),
+                         ast.NumVal(1),
+                         ast.CompOpType.EQ),
+            ast.NumVal(1),
+            ast.NumVal(2)),
+        ast.NumVal(2),
+        ast.BinNumOpType.ADD)
+
+    right = ast.BinNumExpr(ast.NumVal(1), ast.NumVal(2), ast.BinNumOpType.DIV)
+    bool_test = ast.CompExpr(left, right, ast.CompOpType.GTE)
+
+    expr = ast.IfExpr(bool_test, ast.NumVal(1), ast.FeatureRef(0))
+
+    expected_code = """
+def  score(input):
+    if (1) == (1):
+        var1 = 1
+    else:
+        var1 = 2
+    if ((var1) + (2)) >= ((1) / (2)):
+        var0 = 1
+    else:
+        var0 = input[0]
+    return var0
+    """
+
+    interpreter = interpreters.PythonInterpreter()
+
+    utils.assert_code_equal(interpreter.interpret(expr)[0][1], expected_code)
+
+
+def test_nested_condition():
+    left = ast.BinNumExpr(
+        ast.IfExpr(
+            ast.CompExpr(ast.NumVal(1),
+                         ast.NumVal(1),
+                         ast.CompOpType.EQ),
+            ast.NumVal(1),
+            ast.NumVal(2)),
+        ast.NumVal(2),
+        ast.BinNumOpType.ADD)
+
+    bool_test = ast.CompExpr(ast.NumVal(1), left, ast.CompOpType.EQ)
+
+    expr_nested = ast.IfExpr(bool_test, ast.FeatureRef(2), ast.NumVal(2))
+
+    expr = ast.IfExpr(bool_test, expr_nested, ast.NumVal(2))
+
+    expected_code = """
+def  score(input):
+    if (1) == (1):
+        var1 = 1
+    else:
+        var1 = 2
+    if (1) == ((var1) + (2)):
+        if (1) == (1):
+            var2 = 1
+        else:
+            var2 = 2
+        if (1) == ((var2) + (2)):
+            var0 = input[2]
+        else:
+            var0 = 2
+    else:
+        var0 = 2
+    return var0
+    """
+
+    interpreter = interpreters.PythonInterpreter()
+    utils.assert_code_equal(interpreter.interpret(expr)[0][1], expected_code)

tests/utils.py

@@ -30,6 +30,10 @@ def cmp_exprs(left, right):
     return False
 
 
+def assert_code_equal(actual, expected):
+    assert actual.strip() == expected.strip()
+
+
 def train_model(estimator, test_fraction=0.1):
     boston = load_boston()