Repository navigation
Expand file tree
/
Copy pathprogress_tracker.py
More file actions
80 lines (58 loc) · 1.96 KB
/
Copy pathprogress_tracker.py
File metadata and controls
80 lines (58 loc) · 1.96 KB
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
import enum
from typing import Generic, Optional, Union
from taskiq_dependencies import Depends
from typing_extensions import TypeVar
from taskiq.compat import IS_PYDANTIC2
from taskiq.context import Context
_ProgressType = TypeVar("_ProgressType")
class TaskState(str, enum.Enum):
"""State of task execution."""
STARTED = "STARTED"
FAILURE = "FAILURE"
SUCCESS = "SUCCESS"
RETRY = "RETRY"
if IS_PYDANTIC2:
from pydantic import BaseModel, ConfigDict
class _TaskProgressConfig(BaseModel):
model_config = ConfigDict(arbitrary_types_allowed=True)
else:
from pydantic.generics import GenericModel
class _TaskProgressConfig(GenericModel): # type: ignore[no-redef]
class Config:
arbitrary_types_allowed = True
class TaskProgress(_TaskProgressConfig, Generic[_ProgressType]):
"""Progress of task execution."""
state: Union[TaskState, str]
meta: Optional[_ProgressType]
class ProgressTracker(Generic[_ProgressType]):
"""Task's dependency to set progress."""
def __init__(
self,
context: Context = Depends(),
) -> None:
self.context = context
async def set_progress(
self,
state: Union[TaskState, str],
meta: Optional[_ProgressType] = None,
) -> None:
"""Set progress.
:param state: TaskState or str
:param meta: progress data
"""
if meta is None:
progress = await self.get_progress()
meta = progress.meta if progress else None
progress = TaskProgress(
state=state,
meta=meta,
)
await self.context.broker.result_backend.set_progress(
self.context.message.task_id,
progress,
)
async def get_progress(self) -> Optional[TaskProgress[_ProgressType]]:
"""Get progress."""
return await self.context.broker.result_backend.get_progress(
self.context.message.task_id,
)