diff --git a/taskiq_pipelines/pipeliner.py b/taskiq_pipelines/pipeliner.py index 44562ac..f7c3eb8 100644 --- a/taskiq_pipelines/pipeliner.py +++ b/taskiq_pipelines/pipeliner.py @@ -1,3 +1,4 @@ +from types import CoroutineType from typing import ( Any, Coroutine, @@ -69,7 +70,9 @@ def call_next( self: "Pipeline[_FuncParams, _ReturnType]", task: Union[ AsyncKicker[Any, Coroutine[Any, Any, _T2]], + AsyncKicker[Any, "CoroutineType[Any, Any, _T2]"], AsyncTaskiqDecoratedTask[Any, Coroutine[Any, Any, _T2]], + AsyncTaskiqDecoratedTask[Any, "CoroutineType[Any, Any, _T2]"], ], param_name: Union[Optional[str], Literal[-1]] = None, **additional_kwargs: Any, @@ -128,7 +131,9 @@ def call_after( self: "Pipeline[_FuncParams, _ReturnType]", task: Union[ AsyncKicker[Any, Coroutine[Any, Any, _T2]], + AsyncKicker[Any, "CoroutineType[Any, Any, _T2]"], AsyncTaskiqDecoratedTask[Any, Coroutine[Any, Any, _T2]], + AsyncTaskiqDecoratedTask[Any, "CoroutineType[Any, Any, _T2]"], ], **additional_kwargs: Any, ) -> "Pipeline[_FuncParams, _T2]": ... @@ -182,7 +187,9 @@ def map( self: "Pipeline[_FuncParams, _ReturnType]", task: Union[ AsyncKicker[Any, Coroutine[Any, Any, _T2]], + AsyncKicker[Any, "CoroutineType[Any, Any, _T2]"], AsyncTaskiqDecoratedTask[Any, Coroutine[Any, Any, _T2]], + AsyncTaskiqDecoratedTask[Any, "CoroutineType[Any, Any, _T2]"], ], param_name: Optional[str] = None, skip_errors: bool = False, @@ -251,7 +258,9 @@ def filter( self: "Pipeline[_FuncParams, _ReturnType]", task: Union[ AsyncKicker[Any, Coroutine[Any, Any, bool]], + AsyncKicker[Any, "CoroutineType[Any, Any, bool]"], AsyncTaskiqDecoratedTask[Any, Coroutine[Any, Any, bool]], + AsyncTaskiqDecoratedTask[Any, "CoroutineType[Any, Any, bool]"], ], param_name: Optional[str] = None, skip_errors: bool = False,