removed test= parameter

Author: jlwalke2

src/sasctl/utils/pyml2ds/basic/tree.py

@@ -97,7 +97,7 @@ def _remove_diacritic(self, input):
 
         return output
 
-    def parse_node(self, f, node=None, test=False):
+    def parse_node(self, f, node=None):
         """Recursively parse tree node and write generated SAS code to file.
 
         Attributes
@@ -108,8 +108,6 @@ def parse_node(self, f, node=None, test=False):
             Tree node to process.
 
         """
-        if test:
-            self.test = True
 
         pnode = self._node
         if node is not None:
@@ -128,8 +126,6 @@ def parse_node(self, f, node=None, test=False):
             var = self._remove_diacritic(var)
 
             split_value = self._split_value()
-            if self.test:
-                split_value = int(float(split_value))
 
             cond = ""
             if self._go_left():
@@ -151,8 +147,6 @@ def parse_node(self, f, node=None, test=False):
             f.write(self._get_indent() + "end;\n")
         else:
             leaf_value = self._leaf_value()
-            if self.test:
-                leaf_value = int(float(leaf_value))
 
             f.write(self._get_indent() + "treeValue%s = %s;\n"
                     % (self._tree_id, leaf_value))

src/sasctl/utils/pyml2ds/connectors/ensembles/core.py

@@ -29,7 +29,7 @@ def _aggregate(self, booster_count):
         return "treeValue = sum({});\n".format(', '.join(
             ["treeValue%d" % i for i in range(booster_count)]))
 
-    def translate(self, f, test=False):
+    def translate(self, f):
         """Translate a gradient boosting model and write SAS scoring code to a file.
 
         Attributes
@@ -42,7 +42,7 @@ def translate(self, f, test=False):
             f.write("/* Parsing tree {}*/\n".format(booster_id))
 
             self._tree_parser.init(tree, booster_id)
-            self._tree_parser.parse_node(f, test=test)
+            self._tree_parser.parse_node(f)
 
             f.write("\n")
 

src/sasctl/utils/pyml2ds/pyml2ds.py

@@ -38,12 +38,12 @@ def _check_type(model):
     else:
         raise RuntimeError("Unknown booster type: %s. Compatible types are: %s."
                            " Check if corresponding library is installed."
-                           % type(model).__name__)
+                           % (type(model).__name__, ', '.join(comp_types)))
 
     return parser
 
 
-def pyml2ds(in_file, out_file, out_var_name="P_TARGET", test=False):
+def pyml2ds(in_file, out_file, out_var_name="P_TARGET"):
     """Translate a gradient boosting model and write SAS scoring code to file.
 
     Supported models are: xgboost, lightgbm and pmml gradient boosting.
@@ -69,4 +69,4 @@ def pyml2ds(in_file, out_file, out_var_name="P_TARGET", test=False):
     parser = _check_type(model)
     parser.out_var_name = out_var_name
     with open(out_file, "w") as f:
-        parser.translate(f, test=test)
+        parser.translate(f)

tests/unit/test_pyml2ds.py

@@ -6,12 +6,14 @@
 
 
 import os
-dirname = os.path.dirname
 
 import pytest
+from six.moves import mock
+
 from sasctl.utils.pyml2ds import pyml2ds
 
 
+dirname = os.path.dirname
 DATA_PATH = os.path.join(dirname(dirname(__file__)), 'pyml2ds_data')
 
 
@@ -22,7 +24,25 @@ def test_xgb2ds(tmpdir):
     OUT_SAS = os.path.join(str(tmpdir), 'xgb.sas')
     EXPECTED_SAS = os.path.join(DATA_PATH, 'xgb.sas')
 
-    pyml2ds(IN_PKL, OUT_SAS, test=True)
+    from sasctl.utils.pyml2ds.connectors.ensembles.xgb import XgbTreeParser
+
+    # Expected output contains integer values instead of floats.
+    # Convert to ensure match.
+    class TestXgbTreeParser(XgbTreeParser):
+        def _split_value(self):
+            val = super(TestXgbTreeParser, self)._split_value()
+            return int(float(val))
+
+        def _leaf_value(self):
+            val = super(TestXgbTreeParser, self)._leaf_value()
+            return int(float(val))
+
+    test_parser = TestXgbTreeParser()
+
+    with mock.patch('sasctl.utils.pyml2ds.connectors.ensembles.xgb.XgbTreeParser') as parser:
+        parser.return_value = test_parser
+        pyml2ds(IN_PKL, OUT_SAS)
+
     result = open(OUT_SAS, 'rb').read()
     expected = open(EXPECTED_SAS, 'rb').read()
     assert result == expected
@@ -35,7 +55,26 @@ def test_lgb2ds(tmpdir):
     OUT_SAS = os.path.join(str(tmpdir), 'lgb.sas')
     EXPECTED_SAS = os.path.join(DATA_PATH, 'lgb.sas')
 
-    pyml2ds(IN_PKL, OUT_SAS, test=True)
+    from sasctl.utils.pyml2ds.connectors.ensembles.lgb import LightgbmTreeParser
+
+    # Expected output contains integer values instead of floats.
+    # Convert to ensure match.
+    class TestLightgbmTreeParser(LightgbmTreeParser):
+        def _split_value(self):
+            val = super(TestLightgbmTreeParser, self)._split_value()
+            return int(float(val))
+
+        def _leaf_value(self):
+            val = super(TestLightgbmTreeParser, self)._leaf_value()
+            return int(float(val))
+
+    test_parser = TestLightgbmTreeParser()
+
+    with mock.patch('sasctl.utils.pyml2ds.connectors.ensembles.lgb.LightgbmTreeParser') as parser:
+        parser.return_value = test_parser
+        pyml2ds(IN_PKL, OUT_SAS)
+
+
     result = open(OUT_SAS, 'rb').read()
     expected = open(EXPECTED_SAS, 'rb').read()
     assert result == expected
@@ -46,7 +85,27 @@ def test_gbm2ds(tmpdir):
     OUT_SAS = os.path.join(str(tmpdir), 'gbm.sas')
     EXPECTED_SAS = os.path.join(DATA_PATH, 'gbm.sas')
 
-    pyml2ds(IN_PKL, OUT_SAS, test=True)
+    from sasctl.utils.pyml2ds.connectors.ensembles.pmml import PmmlTreeParser
+
+    # Expected output contains integer values instead of floats.
+    # Convert to ensure match.
+    class TestPmmlTreeParser(PmmlTreeParser):
+        def _split_value(self):
+            val = super(TestPmmlTreeParser, self)._split_value()
+            return int(float(val))
+
+        def _leaf_value(self):
+            val = super(TestPmmlTreeParser, self)._leaf_value()
+            return int(float(val))
+
+    test_parser = TestPmmlTreeParser()
+
+    with mock.patch('sasctl.utils.pyml2ds.connectors.ensembles.pmml.PmmlTreeParser') as parser:
+        parser.return_value = test_parser
+        pyml2ds(IN_PKL, OUT_SAS)
+
     result = open(OUT_SAS, 'rb').read()
     expected = open(EXPECTED_SAS, 'rb').read()
     assert result == expected
+
+