import model from score code table

Author: jlwalke2

src/sasctl/tasks.py

@@ -187,7 +187,7 @@ def get_version(x):
     # If model is a CASTable then assume it holds an ASTORE model.
     # Import these via a ZIP file.
     if 'swat.cas.table.CASTable' in str(type(model)):
-        zipfile = utils.create_package_from_astore(model)
+        zipfile = utils.create_package(model)
 
         if create_project:
             project = mr.create_project(project, repository)

src/sasctl/utils/__init__.py

@@ -4,6 +4,6 @@
 # Copyright © 2019, SAS Institute Inc., Cary, NC, USA.  All Rights Reserved.
 # SPDX-License-Identifier: Apache-2.0
 
-from .astore import create_package_from_astore
+from .astore import create_package, create_package_from_astore
 
 

src/sasctl/utils/astore.py

@@ -12,12 +12,77 @@
 import uuid
 import zipfile
 
+import six
+
 try:
     import swat
 except ImportError:
     swat = None
 
 
+def create_package(table):
+    """Create an importable model package from a CAS table.
+
+    Parameters
+    ----------
+    table : swat.CASTable
+        The CAS table containing an ASTORE or score code.
+
+    Returns
+    -------
+    BytesIO
+        A byte stream representing a ZIP archive which can be imported.
+
+    See Also
+    --------
+    :meth:`model_repository.import_model_from_zip <.ModelRepository.import_model_from_zip>`
+
+    """
+    if swat is None:
+        raise RuntimeError(
+            "The 'swat' package is required to work with SAS models.")
+
+    assert isinstance(table, swat.CASTable)
+
+    if 'DataStepSrc' in table.columns:
+        return create_package_from_datastep(table)
+    else:
+        return create_package_from_astore(table)
+
+
+def create_package_from_datastep(table):
+    """Create an importable model package from a score code table.
+
+    Parameters
+    ----------
+    table : swat.CASTable
+        The CAS table containing the score code.
+
+    Returns
+    -------
+    BytesIO
+        A byte stream representing a ZIP archive which can be imported.
+
+    See Also
+    --------
+    :meth:`model_repository.import_model_from_zip <.ModelRepository.import_model_from_zip>`
+
+    """
+    assert 'DataStepSrc' in table.columns
+    sess = table.session.get_connection()
+
+    dscode = table.to_frame().loc[0, 'DataStepSrc']
+
+    file_metadata = [{'role': 'score', 'name': 'dmcas_scorecode.sas'}]
+
+    zip_file = _build_zip_from_files({
+        'fileMetadata.json': file_metadata,
+        'dmcas_scorecode.sas': dscode
+    })
+
+    return zip_file
+
+
 def create_package_from_astore(table):
     """Create an importable model package from an ASTORE.
 
@@ -81,45 +146,62 @@ def create_package_from_astore(table):
                         'uri': '/dataTables/dataSources/cas~fs~cas-shared-default~fs~ModelStore/tables/{}'.format(astore_filename),
                         'key': astore_key}]
 
+    zip_file = _build_zip_from_files({
+        'dmcas_packagescorecode.sas': '\n'.join(package_ds2),
+        'dmcas_epscorecode.sas': ep_ds2,
+        astore_filename: astore,
+        'ModelProperties.json': model_properties,
+        'fileMetadata.json': file_metadata,
+        'AstoreMetadata.json': astore_metadata,
+        'inputVar.json': input_vars,
+        'outputVar.json': output_vars
+    })
+
+    return zip_file
+
+
+def _build_zip_from_files(files):
+    """Create a ZIP file containing the provided files.
+
+    Parameters
+    ----------
+    files : dict
+        Dictionary of filename: content to be added to the .zip file.
+
+    Returns
+    -------
+    BytesIO
+        Byte stream representation of the .zip file.
+
+    """
     try:
         # Create a temp folder
         folder = tempfile.mkdtemp()
 
-        # Closure for easily adding JSON files
-        def json_file(data, filename):
-            filename = os.path.join(folder, filename)
-            with open(filename, 'w') as f:
-                json.dump(data, f, indent=1)
+        for k, v in six.iteritems(files):
+            filename = os.path.join(folder, k)
 
-        filename = os.path.join(folder, 'dmcas_packagescorecode.sas')
-        with open(filename, 'w') as f:
-            f.write('\n'.join(package_ds2))
+            # Write JSON file
+            if os.path.splitext(k)[-1].lower() == '.json':
+                with open(filename, 'w') as f:
+                    json.dump(v, f, indent=1)
+            else:
+                mode = 'wb' if isinstance(v, bytes) else 'w'
 
-        filename = os.path.join(folder, 'dmcas_epscorecode.sas')
-        with open(filename, 'w') as f:
-            f.write(ep_ds2)
-
-        filename = os.path.join(folder, astore_filename)
-        with open(filename, 'wb') as f:
-            f.write(astore)
-
-        json_file(model_properties, 'ModelProperties.json')
-        json_file(file_metadata, 'fileMetadata.json')
-        json_file(astore_metadata, 'AstoreMetadata.json')
-        json_file(input_vars, 'inputVar.json')
-        json_file(output_vars, 'outputVar.json')
+                with open(filename, mode) as f:
+                    f.write(v)
 
         files = os.listdir(folder)
 
         with zipfile.ZipFile(os.path.join(folder, 'model.zip'), 'w') as z:
             for file in files:
                 z.write(os.path.join(folder, file), file)
 
-        # Need to return the ZIP file data but also need to ensure the directory is cleaned up.
+        # Need to return the ZIP file data but also need to ensure the
+        # directory is cleaned up.
         # Read the bytes from disk and return an in memory "file".
         with open(os.path.join(folder, 'model.zip'), 'rb') as z:
             return io.BytesIO(z.read())
-
     finally:
         shutil.rmtree(folder)