Updates to preprocess function

sassoftware · Oct 21, 2024 · 9158ab2 · 9158ab2
1 parent d73f6b6
commit 9158ab2
Show file tree

Hide file tree

Showing 2 changed files with 4 additions and 4 deletions.
diff --git a/src/sasctl/pzmm/import_model.py b/src/sasctl/pzmm/import_model.py
@@ -213,7 +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,
+        preprocess_function: Optional[Callable[[DataFrame], DataFrame]] = None,
         **kwargs,
     ) -> Tuple[RestObj, Union[dict, str, Path]]:
         """

diff --git a/src/sasctl/pzmm/write_score_code.py b/src/sasctl/pzmm/write_score_code.py
@@ -36,7 +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,
+        preprocess_function: Optional[Callable[[DataFrame], DataFrame]] = None,
         **kwargs,
     ) -> Union[dict, None]:
         """
@@ -763,7 +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,
+        preprocess_function: Optional[Callable[[DataFrame], DataFrame]] = None,
     ) -> None:
         """
         Write the model prediction section of the score code.
@@ -2260,7 +2260,7 @@ def _viya35_score_code_import(
                     mr.update_model(model)
                     return mas_code, cas_code
 
-    def _add_preprocess_code(self, preprocess_function: Callable[DataFrame, DataFrame]):
+    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