added capability to have preprocessing function in score code

Author: djm21

src/sasctl/pzmm/import_model.py

@@ -213,6 +213,7 @@ def import_model(
         target_values: Optional[List[str]] = None,
         overwrite_project_properties: Optional[bool] = False,
         target_index: Optional[int] = None,
+        preprocess_function: Optional[Callable[DataFrame, DataFrame]] = None,
         **kwargs,
     ) -> Tuple[RestObj, Union[dict, str, Path]]:
         """
@@ -371,6 +372,7 @@ def import_model(
                 missing_values=missing_values,
                 score_cas=score_cas,
                 target_index=target_index,
+                preprocess_function=preprocess_function,
                 **kwargs,
             )
             if score_code_dict:
@@ -471,6 +473,7 @@ def import_model(
                 missing_values=missing_values,
                 score_cas=score_cas,
                 target_index=target_index,
+                preprocess_function=preprocess_function,
                 **kwargs,
             )
             if score_code_dict:

src/sasctl/pzmm/write_score_code.py

@@ -36,6 +36,7 @@ def write_score_code(
         score_cas: Optional[bool] = True,
         score_code_path: Union[Path, str, None] = None,
         target_index: Optional[int] = None,
+        preprocess_function: Optional[Callable[DataFrame, DataFrame]] = None,
         **kwargs,
     ) -> Union[dict, None]:
         """
@@ -292,6 +293,9 @@ def score(var1, var2, var3, var4):
         if missing_values:
             self._impute_missing_values(input_data, missing_values)
 
+        if preprocess_function:
+            self._add_preprocess_code(preprocess_function)
+        
         # SAS Viya 3.5 model
         if model_id:
             mas_code, cas_code = self._viya35_score_code_import(
@@ -759,6 +763,7 @@ def _predict_method(
         missing_values: Optional[Any] = None,
         statsmodels_model: Optional[bool] = False,
         tf_model: Optional[bool] = False,
+        preprocess_function: Optional[Callable[DataFrame, DataFrame]] = None,
     ) -> None:
         """
         Write the model prediction section of the score code.
@@ -809,6 +814,10 @@ def _predict_method(
             input_frame = f'{{{", ".join(input_dict)}}}, index=index'
             self.score_code += self._wrap_indent_string(input_frame, 8)
             self.score_code += f"\n{'':4})\n"
+            if preprocess_function:
+                self.score_code += (
+                    f"{'':4}input_array = {preprocess_function.__name__}(input_array)"
+                )
             if missing_values:
                 self.score_code += (
                     f"{'':4}input_array = impute_missing_values(input_array)\n"
@@ -851,6 +860,10 @@ def _predict_method(
             input_frame = f'{{{", ".join(input_dict)}}}, index=index'
             self.score_code += self._wrap_indent_string(input_frame, 8)
             self.score_code += f"\n{'':4})\n"
+            if preprocess_function:
+                self.score_code += (
+                    f"{'':4}input_array = {preprocess_function.__name__}(input_array)"
+                )
             if missing_values:
                 self.score_code += (
                     f"{'':4}input_array = impute_missing_values(input_array)\n"
@@ -872,6 +885,10 @@ def _predict_method(
             input_frame = f'{{{", ".join(input_dict)}}}, index=index'
             self.score_code += self._wrap_indent_string(input_frame, 8)
             self.score_code += f"\n{'':4})\n"
+            if preprocess_function:
+                self.score_code += (
+                    f"{'':4}input_array = {preprocess_function.__name__}(input_array)"
+                )
             if missing_values:
                 self.score_code += (
                     f"{'':4}input_array = impute_missing_values(input_array)\n"
@@ -904,6 +921,10 @@ def _predict_method(
             input_frame = f'{{{", ".join(input_dict)}}}, index=index'
             self.score_code += self._wrap_indent_string(input_frame, 8)
             self.score_code += f"\n{'':4})\n"
+            if preprocess_function:
+                self.score_code += (
+                    f"{'':4}input_array = {preprocess_function.__name__}(input_array)"
+                )
             if missing_values:
                 self.score_code += (
                     f"{'':4}input_array = impute_missing_values(input_array)\n"
@@ -2238,3 +2259,30 @@ def _viya35_score_code_import(
                     model["scoreCodeType"] = "ds2MultiType"
                     mr.update_model(model)
                     return mas_code, cas_code
+
+    def _add_preprocess_code(
+        self, 
+        preprocess_function: Callable[DataFrame, DataFrame]
+    ):
+        """
+        Places the given preprocess function, which must both take a DataFrame as an argument
+        and return a DataFrame, into the score code. If the preprocess function does not   
+        return anything, an error is thrown.
+
+        Parameters
+        ----------
+        preprocess_function: function
+            The preprocess function to be added to the score code.
+        """
+        import inspect
+        preprocess_code = inspect.getsource(preprocess_function)
+        if not "return" in preprocess_code:
+            raise ValueError(
+                "The given score code does not return a value. " + 
+                "To allow for the score code to work correctly, please ensure the preprocessed " +
+                "data is returned."
+            )
+        if self.score_code[-1] == '\n':
+            self.score_code += preprocess_code
+        else:
+            self.score_code += '\n' + preprocess_code

tests/unit/test_write_score_code.py

@@ -188,6 +188,29 @@ def test_impute_missing_values():
     assert "'b': 'test'" in sc.score_code
     assert "'c': 1" in sc.score_code or "'c': np.int64(1)" in sc.score_code
 
+def test_preprocess_function():
+    """
+    Test Cases:
+    - function
+    - function with no return
+    """
+    test_df = pd.DataFrame(
+        data=[[0, "a", 1], [2, "b", 0]], columns=["num", "char", "bin"]
+    )
+    sc = ScoreCode()
+    def preprocess_function_one(data: pd.DataFrame):
+        print("preprocessing happens here")
+        return data
+    sc._add_preprocess_code(preprocess_function_one)
+    assert "preprocessing happens here" in sc.score_code
+    assert "preprocess_function_one" in sc.score_code
+
+    sc = ScoreCode()
+    def preprocess_function_two(data: pd.DataFrame):
+        print("preprocessing happens here?")
+    with pytest.raises(ValueError):
+        sc._add_preprocess_code(preprocess_function_two)
+    
 
 def test_predict_method():
     """