Merge pull request #14 from asvetlov/serializer

Author: s3rius

poetry.lock

@@ -25,7 +25,7 @@ keywords = ["taskiq", "pipelines", "tasks", "distributed", "async"]
 
 [tool.poetry.dependencies]
 python = "^3.8.1"
-taskiq = ">=0.0.8, <1"
+taskiq = ">=0.10.1, <1"
 typing-extensions = "^4.3.0"
 pydantic = "^2"
 

pyproject.toml

@@ -20,24 +20,6 @@ def __init_subclass__(cls, step_name: str, **kwargs: Any) -> None:
         # known steps.
         cls._known_steps[step_name] = cls
 
-    @abstractmethod
-    def dumps(self) -> str:
-        """
-        Generate parsable string.
-
-        :return: dumped object.
-        """
-
-    @classmethod
-    @abstractmethod
-    def loads(cls: Type[_T], data: str) -> _T:
-        """
-        Method to load previously dumped data.
-
-        :param data: dumped data.
-        :return: instance of a class.
-        """
-
     @abstractmethod
     async def act(
         self,

taskiq_pipelines/abc.py

@@ -15,7 +15,7 @@
 class PipelineMiddleware(TaskiqMiddleware):
     """Pipeline middleware."""
 
-    async def post_save(  # noqa: C901, WPS212
+    async def post_save(  # noqa: C901, WPS210, WPS212
         self,
         message: "TaskiqMessage",
         result: "TaskiqResult[Any]",
@@ -41,9 +41,21 @@ async def post_save(  # noqa: C901, WPS212
             logger.warn("Pipline data not found. Execution flow is broken.")
             return
         pipeline_data = message.labels[PIPELINE_DATA]
+        # workaround for obligatory casting label values to `str`
+        # in `AsyncKicker._prepare_message`.
+        # The trick can be removed later after adding explicit `bytes` support.
+        if (  # noqa: WPS337
+            isinstance(pipeline_data, str)
+            and pipeline_data.startswith("b'")
+            and pipeline_data.endswith("'")
+        ):
+            pipeline_data2 = pipeline_data[2:-1].encode()
+        else:
+            pipeline_data2 = pipeline_data
+        parsed_data = self.broker.serializer.loadb(pipeline_data2)
         try:
-            steps_data = pydantic.TypeAdapter(List[DumpedStep]).validate_json(
-                pipeline_data,
+            steps_data = pydantic.TypeAdapter(List[DumpedStep]).validate_python(
+                parsed_data,
             )
         except ValueError:
             return

taskiq_pipelines/middleware.py

@@ -1,7 +1,7 @@
-import json
 from typing import (
     Any,
     Coroutine,
+    Dict,
     Generic,
     List,
     Literal,
@@ -29,10 +29,13 @@ class DumpedStep(pydantic.BaseModel):
     """Dumped state model."""
 
     step_type: str
-    step_data: str
+    step_data: Dict[str, Any]
     task_id: str
 
 
+DumpedSteps = pydantic.RootModel[List[DumpedStep]]
+
+
 class Pipeline(Generic[_FuncParams, _ReturnType]):
     """
     Pipeline constructor.
@@ -116,7 +119,7 @@ def call_next(
                     task=task,
                     param_name=param_name,
                     **additional_kwargs,
-                ).dumps(),
+                ).model_dump(),
                 task_id="",
             ),
         )
@@ -172,7 +175,7 @@ def call_after(
                     task=task,
                     param_name=EMPTY_PARAM_NAME,
                     **additional_kwargs,
-                ).dumps(),
+                ).model_dump(),
                 task_id="",
             ),
         )
@@ -243,7 +246,7 @@ def map(
                     skip_errors=skip_errors,
                     check_interval=check_interval,
                     **additional_kwargs,
-                ).dumps(),
+                ).model_dump(),
                 task_id="",
             ),
         )
@@ -315,24 +318,24 @@ def filter(
                     skip_errors=skip_errors,
                     check_interval=check_interval,
                     **additional_kwargs,
-                ).dumps(),
+                ).model_dump(),
                 task_id="",
             ),
         )
         return self
 
-    def dumps(self) -> str:
+    def dumpb(self) -> bytes:
         """
         Dumps current pipeline as string.
 
         :returns: serialized pipeline.
         """
-        return json.dumps(
-            [step.model_dump() for step in self.steps],
+        return self.broker.serializer.dumpb(
+            DumpedSteps.model_validate(self.steps).model_dump(),
         )
 
     @classmethod
-    def loads(cls, broker: AsyncBroker, pipe_data: str) -> "Pipeline[Any, Any]":
+    def loadb(cls, broker: AsyncBroker, pipe_data: bytes) -> "Pipeline[Any, Any]":
         """
         Parses serialized pipeline.
 
@@ -344,7 +347,8 @@ def loads(cls, broker: AsyncBroker, pipe_data: str) -> "Pipeline[Any, Any]":
         :return: new
         """
         pipe: "Pipeline[Any, Any]" = Pipeline(broker)
-        pipe.steps = pydantic.TypeAdapter(List[DumpedStep]).validate_json(pipe_data)
+        data = broker.serializer.loadb(pipe_data)
+        pipe.steps = DumpedSteps.model_validate(data)  # type: ignore[assignment]
         return pipe
 
     async def kiq(
@@ -383,7 +387,7 @@ async def kiq(
             )
             .with_task_id(step.task_id)
             .with_labels(
-                **{CURRENT_STEP: 0, PIPELINE_DATA: self.dumps()},  # type: ignore
+                **{CURRENT_STEP: 0, PIPELINE_DATA: self.dumpb()},  # type: ignore
             )
         )
         taskiq_task = await kicker.kiq(*args, **kwargs)

taskiq_pipelines/pipeliner.py

@@ -1,5 +1,6 @@
 """Package with default pipeline steps."""
 from logging import getLogger
+from typing import Any, Dict
 
 from taskiq_pipelines.abc import AbstractStep
 from taskiq_pipelines.steps.filter import FilterStep
@@ -9,12 +10,12 @@
 logger = getLogger(__name__)
 
 
-def parse_step(step_type: str, step_data: str) -> AbstractStep:
+def parse_step(step_type: str, step_data: Dict[str, Any]) -> AbstractStep:
     step_cls = AbstractStep._known_steps.get(step_type)  # noqa: WPS437
     if step_cls is None:
         logger.warning(f"Unknown step type: {step_type}")
         raise ValueError("Unknown step type.")
-    return step_cls.loads(step_data)
+    return step_cls(**step_data)
 
 
 __all__ = [

taskiq_pipelines/steps/__init__.py

@@ -78,24 +78,6 @@ class FilterStep(pydantic.BaseModel, AbstractStep, step_name="filter"):
     skip_errors: bool
     check_interval: float
 
-    def dumps(self) -> str:
-        """
-        Dumps step as string.
-
-        :return: returns json.
-        """
-        return self.model_dump_json()
-
-    @classmethod
-    def loads(cls, data: str) -> "FilterStep":
-        """
-        Parses mapper step from string.
-
-        :param data: dumped data.
-        :return: parsed step.
-        """
-        return pydantic.TypeAdapter(FilterStep).validate_json(data)
-
     async def act(
         self,
         broker: AsyncBroker,

taskiq_pipelines/steps/filter.py

@@ -75,24 +75,6 @@ class MapperStep(pydantic.BaseModel, AbstractStep, step_name="mapper"):
     skip_errors: bool
     check_interval: float
 
-    def dumps(self) -> str:
-        """
-        Dumps step as string.
-
-        :return: returns json.
-        """
-        return self.model_dump_json()
-
-    @classmethod
-    def loads(cls, data: str) -> "MapperStep":
-        """
-        Parses mapper step from string.
-
-        :param data: dumped data.
-        :return: parsed step.
-        """
-        return pydantic.TypeAdapter(MapperStep).validate_json(data)
-
     async def act(
         self,
         broker: AsyncBroker,

taskiq_pipelines/steps/mapper.py

@@ -24,24 +24,6 @@ class SequentialStep(pydantic.BaseModel, AbstractStep, step_name="sequential"):
     param_name: Union[Optional[int], str]
     additional_kwargs: Dict[str, Any]
 
-    def dumps(self) -> str:
-        """
-        Dumps step as string.
-
-        :return: returns json.
-        """
-        return self.model_dump_json()
-
-    @classmethod
-    def loads(cls, data: str) -> "SequentialStep":
-        """
-        Parses sequential step from string.
-
-        :param data: dumped data.
-        :return: parsed step.
-        """
-        return pydantic.TypeAdapter(SequentialStep).validate_json(data)
-
     async def act(
         self,
         broker: AsyncBroker,