support pickle and file inputs

Author: jlwalke2

src/sasctl/utils/pyml2ds/__init__.py

@@ -1 +1 @@
-from .pyml2ds import pyml2ds
+from .core import pyml2ds

src/sasctl/utils/pyml2ds/pyml2ds.py renamed to src/sasctl/utils/pyml2ds/core.py

@@ -1,10 +1,8 @@
+import os
+import pickle
 import xml.etree.ElementTree as etree
 
-try:
-    import pickle
-except ImportError:
-    pickle = None
-
+import six
 try:
     import xgboost
 except ImportError:
@@ -50,23 +48,43 @@ def pyml2ds(in_file, out_file, out_var_name="P_TARGET"):
 
     Parameters
     ----------
-    in_file : str
-        Path to file to be translated.
+    in_file : str or bytes or file-like
+        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.
 
     """
-    # Load model file
-    ext = ".pmml"
-    if in_file[-len(ext):] == ext:
-        model = etree.parse(in_file)
+
+    try:
+        # In Python2 str could either be a path or the binary pickle data,
+        # so check if its a valid filepath too.
+        is_file_path = isinstance(in_file, six.string_types) and os.path.isfile(in_file)
+    except TypeError:
+        is_file_path = False
+
+    # Path to a PMML or pickle file
+    if is_file_path:
+        # Parse PMML files
+        if os.path.splitext(in_file)[-1] == '.pmml':
+            model = etree.parse(in_file)
+        else:
+            # Read pickled files
+            with open(in_file, 'rb') as f:
+                model = pickle.load(f)
+
+    elif isinstance(in_file, bytes):
+        # Assume byte string is the actual pickled bytes
+        model = pickle.loads(in_file)
     else:
-        with open(in_file, 'rb') as mf:
-            model = pickle.load(mf)
+        # Assume a file object containing the pickled object
+        model = pickle.load(in_file)
 
+    # 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.translate(f)

tests/unit/test_pyml2ds.py

@@ -109,3 +109,73 @@ def _leaf_value(self):
     assert result == expected
 
 
+def test_path_input(tmpdir_factory):
+    """pyml2ds should accept a file path (str) as input."""
+    import pickle
+    from sasctl.utils.pyml2ds import pyml2ds
+
+    # The target "model" to use
+    target = {'msg': 'hello world'}
+
+    # Pickle the "model" to a file
+    temp_dir = tmpdir_factory.mktemp('pyml2ds')
+    in_file = str(temp_dir.join('model.pkl'))
+    out_file = str(temp_dir.join('model.sas'))
+    with open(in_file, 'wb') as f:
+        pickle.dump(target, f)
+
+    with mock.patch('sasctl.utils.pyml2ds.core._check_type') as check:
+        check.translate.return_value = 'translated'
+        pyml2ds(in_file, out_file)
+
+    # Verify _check_type should have been called with the "model"
+    assert check.call_count == 1
+    assert check.call_args[0][0] == target
+
+
+def test_file_input():
+    """pyml2ds should accept a file-like obj as input."""
+    import io
+    import pickle
+    from sasctl.utils.pyml2ds import pyml2ds
+
+    # The target "model" to use
+    target = {'msg': 'hello world'}
+
+    # Pickle the "model" to a file-like object
+    in_file = io.BytesIO(pickle.dumps(target))
+    out_file = 'model.sas'
+
+    with mock.patch('sasctl.utils.pyml2ds.core._check_type') as check:
+        check.translate.return_value = 'translated'
+        pyml2ds(in_file, out_file)
+
+    # Verify _check_type should have been called with the "model"
+    assert check.call_count == 1
+    assert check.call_args[0][0] == target
+
+def test_pickle_input():
+    """pyml2ds should accept a binary pickle string as input."""
+    import pickle
+    from sasctl.utils.pyml2ds import pyml2ds
+
+    # The target "model" to use
+    target = {'msg': 'hello world'}
+
+    # Pickle the "model" to a file-like object
+    in_file = pickle.dumps(target)
+    out_file = 'model.sas'
+
+    with mock.patch('sasctl.utils.pyml2ds.core._check_type') as check:
+        check.translate.return_value = 'translated'
+        pyml2ds(in_file, out_file)
+
+    # Verify _check_type should have been called with the "model"
+    assert check.call_count == 1
+    assert check.call_args[0][0] == target
+
+
+
+
+
+