From 5a01b87f69834adfc4a6423e57c93446e301c8b4 Mon Sep 17 00:00:00 2001 From: Pratyush Sharma <56130065+pratyush618@users.noreply.github.com> Date: Sat, 30 May 2026 01:03:04 +0530 Subject: [PATCH 1/3] feat(serializers): add SmartSerializer (msgpack + cloudpickle) Tagged envelope: msgpack for plain data, cloudpickle fallback for arbitrary objects. A custom ExtType preserves tuples so payloads keep exact Python semantics. msgpack becomes a base dependency. --- py_src/taskito/__init__.py | 2 + py_src/taskito/serializers.py | 85 +++++++++++++++++++++++++++++++--- pyproject.toml | 1 + tests/core/test_serializers.py | 57 ++++++++++++++++++++++- 4 files changed, 137 insertions(+), 8 deletions(-) diff --git a/py_src/taskito/__init__.py b/py_src/taskito/__init__.py index ee63fdfa..491707e7 100644 --- a/py_src/taskito/__init__.py +++ b/py_src/taskito/__init__.py @@ -44,6 +44,7 @@ JsonSerializer, MsgPackSerializer, Serializer, + SmartSerializer, ) from taskito.task import TaskWrapper from taskito.testing import MockResource, TestMode, TestResult, TestResults @@ -83,6 +84,7 @@ "SerializationError", "Serializer", "Signature", + "SmartSerializer", "SoftTimeoutError", "TaskCancelledError", "TaskFailedError", diff --git a/py_src/taskito/serializers.py b/py_src/taskito/serializers.py index e83b0c83..7195e124 100644 --- a/py_src/taskito/serializers.py +++ b/py_src/taskito/serializers.py @@ -6,6 +6,47 @@ from typing import Any, Protocol, runtime_checkable import cloudpickle +import msgpack + +# Format tags for the envelope written by ``SmartSerializer``. A legacy +# cloudpickle payload starts with the pickle protocol-2+ opcode ``\x80``, which +# never collides with these tags — so untagged bytes are unambiguously legacy. +_CODEC_CLOUDPICKLE = b"\x00" +_CODEC_MSGPACK = b"\x01" + +# msgpack has no native tuple type and would silently flatten tuples to lists. +# A custom ExtType preserves them so payloads round-trip with exact Python +# semantics (a task returning ``(1, 2)`` gets back ``(1, 2)``, not ``[1, 2]``). +_EXT_TUPLE = 0 + + +def _msgpack_default(obj: Any) -> Any: + """Encode types msgpack can't represent natively. + + Tuples become a tagged ExtType (recursively packed). Anything else raises, + which ``SmartSerializer`` catches to fall back to cloudpickle. + """ + if isinstance(obj, tuple): + return msgpack.ExtType(_EXT_TUPLE, _msgpack_packb(list(obj))) + raise TypeError(f"Cannot msgpack-encode {type(obj).__name__}") + + +def _msgpack_ext_hook(code: int, data: bytes) -> Any: + if code == _EXT_TUPLE: + return tuple(_msgpack_unpackb(data)) + return msgpack.ExtType(code, data) + + +def _msgpack_packb(obj: Any) -> bytes: + # ``strict_types`` ensures tuples reach ``default`` instead of being + # auto-coerced to arrays; subclasses also route to the cloudpickle fallback. + return bytes( + msgpack.packb(obj, use_bin_type=True, strict_types=True, default=_msgpack_default) + ) + + +def _msgpack_unpackb(data: bytes) -> Any: + return msgpack.unpackb(data, raw=False, ext_hook=_msgpack_ext_hook) @runtime_checkable @@ -47,22 +88,52 @@ def loads(self, data: bytes) -> Any: class MsgPackSerializer: """MsgPack-based serializer for compact, cross-language payloads. - Requires the ``msgpack`` extra:: - - pip install taskito[msgpack] + Only handles msgpack-native types. For arbitrary Python objects use + :class:`SmartSerializer`, which falls back to cloudpickle. """ def dumps(self, obj: Any) -> bytes: - import msgpack - return bytes(msgpack.packb(obj, use_bin_type=True)) def loads(self, data: bytes) -> Any: - import msgpack - return msgpack.unpackb(data, raw=False) +class SmartSerializer: + """Default serializer: msgpack for plain payloads, cloudpickle fallback. + + Plain data (the common case) serializes via msgpack — faster and more + compact than cloudpickle. Anything msgpack can't encode (lambdas, closures, + arbitrary class instances) transparently falls back to cloudpickle. A + one-byte tag records which codec produced each payload. + + Tuples are preserved (via a msgpack ExtType), so payloads round-trip with + exact Python semantics. + + Backward compatible: untagged payloads (written by older versions, raw + cloudpickle) are detected and loaded as cloudpickle. + """ + + def dumps(self, obj: Any) -> bytes: + try: + return _CODEC_MSGPACK + _msgpack_packb(obj) + except Exception: + # msgpack rejects non-native types (lambdas, custom classes, …); + # cloudpickle handles them. + return _CODEC_CLOUDPICKLE + bytes(cloudpickle.dumps(obj)) + + def loads(self, data: bytes) -> Any: + if not data: + raise ValueError("Cannot deserialize empty payload") + tag, body = data[:1], data[1:] + if tag == _CODEC_MSGPACK: + return _msgpack_unpackb(body) + if tag == _CODEC_CLOUDPICKLE: + return cloudpickle.loads(body) + # Untagged: a legacy cloudpickle payload from before the envelope existed. + return cloudpickle.loads(data) + + class EncryptedSerializer: """Wraps another serializer with AES-256-GCM encryption at rest. diff --git a/pyproject.toml b/pyproject.toml index 7216f826..5479b7a2 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -24,6 +24,7 @@ classifiers = [ ] dependencies = [ "cloudpickle>=3.0", + "msgpack>=1.0", 'tzdata; platform_system == "Windows"', ] [project.urls] diff --git a/tests/core/test_serializers.py b/tests/core/test_serializers.py index c78f8391..2bcb1616 100644 --- a/tests/core/test_serializers.py +++ b/tests/core/test_serializers.py @@ -5,7 +5,12 @@ import pytest -from taskito.serializers import CloudpickleSerializer, JsonSerializer, Serializer +from taskito.serializers import ( + CloudpickleSerializer, + JsonSerializer, + Serializer, + SmartSerializer, +) class TestJsonSerializer: @@ -148,3 +153,53 @@ def test_short_ciphertext_fails(self) -> None: s = EncryptedSerializer(JsonSerializer(), os.urandom(32)) with pytest.raises(ValueError, match="too short"): s.loads(b"only-twelve-") + + +class TestSmartSerializer: + def test_plain_data_uses_msgpack_tag(self) -> None: + s = SmartSerializer() + encoded = s.dumps({"x": 1, "y": ["a", "b"]}) + assert encoded[:1] == b"\x01" + assert s.loads(encoded) == {"x": 1, "y": ["a", "b"]} + + def test_falls_back_to_cloudpickle_for_lambda(self) -> None: + s = SmartSerializer() + encoded = s.dumps(lambda x: x + 1) + assert encoded[:1] == b"\x00" + assert s.loads(encoded)(41) == 42 + + def test_falls_back_to_cloudpickle_for_custom_class(self) -> None: + s = SmartSerializer() + + class Point: + def __init__(self, x: int) -> None: + self.x = x + + encoded = s.dumps(Point(7)) + assert encoded[:1] == b"\x00" + assert s.loads(encoded).x == 7 + + def test_loads_legacy_untagged_cloudpickle(self) -> None: + """Payloads written before the envelope existed are raw cloudpickle + and must still deserialize.""" + legacy = CloudpickleSerializer().dumps({"legacy": True}) + # A protocol-2+ pickle starts with 0x80, never a codec tag. + assert legacy[:1] not in (b"\x00", b"\x01") + assert SmartSerializer().loads(legacy) == {"legacy": True} + + def test_preserves_tuples(self) -> None: + """Tuples round-trip as tuples (not lists) via the msgpack ExtType, + including nested ones — so task args/results keep exact semantics.""" + s = SmartSerializer() + payload = ((1, "two", 3.0), {"k": ("nested", "tuple")}) + restored = s.loads(s.dumps(payload)) + assert restored == payload + assert isinstance(restored[0], tuple) + assert isinstance(restored[1]["k"], tuple) + + def test_empty_payload_rejected(self) -> None: + with pytest.raises(ValueError, match="empty payload"): + SmartSerializer().loads(b"") + + def test_satisfies_protocol(self) -> None: + assert isinstance(SmartSerializer(), Serializer) From 9eac7bacb99dcaa39d8e0491eebcccd45405b93e Mon Sep 17 00:00:00 2001 From: Pratyush Sharma <56130065+pratyush618@users.noreply.github.com> Date: Sat, 30 May 2026 01:03:21 +0530 Subject: [PATCH 2/3] perf(serializers): default queues to SmartSerializer Normalize _deserialize_payload to an (args, kwargs) tuple so the worker stays robust to serializers that flatten tuples to lists (msgpack, JSON). --- py_src/taskito/app.py | 18 +++++++++++++----- 1 file changed, 13 insertions(+), 5 deletions(-) diff --git a/py_src/taskito/app.py b/py_src/taskito/app.py index 60bc0fde..c25b0150 100644 --- a/py_src/taskito/app.py +++ b/py_src/taskito/app.py @@ -52,7 +52,7 @@ from taskito.proxies.built_in import register_builtin_handlers from taskito.proxies.metrics import ProxyMetrics from taskito.result import JobResult -from taskito.serializers import CloudpickleSerializer, Serializer +from taskito.serializers import Serializer, SmartSerializer from taskito.webhooks import WebhookManager if TYPE_CHECKING: @@ -147,7 +147,8 @@ def __init__( default_priority: Default task priority (higher = more urgent). result_ttl: Auto-cleanup completed/dead jobs older than this many seconds. None disables auto-cleanup. - serializer: Serializer for task payloads. Defaults to CloudpickleSerializer. + serializer: Serializer for task payloads. Defaults to SmartSerializer + (msgpack with cloudpickle fallback). middleware: List of global middleware instances applied to all tasks. backend: Storage backend — ``"sqlite"`` (default) or ``"postgres"``. db_url: PostgreSQL connection URL (required when backend is ``"postgres"``). @@ -218,7 +219,7 @@ def __init__( "on_success": [], "on_failure": [], } - self._serializer: Serializer = serializer or CloudpickleSerializer() + self._serializer: Serializer = serializer or SmartSerializer() self._task_serializers: dict[str, Serializer] = {} self._task_idempotent: dict[str, bool] = {} self._task_compensates: dict[str, str] = {} @@ -386,8 +387,15 @@ def close(self) -> None: self._batch_accumulator = None def _deserialize_payload(self, task_name: str, payload: bytes) -> tuple: - """Deserialize a job payload using the per-task or queue-level serializer.""" - return self._get_serializer(task_name).loads(payload) # type: ignore[no-any-return] + """Deserialize a job payload into an ``(args, kwargs)`` pair. + + Serializers without a native tuple type (msgpack, JSON) flatten the + payload to ``[args, kwargs]`` with ``args`` as a list. The worker hands + ``args`` straight to PyO3, which requires a real tuple, so the shape is + normalized back to ``(tuple, dict)`` here regardless of serializer. + """ + args, kwargs = self._get_serializer(task_name).loads(payload) + return tuple(args), kwargs def enqueue( self, From d8e66acc7975553b606f9210b41857eb79db6a49 Mon Sep 17 00:00:00 2001 From: Pratyush Sharma <56130065+pratyush618@users.noreply.github.com> Date: Sat, 30 May 2026 09:24:16 +0530 Subject: [PATCH 3/3] fix(serializers): keep namedtuples out of the tuple ExtType path --- py_src/taskito/serializers.py | 8 +++++--- tests/core/test_serializers.py | 14 ++++++++++++++ 2 files changed, 19 insertions(+), 3 deletions(-) diff --git a/py_src/taskito/serializers.py b/py_src/taskito/serializers.py index 7195e124..51c8f95d 100644 --- a/py_src/taskito/serializers.py +++ b/py_src/taskito/serializers.py @@ -23,10 +23,12 @@ def _msgpack_default(obj: Any) -> Any: """Encode types msgpack can't represent natively. - Tuples become a tagged ExtType (recursively packed). Anything else raises, - which ``SmartSerializer`` catches to fall back to cloudpickle. + Exact tuples become a tagged ExtType (recursively packed). Tuple + *subclasses* (e.g. namedtuples) are not exact tuples — ``list(obj)`` would + drop their type — so they raise here and fall back to cloudpickle, which + preserves them. Anything else also raises and falls back. """ - if isinstance(obj, tuple): + if type(obj) is tuple: return msgpack.ExtType(_EXT_TUPLE, _msgpack_packb(list(obj))) raise TypeError(f"Cannot msgpack-encode {type(obj).__name__}") diff --git a/tests/core/test_serializers.py b/tests/core/test_serializers.py index 2bcb1616..f5da3935 100644 --- a/tests/core/test_serializers.py +++ b/tests/core/test_serializers.py @@ -162,6 +162,20 @@ def test_plain_data_uses_msgpack_tag(self) -> None: assert encoded[:1] == b"\x01" assert s.loads(encoded) == {"x": 1, "y": ["a", "b"]} + def test_namedtuple_preserved_via_cloudpickle(self) -> None: + """Tuple subclasses (namedtuples) must keep their type, not be + demoted to a plain tuple by the ExtType path.""" + from collections import namedtuple + + Point = namedtuple("Point", ["x", "y"]) + s = SmartSerializer() + encoded = s.dumps(Point(1, 2)) + assert encoded[:1] == b"\x00" # cloudpickle fallback, not msgpack + restored = s.loads(encoded) + assert restored == Point(1, 2) + assert type(restored) is Point + assert restored.x == 1 + def test_falls_back_to_cloudpickle_for_lambda(self) -> None: s = SmartSerializer() encoded = s.dumps(lambda x: x + 1)