DFP-3901: Fixed lint issue

Author: deepinsight1

Diff for: spark_pipeline_framework/transformers/framework_exception_handler/v1/framework_exception_handler.py

@@ -1,6 +1,5 @@
-from typing import Dict, Any, Optional, Union, List, Callable
+from typing import Dict, Any, Optional, Type, Union, List, Callable
 
-from spark_pipeline_framework.logger.log_level import LogLevel
 from spark_pipeline_framework.utilities.capture_parameters import capture_parameters
 from pyspark.ml import Transformer
 from pyspark.sql.dataframe import DataFrame
@@ -18,7 +17,7 @@ def __init__(
         self,
         *,
         raise_on_exception: Optional[Union[bool, Callable[[DataFrame], bool]]] = True,
-        error_exception: Optional[Exception] = Exception,
+        error_exception: Type[BaseException] = BaseException,
         stages: Union[List[Transformer], Callable[[], List[Transformer]]],
         exception_stages: Optional[
             Union[List[Transformer], Callable[[], List[Transformer]]]
@@ -50,18 +49,18 @@ def __init__(
             raise_on_exception
         )
 
-        self.error_exception: Optional[Exception] = error_exception
+        self.error_exception: Type[BaseException] = error_exception
         self.stages: Union[List[Transformer], Callable[[], List[Transformer]]] = stages
-        self.exception_stages: Optional[
-            Union[List[Transformer], Callable[[], List[Transformer]]]
+        self.exception_stages: Union[
+            List[Transformer], Callable[[], List[Transformer]]
         ] = (exception_stages or [])
 
         self.loop_id: Optional[str] = None
 
         kwargs = self._input_kwargs
         self.setParams(**kwargs)
 
-    async def _transform_async(self, df):
+    async def _transform_async(self, df: DataFrame) -> DataFrame:
         """
         Executes the transformation pipeline asynchronously.
 
@@ -81,8 +80,8 @@ async def run_pipeline(
             df: DataFrame,
             stages: Union[List[Transformer], Callable[[], List[Transformer]]],
             progress_logger: Optional[ProgressLogger],
-        ):
-            stages: List[Transformer] = stages if not callable(stages) else stages()
+        ) -> None:
+            stages = stages if not callable(stages) else stages()
             nonlocal stage_name
 
             for stage in stages:
@@ -112,7 +111,6 @@ async def run_pipeline(
                 progress_logger.write_to_log(
                     self.getName() or "FrameworkExceptionHandlerTransformer",
                     f"Failed while running steps with error: {e}. Run execution steps: {isinstance(e, self.error_exception)}",
-                    log_level=LogLevel.INFO,
                 )
 
             try:
@@ -140,8 +138,8 @@ def as_dict(self) -> Dict[str, Any]:
                 else str(self.stages)
             ),
             "exception_stages": (
-                [s.as_dict() for s in self.else_stages]  # type: ignore
-                if self.else_stages and not callable(self.else_stages)
-                else str(self.else_stages)
+                [s.as_dict() for s in self.exception_stages]  # type: ignore
+                if self.exception_stages and not callable(self.exception_stages)
+                else str(self.exception_stages)
             ),
         }

Diff for: spark_pipeline_framework/transformers/framework_exception_handler/v1/test/test_framework_exception_handler.py

@@ -16,13 +16,13 @@
 
 
 def test_framework_exception_handle(spark_session: SparkSession) -> None:
-    
+
     # create a dataframe with the test data
     data_dir: Path = Path(__file__).parent.joinpath("./")
     df: DataFrame = create_empty_dataframe(spark_session=spark_session)
     invalid_view: str = "invalid_view"
     valid_view = "view"
-    
+
     with ProgressLogger() as progress_logger:
         FrameworkExceptionHandlerTransformer(
             name="Exception Handler Test",
@@ -42,12 +42,12 @@ def test_framework_exception_handle(spark_session: SparkSession) -> None:
                 )
             ],
             raise_on_exception=False,
-            progress_logger=progress_logger
+            progress_logger=progress_logger,
         ).transform(df)
     result_df: DataFrame = spark_session.table(valid_view)
-    
+
     # Assert that the exception-handling stage has successfully run
     assert result_df.count() == 3
-    
+
     # Verify that the invalid view was NOT created, confirming that the original stage failed
     assert not spark_session.catalog.tableExists(invalid_view)