|
| 1 | +from typing import Any, Dict |
| 2 | + |
| 3 | +from taskiq.abc.broker import AsyncBroker |
| 4 | +from taskiq.abc.formatter import TaskiqFormatter |
| 5 | +from taskiq.compat import IS_PYDANTIC2, Model, model_dump, model_validate |
| 6 | +from taskiq.message import BrokerMessage, TaskiqMessage |
| 7 | + |
| 8 | +if IS_PYDANTIC2: |
| 9 | + |
| 10 | + def model_dump(instance: Model) -> Dict[str, Any]: |
| 11 | + """Model dump.""" |
| 12 | + return instance.model_dump() |
| 13 | + |
| 14 | +else: |
| 15 | + |
| 16 | + def model_dump(instance: Model) -> Dict[str, Any]: |
| 17 | + """Model dump.""" |
| 18 | + return instance.dict() |
| 19 | + |
| 20 | + |
| 21 | +class PatchedFormatter(TaskiqFormatter): |
| 22 | + """Default taskiq formatter.""" |
| 23 | + |
| 24 | + def __init__(self, broker: AsyncBroker) -> None: |
| 25 | + self.broker = broker |
| 26 | + |
| 27 | + def dumps(self, message: TaskiqMessage) -> BrokerMessage: |
| 28 | + """ |
| 29 | + Dumps taskiq message to some broker message format. |
| 30 | +
|
| 31 | + :param message: message to send. |
| 32 | + :return: Dumped message. |
| 33 | + """ |
| 34 | + return BrokerMessage( |
| 35 | + task_id=message.task_id, |
| 36 | + task_name=message.task_name, |
| 37 | + message=self.broker.serializer.dumpb(model_dump(message)), |
| 38 | + labels=message.labels, |
| 39 | + ) |
| 40 | + |
| 41 | + def loads(self, message: bytes) -> TaskiqMessage: |
| 42 | + """ |
| 43 | + Loads json from message. |
| 44 | +
|
| 45 | + :param message: broker's message. |
| 46 | + :return: parsed taskiq message. |
| 47 | + """ |
| 48 | + return model_validate(TaskiqMessage, self.broker.serializer.loadb(message)) |
0 commit comments