Coverage for async-durable-execution/src/async_durable_execution/serdes.py: 100%
318 statements
« prev ^ index » next coverage.py v7.15.0, created at 2026-07-07 16:54 +0000
1"""Codec-based serialization and deserialization for Python types.
3This module provides comprehensive serialization support using a codec-based
4architecture with recursive encoding/decoding for nested structures.
6Key Features:
7- Plain JSON for primitives and simple lists (performance optimization)
8- Envelope format with type tags for complex types
9- Modular codec architecture
10- Recursive handling of nested structures
12Serialization Strategy:
13- Primitives (None, str, int, float, bool): Plain JSON
14- Simple lists containing only primitives: Plain JSON
15- Everything else: Envelope format with type tags
17Wire Formats:
18 Plain JSON: 42, "hello", [1, 2, 3]
19 Envelope: {"t": "<type_tag>", "v": <encoded_value>}
20"""
22from __future__ import annotations
24import base64
25import json
26import logging
27import uuid
28from abc import ABC, abstractmethod
29from dataclasses import dataclass
30from datetime import date, datetime
31from decimal import Decimal
32from enum import Enum
33from typing import TYPE_CHECKING, Any, Generic, Protocol, TypeVar, cast
35from .context import bind_current_context
36from .exceptions import (
37 DurableExecutionsError,
38 ExecutionError,
39 SerDesError,
40)
42if TYPE_CHECKING:
43 from .composite.parallel import BatchResult
46logger = logging.getLogger(__name__)
48T = TypeVar("T")
50TYPE_TOKEN: str = "t"
51VALUE_TOKEN: str = "v"
54def _get_batch_result_type() -> type[BatchResult]:
55 from .composite.parallel import BatchResult
57 return BatchResult
60class TypeTag(str, Enum):
61 """Type tags for envelope format."""
63 NONE = "n"
64 STR = "s"
65 INT = "i"
66 FLOAT = "f"
67 BOOL = "b"
68 BYTES = "B"
69 UUID = "u"
70 DECIMAL = "d"
71 DATETIME = "dt"
72 DATE = "D"
73 TUPLE = "t"
74 LIST = "l"
75 DICT = "m"
76 BATCH_RESULT = "br"
79@dataclass(frozen=True)
80class EncodedValue:
81 """Encoded value with type tag."""
83 tag: TypeTag
85 value: Any
88class Codec(Protocol):
89 """Protocol for type-specific codecs."""
91 def encode(self, obj: Any) -> EncodedValue: ...
93 def decode(self, tag: TypeTag, value: Any) -> Any: ...
96class PrimitiveCodec:
97 """Codec for primitive types."""
99 def encode(self, obj: Any) -> EncodedValue: # noqa: PLR6301
100 match obj:
101 case None:
102 return EncodedValue(TypeTag.NONE, None)
103 case str():
104 return EncodedValue(TypeTag.STR, obj)
105 case bool(): # Must come before int
106 return EncodedValue(TypeTag.BOOL, obj)
107 case int():
108 return EncodedValue(TypeTag.INT, obj)
109 case float():
110 return EncodedValue(TypeTag.FLOAT, obj)
111 case _:
112 msg = f"Unsupported primitive type: {type(obj)!r}"
113 raise SerDesError(msg)
115 def decode(self, tag: TypeTag, value: Any) -> Any: # noqa: PLR6301
116 match tag:
117 case TypeTag.NONE:
118 return None
119 case TypeTag.STR:
120 return str(value)
121 case TypeTag.BOOL:
122 return bool(value)
123 case TypeTag.INT:
124 return int(value)
125 case TypeTag.FLOAT:
126 return float(value)
127 case _:
128 msg = f"Unknown primitive tag: {tag}"
129 raise SerDesError(msg)
132class BytesCodec:
133 """Codec for bytes, bytearray, and memoryview."""
135 def encode(self, obj: Any) -> EncodedValue: # noqa: PLR6301
136 encoded = base64.b64encode(bytes(obj)).decode("utf-8")
137 return EncodedValue(TypeTag.BYTES, encoded)
139 def decode(self, tag: TypeTag, value: Any) -> Any: # noqa: PLR6301
140 if tag != TypeTag.BYTES:
141 msg = f"Expected BYTES tag, got {tag}"
142 raise SerDesError(msg)
143 return base64.b64decode(value.encode("utf-8"))
146class UuidCodec:
147 """Codec for UUID objects."""
149 def encode(self, obj: Any) -> EncodedValue: # noqa: PLR6301
150 return EncodedValue(TypeTag.UUID, str(obj))
152 def decode(self, tag: TypeTag, value: Any) -> Any: # noqa: PLR6301
153 if tag != TypeTag.UUID:
154 msg = f"Expected UUID tag, got {tag}"
155 raise SerDesError(msg)
156 return uuid.UUID(value)
159class DecimalCodec:
160 """Codec for Decimal objects."""
162 def encode(self, obj: Any) -> EncodedValue: # noqa: PLR6301
163 return EncodedValue(TypeTag.DECIMAL, str(obj))
165 def decode(self, tag: TypeTag, value: Any) -> Any: # noqa: PLR6301
166 if tag != TypeTag.DECIMAL:
167 msg = f"Expected DECIMAL tag, got {tag}"
168 raise SerDesError(msg)
169 return Decimal(value)
172class DateTimeCodec:
173 """Codec for datetime and date objects."""
175 def encode(self, obj: Any) -> EncodedValue: # noqa: PLR6301
176 match obj:
177 case datetime():
178 return EncodedValue(TypeTag.DATETIME, obj.isoformat())
179 case date():
180 return EncodedValue(TypeTag.DATE, obj.isoformat())
181 case _:
182 msg = f"Unsupported datetime type: {type(obj)!r}"
183 raise SerDesError(msg)
185 def decode(self, tag: TypeTag, value: Any) -> Any: # noqa: PLR6301
186 match tag:
187 case TypeTag.DATETIME:
188 # Handle Z suffix for UTC
189 s = value
190 if isinstance(s, str) and s.endswith("Z"):
191 s = s[:-1] + "+00:00"
192 return datetime.fromisoformat(s)
193 case TypeTag.DATE:
194 return date.fromisoformat(value)
195 case _:
196 msg = f"Unknown datetime tag: {tag}"
197 raise SerDesError(msg)
200class ContainerCodec(Codec):
201 """Codec for container types with recursive encoding/decoding."""
203 def __init__(self) -> None:
204 self._dispatcher: TypeCodec | None = None
206 def set_dispatcher(self, dispatcher) -> None:
207 """Set the main codec dispatcher for recursive encoding."""
208 self._dispatcher = dispatcher
210 @property
211 def dispatcher(self):
212 """Get the dispatcher, raising error if not set."""
213 if self._dispatcher is None:
214 msg = "ContainerCodec not linked to a TypeCodec dispatcher."
215 raise DurableExecutionsError(msg)
216 return self._dispatcher
218 def encode(self, obj: Any) -> EncodedValue:
219 """Encode container using dispatcher for recursive elements."""
221 match obj:
222 case obj if isinstance(obj, _get_batch_result_type()):
223 # Encode BatchResult as dict with special tag
224 return EncodedValue(
225 TypeTag.BATCH_RESULT,
226 self._wrap(obj.to_dict(), self.dispatcher).value,
227 )
228 case list():
229 return EncodedValue(
230 TypeTag.LIST, [self._wrap(v, self.dispatcher) for v in obj]
231 )
232 case tuple():
233 return EncodedValue(
234 TypeTag.TUPLE, [self._wrap(v, self.dispatcher) for v in obj]
235 )
236 case dict():
237 for k in obj:
238 if isinstance(k, tuple):
239 msg = "Tuple keys not supported"
240 raise SerDesError(msg)
241 return EncodedValue(
242 TypeTag.DICT,
243 {k: self._wrap(v, self.dispatcher) for k, v in obj.items()},
244 )
245 case _:
246 msg = f"Unsupported container type: {type(obj)!r}"
247 raise SerDesError(msg)
249 def decode(self, tag: TypeTag, value: Any) -> Any:
250 """Decode container using dispatcher for recursive elements."""
252 match tag:
253 case TypeTag.BATCH_RESULT:
254 # Decode BatchResult from dict - value is already the dict structure
255 # First decode it as a dict to unwrap all nested EncodedValues
256 decoded_dict = self.decode(TypeTag.DICT, value)
257 return _get_batch_result_type().from_dict(decoded_dict)
258 case TypeTag.LIST:
259 if not isinstance(value, list):
260 msg = f"Expected list, got {type(value)}"
261 raise SerDesError(msg)
262 return [self._unwrap(v, self.dispatcher) for v in value]
263 case TypeTag.TUPLE:
264 if not isinstance(value, list):
265 msg = f"Expected list, got {type(value)}"
266 raise SerDesError(msg)
267 return tuple(self._unwrap(v, self.dispatcher) for v in value)
268 case TypeTag.DICT:
269 if not isinstance(value, dict):
270 msg = f"Expected dict, got {type(value)}"
271 raise SerDesError(msg)
272 return {k: self._unwrap(v, self.dispatcher) for k, v in value.items()}
273 case _:
274 msg = f"Unknown container tag: {tag}"
275 raise SerDesError(msg)
277 @staticmethod
278 def _wrap(obj: Any, dispatcher) -> EncodedValue:
279 """Wrap object using dispatcher."""
280 return dispatcher.encode(obj)
282 @staticmethod
283 def _unwrap(obj: Any, dispatcher) -> Any:
284 """Unwrap object using dispatcher."""
285 match obj:
286 case EncodedValue():
287 return dispatcher.decode(obj.tag, obj.value)
288 case dict() if TYPE_TOKEN in obj and VALUE_TOKEN in obj:
289 tag = TypeTag(obj[TYPE_TOKEN])
290 return dispatcher.decode(tag, obj[VALUE_TOKEN])
291 case _:
292 return obj
295class TypeCodec(Codec):
296 """Main codec dispatcher."""
298 def __init__(self):
299 self.primitive_codec = PrimitiveCodec()
300 self.bytes_codec = BytesCodec()
301 self.uuid_codec = UuidCodec()
302 self.decimal_codec = DecimalCodec()
303 self.datetime_codec = DateTimeCodec()
304 self.container_codec = ContainerCodec()
305 self.container_codec.set_dispatcher(self)
307 def encode(self, obj: Any) -> EncodedValue:
308 match obj:
309 case None | str() | bool() | int() | float():
310 return self.primitive_codec.encode(obj)
311 case bytes() | bytearray() | memoryview():
312 return self.bytes_codec.encode(bytes(obj))
313 case uuid.UUID():
314 return self.uuid_codec.encode(obj)
315 case Decimal():
316 return self.decimal_codec.encode(obj)
317 case datetime() | date():
318 return self.datetime_codec.encode(obj)
319 case obj if isinstance(obj, list | tuple | dict) or isinstance(
320 obj, _get_batch_result_type()
321 ):
322 return self.container_codec.encode(obj)
323 case _:
324 msg = f"Unsupported type: {type(obj)}"
325 raise SerDesError(msg)
327 def decode(self, tag: TypeTag, value: Any) -> Any:
328 match tag:
329 case (
330 TypeTag.NONE | TypeTag.STR | TypeTag.BOOL | TypeTag.INT | TypeTag.FLOAT
331 ):
332 return self.primitive_codec.decode(tag, value)
333 case TypeTag.BYTES:
334 return self.bytes_codec.decode(tag, value)
335 case TypeTag.UUID:
336 return self.uuid_codec.decode(tag, value)
337 case TypeTag.DECIMAL:
338 return self.decimal_codec.decode(tag, value)
339 case TypeTag.DATETIME | TypeTag.DATE:
340 return self.datetime_codec.decode(tag, value)
341 case TypeTag.LIST | TypeTag.TUPLE | TypeTag.DICT | TypeTag.BATCH_RESULT:
342 return self.container_codec.decode(tag, value)
343 case _:
344 msg = f"Unknown type tag: {tag}"
345 raise SerDesError(msg)
348TYPE_CODEC = TypeCodec()
351@dataclass(frozen=True)
352class SerDesContext:
353 """Context for serialization operations."""
355 operation_id: str = ""
357 durable_execution_arn: str = ""
360class SerDes(ABC, Generic[T]):
361 """Abstract serializer interface for durable operation payloads and results."""
363 @abstractmethod
364 async def serialize(self, value: T) -> str:
365 """Convert a Python value into the wire format stored by the SDK."""
366 pass
368 @abstractmethod
369 async def deserialize(self, data: str) -> T:
370 """Reconstruct a Python value from the durable wire format."""
371 pass
373 @staticmethod
374 def is_primitive(obj: Any) -> bool:
375 """Check if object contains only JSON-serializable primitives."""
376 if obj is None or isinstance(obj, str | int | float | bool):
377 return True
378 if isinstance(obj, list):
379 return all(SerDes.is_primitive(item) for item in obj)
380 return False
383class PassThroughSerDes(SerDes[T]):
384 """Serializer that leaves already-serialized string payloads unchanged."""
386 async def serialize(self, value: T) -> str: # noqa: PLR6301
387 return cast("str", value)
389 async def deserialize(self, data: str) -> T: # noqa: PLR6301
390 return cast("T", data)
393class JsonSerDes(SerDes[T]):
394 """Serializer that uses the standard library `json` module."""
396 async def serialize(self, value: T) -> str: # noqa: PLR6301
397 return json.dumps(value)
399 async def deserialize(self, data: str) -> T: # noqa: PLR6301
400 return json.loads(data)
403class ExtendedTypeSerDes(SerDes[T]):
404 """Main serializer class."""
406 def __init__(self):
407 self._codec = TYPE_CODEC
409 async def serialize(self, value: Any) -> str:
410 """Serialize value to JSON string."""
411 return self.serialize_sync(value)
413 async def deserialize(self, data: str) -> Any:
414 """Deserialize JSON string to Python object."""
415 return self.deserialize_sync(data)
417 def serialize_sync(self, value: Any) -> str:
418 """Serialize value to JSON string without awaiting."""
419 # Fast path for primitives
420 if SerDes.is_primitive(value):
421 return json.dumps(value, separators=(",", ":"))
423 encoded = self._codec.encode(value)
424 wrapped = self._to_json_serializable(encoded)
425 return json.dumps(wrapped, separators=(",", ":"))
427 def deserialize_sync(self, data: str) -> Any:
428 """Deserialize JSON string to Python object without awaiting."""
429 obj = json.loads(data)
431 # Fast path for primitives
432 if SerDes.is_primitive(obj):
433 return obj
435 if not (isinstance(obj, dict) and TYPE_TOKEN in obj and VALUE_TOKEN in obj):
436 msg = 'Malformed envelope: missing "t" or "v" at root.'
437 raise SerDesError(msg)
438 # Python 3.11 compatibility: Using try-except instead of 'in' operator
439 # because checking 'str in EnumType' raises TypeError in Python 3.11
440 try:
441 tag = TypeTag(obj[TYPE_TOKEN])
442 except ValueError:
443 msg = f'Unknown type tag: "{obj[TYPE_TOKEN]}"'
444 raise SerDesError(msg) from None
446 return self._codec.decode(tag, obj[VALUE_TOKEN])
448 def _to_json_serializable(self, obj: Any) -> Any:
449 """Convert EncodedValue objects to JSON-serializable format."""
450 match obj:
451 case EncodedValue():
452 return {
453 TYPE_TOKEN: obj.tag,
454 VALUE_TOKEN: self._to_json_serializable(obj.value),
455 }
456 case list():
457 return [self._to_json_serializable(x) for x in obj]
458 case dict():
459 return {k: self._to_json_serializable(v) for k, v in obj.items()}
460 case _:
461 return obj
464DEFAULT_JSON_SERDES: SerDes[Any] = JsonSerDes()
465EXTENDED_TYPES_SERDES: SerDes[Any] = ExtendedTypeSerDes()
468async def serialize(
469 serdes: SerDes[T] | None, value: T, operation_id: str, durable_execution_arn: str
470) -> str:
471 """Serialize value using provided or default serializer.
473 Args:
474 serdes: Custom serializer or None for default
475 value: Object to serialize
476 operation_id: Unique operation identifier
477 durable_execution_arn: ARN of durable execution
479 Returns:
480 Serialized string representation
482 Raises:
483 FatalError: If serialization fails
484 """
485 serdes_context: SerDesContext = SerDesContext(operation_id, durable_execution_arn)
486 active_serdes: SerDes[T] = serdes or EXTENDED_TYPES_SERDES
488 async def serialize_value() -> str:
489 return await active_serdes.serialize(value)
491 try:
492 with bind_current_context(serdes_context):
493 return await serialize_value()
494 except Exception as e:
495 logger.exception(
496 "⚠️ Serialization failed for id: %s",
497 operation_id,
498 )
499 msg = f"Serialization failed for id: {operation_id}, error: {e}."
500 raise ExecutionError(msg) from e
503async def deserialize(
504 serdes: SerDes[T] | None, data: str, operation_id: str, durable_execution_arn: str
505) -> T:
506 """Deserialize data using provided or default serializer.
508 Args:
509 serdes: Custom serializer or None for default
510 data: Serialized string data
511 operation_id: Unique operation identifier
512 durable_execution_arn: ARN of durable execution
514 Returns:
515 Deserialized Python object
517 Raises:
518 FatalError: If deserialization fails
519 """
520 serdes_context: SerDesContext = SerDesContext(operation_id, durable_execution_arn)
521 active_serdes: SerDes[T] = serdes or EXTENDED_TYPES_SERDES
523 async def deserialize_value() -> T:
524 return await active_serdes.deserialize(data)
526 try:
527 with bind_current_context(serdes_context):
528 return await deserialize_value()
529 except Exception as e:
530 logger.exception("⚠️ Deserialization failed for id: %s", operation_id)
531 msg = f"Deserialization failed for id: {operation_id}"
532 raise ExecutionError(msg) from e