feat: upgrade to pydantic v2 (#160)

Author: thoas

poetry.lock

@@ -28,7 +28,7 @@ keywords = ["taskiq", "tasks", "distributed", "async"]
 [tool.poetry.dependencies]
 python = "^3.8.1"
 typing-extensions = ">=3.10.0.0"
-pydantic = "^1.6.2"
+pydantic = ">=1.0,<=3.0"
 importlib-metadata = "*"
 pycron = "^3.0.0"
 taskiq_dependencies = "^1"

pyproject.toml

@@ -0,0 +1,52 @@
+# flake8: noqa
+from typing import Any, Dict, Optional, Type, TypeVar, Union
+
+import pydantic
+from importlib_metadata import version
+from packaging.version import Version, parse
+
+PYDANTIC_VER = parse(version("pydantic"))
+
+Model = TypeVar("Model", bound="pydantic.BaseModel")
+
+
+if PYDANTIC_VER >= Version("2.0"):
+    T = TypeVar("T")
+
+    def parse_obj_as(annot: T, obj: Any) -> T:
+        return pydantic.TypeAdapter(annot).validate_python(obj)
+
+    def model_validate_json(
+        model_class: Type[Model],
+        message: Union[str, bytes, bytearray],
+    ) -> Model:
+        return model_class.model_validate_json(message)
+
+    def model_dump_json(instance: Model) -> str:
+        return instance.model_dump_json()
+
+    def model_copy(
+        instance: Model,
+        update: Optional[Dict[str, Any]] = None,
+        deep: bool = False,
+    ) -> Model:
+        return instance.model_copy(update=update, deep=deep)
+
+else:
+    parse_obj_as = pydantic.parse_obj_as  # type: ignore
+
+    def model_validate_json(
+        model_class: Type[Model],
+        message: Union[str, bytes, bytearray],
+    ) -> Model:
+        return model_class.parse_raw(message)
+
+    def model_dump_json(instance: Model) -> str:
+        return instance.json()
+
+    def model_copy(
+        instance: Model,
+        update: Optional[Dict[str, Any]] = None,
+        deep: bool = False,
+    ) -> Model:
+        return instance.copy(update=update, deep=deep)

taskiq/compat.py

@@ -1,4 +1,5 @@
 from taskiq.abc.formatter import TaskiqFormatter
+from taskiq.compat import model_dump_json, model_validate_json
 from taskiq.message import BrokerMessage, TaskiqMessage
 
 
@@ -15,7 +16,7 @@ def dumps(self, message: TaskiqMessage) -> BrokerMessage:
         return BrokerMessage(
             task_id=message.task_id,
             task_name=message.task_name,
-            message=message.json().encode(),
+            message=model_dump_json(message).encode(),
             labels=message.labels,
         )
 
@@ -26,4 +27,4 @@ def loads(self, message: bytes) -> TaskiqMessage:
         :param message: broker's message.
         :return: parsed taskiq message.
         """
-        return TaskiqMessage.parse_raw(message)
+        return model_validate_json(TaskiqMessage, message)

taskiq/formatters/json_formatter.py

@@ -2,8 +2,7 @@
 from logging import getLogger
 from typing import Any, Dict, Optional
 
-from pydantic import parse_obj_as
-
+from taskiq.compat import parse_obj_as
 from taskiq.message import TaskiqMessage
 
 logger = getLogger(__name__)

taskiq/receiver/params_parser.py

@@ -0,0 +1,14 @@
+# flake8: noqa
+from packaging.version import Version
+
+from taskiq.compat import PYDANTIC_VER
+
+if PYDANTIC_VER >= Version("2.0"):
+    from .v2 import TaskiqResult
+else:
+    from .v1 import TaskiqResult
+
+
+__all__ = [
+    "TaskiqResult",
+]

taskiq/result/__init__.py

@@ -0,0 +1,67 @@
+import json
+import pickle  # noqa: S403
+from typing import Any, Dict, Generic, Optional, TypeVar
+
+from pydantic import BaseModel, ConfigDict, Field, field_serializer, field_validator
+from typing_extensions import Self
+
+from taskiq.serialization import exception_to_python, prepare_exception
+
+_ReturnType = TypeVar("_ReturnType")
+
+
+class TaskiqResult(BaseModel, Generic[_ReturnType]):
+    """Result of a remote task invocation."""
+
+    is_err: bool
+    # Log is a deprecated field. It would be removed in future
+    # releases of not, if we find a way to capture logs in async
+    # environment.
+    log: Optional[str] = None
+    return_value: _ReturnType
+    execution_time: float
+    labels: Dict[str, str] = Field(default_factory=dict)
+
+    error: Optional[BaseException] = None
+
+    model_config = ConfigDict(arbitrary_types_allowed=True)
+
+    @field_serializer("error")
+    def serialize_error(self, value: BaseException) -> Any:
+        """
+        Serialize error field.
+
+        :returns: Any
+        :param value: exception to serialize.
+        """
+        if value:
+            return prepare_exception(value, json)
+
+        return None
+
+    def raise_for_error(self) -> "Self":
+        """Raise exception if `error`.
+
+        :raises error: task execution exception
+        :returns: TaskiqResult
+        """
+        if self.error is not None:
+            raise self.error
+        return self
+
+    def __getstate__(self) -> Dict[Any, Any]:
+        dict = super().__getstate__()  # noqa: WPS125
+        vals: Dict[str, Any] = dict["__dict__"]
+
+        if "error" in vals and vals["error"] is not None:
+            vals["error"] = prepare_exception(
+                vals["error"],
+                pickle,
+            )
+
+        return dict
+
+    @field_validator("error", mode="before")
+    @classmethod
+    def _validate_error(cls, value: Any) -> Optional[BaseException]:
+        return exception_to_python(value)

taskiq/result.py renamed to taskiq/result/v1.py

@@ -297,7 +297,7 @@ def _prepare_exception(
         SEEN_EXCEPTIONS_CACHE.discard(id(exc))
 
 
-@pydantic.validate_arguments(config={"arbitrary_types_allowed": True})
+@pydantic.validate_arguments(config={"arbitrary_types_allowed": True})  # type: ignore
 def prepare_exception(
     exc: BaseException,
     coder: Coder[Any, Any],
@@ -312,7 +312,7 @@ def prepare_exception(
     return _prepare_exception(exc, coder)  # type: ignore
 
 
-@pydantic.validate_arguments(config={"arbitrary_types_allowed": True})
+@pydantic.validate_arguments(config={"arbitrary_types_allowed": True})  # type: ignore
 def exception_to_python(  # noqa: C901, WPS210
     exc: Optional[Union[BaseException, ExceptionRepr]],
 ) -> Optional[BaseException]:

taskiq/result/v2.py

@@ -5,6 +5,7 @@
 import pytest
 from pydantic import BaseModel
 
+from taskiq.compat import model_copy
 from taskiq.message import TaskiqMessage
 from taskiq.receiver.params_parser import parse_params
 
@@ -27,7 +28,7 @@ def test_parse_params_no_signature() -> None:
         args=[1, 2],
         kwargs={"a": 1},
     )
-    modify_msg = src_msg.copy(deep=True)
+    modify_msg = model_copy(src_msg, deep=True)
     parse_params(
         signature=None,
         type_hints={},