Fix '==' operator and categorical values for Python generation

Author: aficionado

bigml/model.py

@@ -56,7 +56,7 @@
 
 from bigml.api import FINISHED
 from bigml.util import invert_dictionary, slugify, split
-
+from bigml.fields import type_map
 
 # Map operator str to its corresponding function
 OPERATOR = {
@@ -68,6 +68,17 @@
     ">": operator.gt
 }
 
+# Map operator str to its corresponding python operator
+PYTHON_OPERATOR = {
+    "<": "<",
+    "<=": "<=",
+    "=": "==",
+    "!=": "!=",
+    ">=": ">=",
+    ">": ">"
+}
+
+
 INDENT = '    '
 
 
@@ -226,14 +237,11 @@ def python_body(self, depth=1, cmv=False):
                 body += ("%sif (%s %s %s):\n" %
                         (INDENT * depth,
                          self.fields[child.predicate.field]['slug'],
-                         child.predicate.operator,
-                         child.predicate.value))
+                         PYTHON_OPERATOR[child.predicate.operator],
+                         `child.predicate.value`))
                 body += child.python_body(depth + 1)
         else:
-            if self.fields[self.objective_field]['optype'] == 'numeric':
-                body = "%s return %s\n" % (INDENT * depth, self.output)
-            else:
-                body = "%s return '%s'\n" % (INDENT * depth, self.output)
+            body = "%s return %s\n" % (INDENT * depth, `self.output`)
         return body
 
     def python(self, out):
@@ -447,7 +455,7 @@ def extract_common_path(groups):
                             common_path.append(test_common_path)
                 groups[group]['total'][0] = common_path
                 if len(details) > 0:
-                    groups[group]['details'] = sorted(details, 
+                    groups[group]['details'] = sorted(details,
                                                       key=lambda x: x[1],
                                                       reverse=True)