File tree Expand file tree Collapse file tree 2 files changed +19
-12
lines changed Expand file tree Collapse file tree 2 files changed +19
-12
lines changed Original file line number Diff line number Diff line change 77from taskiq .compat import IS_PYDANTIC2
88from taskiq .context import Context
99
10- if IS_PYDANTIC2 :
11- from pydantic import BaseModel as GenericModel
12- else :
13- from pydantic .generics import GenericModel # type: ignore[no-redef]
14-
15-
1610_ProgressType = TypeVar ("_ProgressType" )
1711
1812
@@ -25,15 +19,26 @@ class TaskState(str, enum.Enum):
2519 RETRY = "RETRY"
2620
2721
28- class TaskProgress (GenericModel , Generic [_ProgressType ]):
22+ if IS_PYDANTIC2 :
23+ from pydantic import BaseModel , ConfigDict
24+
25+ class _TaskProgressConfig (BaseModel ):
26+ model_config = ConfigDict (arbitrary_types_allowed = True )
27+
28+ else :
29+ from pydantic .generics import GenericModel
30+
31+ class _TaskProgressConfig (GenericModel ): # type: ignore[no-redef]
32+ class Config :
33+ arbitrary_types_allowed = True
34+
35+
36+ class TaskProgress (_TaskProgressConfig , Generic [_ProgressType ]):
2937 """Progress of task execution."""
3038
3139 state : Union [TaskState , str ]
3240 meta : Optional [_ProgressType ]
3341
34- class Config :
35- arbitrary_types_allowed = True
36-
3742
3843class ProgressTracker (Generic [_ProgressType ]):
3944 """Task's dependency to set progress."""
Original file line number Diff line number Diff line change @@ -401,8 +401,10 @@ def task_cb(task: "asyncio.Task[Any]") -> None:
401401 self .sem_prefetch .release ()
402402 message = await queue .get ()
403403 if message is QUEUE_DONE :
404- logger .info ("Waiting for running tasks to complete." )
405- await asyncio .wait (tasks , timeout = self .wait_tasks_timeout )
404+ # asyncio.wait will throw an error if there is nothing to wait for
405+ if tasks :
406+ logger .info ("Waiting for running tasks to complete." )
407+ await asyncio .wait (tasks , timeout = self .wait_tasks_timeout )
406408 break
407409
408410 task = asyncio .create_task (
You can’t perform that action at this time.
0 commit comments