diff --git a/karton/core/karton.py b/karton/core/karton.py index f7be257..6fb5992 100644 --- a/karton/core/karton.py +++ b/karton/core/karton.py @@ -17,7 +17,6 @@ from .task import Task, TaskState from .utils import timeout -KartonTaskException = Union[Exception, BaseException] class Producer(KartonBase): @@ -134,7 +133,7 @@ def __init__( self._post_hooks: List[ Tuple[ Optional[str], - Callable[[Task, Optional[KartonTaskException]], None], + Callable[[Task, Optional[BaseException]], None], ] ] = [] @@ -266,7 +265,7 @@ def add_pre_hook( def add_post_hook( self, - callback: Callable[[Task, Optional[KartonTaskException]], None], + callback: Callable[[Task, Optional[BaseException]], None], name: Optional[str] = None, ) -> None: """ @@ -295,7 +294,7 @@ def _run_pre_hooks(self) -> None: else: self.log.exception("Pre-hook failed") - def _run_post_hooks(self, exception: Optional[KartonTaskException]) -> None: + def _run_post_hooks(self, exception: Optional[BaseException]) -> None: """ Run registered postprocessing hooks @@ -437,7 +436,7 @@ def _send_signaling_status_task_begin(self, task: Task) -> None: self._send_signaling_status_task("task_begin") def _send_signaling_status_task_end( - self, task: Task, ex: Optional[KartonTaskException] + self, task: Task, ex: Optional[BaseException] ) -> None: """Send a begin status signaling task.