return code instead of writing to file

Author: jlwalke2

src/sasctl/utils/pyml2ds/core.py

@@ -1,3 +1,4 @@
+import io
 import os
 import pickle
 import xml.etree.ElementTree as etree
@@ -43,7 +44,7 @@ def _check_type(model):
 
 
 @experimental
-def pyml2ds(in_file, out_file, out_var_name="P_TARGET"):
+def pyml2ds(in_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.
@@ -54,11 +55,14 @@ def pyml2ds(in_file, out_file, out_var_name="P_TARGET"):
         Pickled object to translate.  String is assumed to be a path to a picked
         file, file-like is assumed to be an open file handle to a pickle
         object, and bytes is assumed to be the raw pickled bytes.
-    out_file : str
-        Path to output file with SAS code.
     out_var_name : str (optional)
         Output variable name.
 
+    Returns
+    -------
+    str
+        A SAS Data Step program implementing the model.
+
     """
 
     try:
@@ -88,5 +92,12 @@ def pyml2ds(in_file, out_file, out_var_name="P_TARGET"):
     # Verify model is a valid type
     parser = _check_type(model)
     parser.out_var_name = out_var_name
-    with open(out_file, "w") as f:
-        parser.translate(f)
+
+    # Parser is currently written to expect a file input
+    # Until refactored, use StringIO to collect the text in memory
+    with io.StringIO() as f:
+        parser.translate(f)
+
+        # Return contents of "file"
+        f.seek(0)
+        return f.read()

tests/unit/test_pyml2ds.py

@@ -17,11 +17,10 @@
 DATA_PATH = os.path.join(dirname(dirname(__file__)), 'pyml2ds_data')
 
 
-def test_xgb2ds(tmpdir):
+def test_xgb2ds():
     pytest.importorskip('xgboost')
 
     IN_PKL = os.path.join(DATA_PATH, 'xgb.pkl')
-    OUT_SAS = os.path.join(str(tmpdir), 'xgb.sas')
     EXPECTED_SAS = os.path.join(DATA_PATH, 'xgb.sas')
 
     from sasctl.utils.pyml2ds.connectors.ensembles.xgb import XgbTreeParser
@@ -41,18 +40,16 @@ def _leaf_value(self):
 
     with mock.patch('sasctl.utils.pyml2ds.connectors.ensembles.xgb.XgbTreeParser') as parser:
         parser.return_value = test_parser
-        pyml2ds(IN_PKL, OUT_SAS)
+        result = pyml2ds(IN_PKL)
 
-    result = open(OUT_SAS, 'rb').read()
-    expected = open(EXPECTED_SAS, 'rb').read()
+    expected = open(EXPECTED_SAS, 'r').read()
     assert result == expected
 
 
-def test_lgb2ds(tmpdir):
+def test_lgb2ds():
     pytest.importorskip('lightgbm')
 
     IN_PKL = os.path.join(DATA_PATH, 'lgb.pkl')
-    OUT_SAS = os.path.join(str(tmpdir), 'lgb.sas')
     EXPECTED_SAS = os.path.join(DATA_PATH, 'lgb.sas')
 
     from sasctl.utils.pyml2ds.connectors.ensembles.lgb import LightgbmTreeParser
@@ -72,17 +69,14 @@ def _leaf_value(self):
 
     with mock.patch('sasctl.utils.pyml2ds.connectors.ensembles.lgb.LightgbmTreeParser') as parser:
         parser.return_value = test_parser
-        pyml2ds(IN_PKL, OUT_SAS)
+        result = pyml2ds(IN_PKL)
 
-
-    result = open(OUT_SAS, 'rb').read()
-    expected = open(EXPECTED_SAS, 'rb').read()
+    expected = open(EXPECTED_SAS, 'r').read()
     assert result == expected
 
 
-def test_gbm2ds(tmpdir):
+def test_gbm2ds():
     IN_PKL = os.path.join(DATA_PATH, 'gbm.pmml')
-    OUT_SAS = os.path.join(str(tmpdir), 'gbm.sas')
     EXPECTED_SAS = os.path.join(DATA_PATH, 'gbm.sas')
 
     from sasctl.utils.pyml2ds.connectors.ensembles.pmml import PmmlTreeParser
@@ -102,10 +96,9 @@ def _leaf_value(self):
 
     with mock.patch('sasctl.utils.pyml2ds.connectors.ensembles.pmml.PmmlTreeParser') as parser:
         parser.return_value = test_parser
-        pyml2ds(IN_PKL, OUT_SAS)
+        result = pyml2ds(IN_PKL)
 
-    result = open(OUT_SAS, 'rb').read()
-    expected = open(EXPECTED_SAS, 'rb').read()
+    expected = open(EXPECTED_SAS, 'r').read()
     assert result == expected