make mypy happy with mixin classes

Author: aaronsteers

airbyte_cdk/models/airbyte_protocol.py

@@ -1,7 +1,15 @@
 # Copyright (c) 2025 Airbyte, Inc., all rights reserved.
-from collections.abc import Callable
+from collections.abc import Callable, Mapping
 from dataclasses import InitVar, dataclass
-from typing import Annotated, Any, Dict, List, Mapping, Optional, Union
+from typing import (
+    Annotated,
+    Any,
+    Dict,
+    List,
+    Optional,
+    Type,
+    TypeVar,
+)
 
 import orjson
 from airbyte_protocol_dataclasses.models import *  # noqa: F403  # Allow '*'
@@ -48,21 +56,74 @@ def __eq__(self, other: object) -> bool:
         )
 
 
+T = TypeVar("T", bound="SerDeMixin")
+
+
+class SerDeMixin:
+    # allow subclasses to override their resolver if they need one
+    _type_resolver: Callable[[type], CustomType[Any, Any] | None] | None = None
+
+    def __init_subclass__(cls, **kwargs):
+        super().__init_subclass__(**kwargs)
+        # build a Serializer *once* for each subclass
+        cls._serializer = Serializer(
+            cls,
+            omit_none=True,
+            custom_type_resolver=cls._type_resolver,
+        )
+
+    def to_dict(self) -> Dict[str, Any]:
+        return self._serializer.dump(self)
+
+    def to_json(self) -> str:
+        # use to_dict so you only have one canonical dump
+        return orjson.dumps(self.to_dict()).decode("utf-8")
+
+    @classmethod
+    def from_dict(cls: type[T], data: Dict[str, Any]) -> T:
+        return cls._serializer.load(data)
+
+    @classmethod
+    def from_json(cls: type[T], s: str) -> T:
+        return cls._serializer.load(orjson.loads(s))
+
+
+def _custom_state_resolver(t: type) -> CustomType[AirbyteStateBlob, dict[str, Any]] | None:
+    class AirbyteStateBlobType(CustomType[AirbyteStateBlob, Dict[str, Any]]):
+        def serialize(self, value: AirbyteStateBlob) -> Dict[str, Any]:
+            # cant use orjson.dumps() directly because private attributes are excluded, e.g. "__ab_full_refresh_sync_complete"
+            return {k: v for k, v in value.__dict__.items()}
+
+        def deserialize(self, value: Dict[str, Any]) -> AirbyteStateBlob:
+            return AirbyteStateBlob(value)
+
+        def get_json_schema(self) -> Dict[str, Any]:
+            return {"type": "object"}
+
+    return AirbyteStateBlobType() if t is AirbyteStateBlob else None
+
+
 # The following dataclasses have been redeclared to include the new version of AirbyteStateBlob
 @dataclass
-class AirbyteStreamState:
+class AirbyteStreamState(SerDeMixin):
     stream_descriptor: StreamDescriptor  # type: ignore [name-defined]
     stream_state: Optional[AirbyteStateBlob] = None
 
+    # override the resolver for AirbyteStreamState to use the custom one
+    _type_resolver = _custom_state_resolver
+
 
 @dataclass
-class AirbyteGlobalState:
+class AirbyteGlobalState(SerDeMixin):
     stream_states: List[AirbyteStreamState]
     shared_state: Optional[AirbyteStateBlob] = None
 
+    # override the resolver for AirbyteStreamState to use the custom one
+    _type_resolver = _custom_state_resolver
+
 
 @dataclass
-class AirbyteStateMessage:
+class AirbyteStateMessage(SerDeMixin):
     type: Optional[AirbyteStateType] = None  # type: ignore [name-defined]
     stream: Optional[AirbyteStreamState] = None
     global_: Annotated[AirbyteGlobalState | None, Alias("global")] = (
@@ -72,9 +133,12 @@ class AirbyteStateMessage:
     sourceStats: Optional[AirbyteStateStats] = None  # type: ignore [name-defined]
     destinationStats: Optional[AirbyteStateStats] = None  # type: ignore [name-defined]
 
+    # override the resolver for AirbyteStreamState to use the custom one
+    _type_resolver = _custom_state_resolver
+
 
 @dataclass
-class AirbyteMessage:
+class AirbyteMessage(SerDeMixin):
     type: Type  # type: ignore [name-defined]
     log: Optional[AirbyteLogMessage] = None  # type: ignore [name-defined]
     spec: Optional[ConnectorSpecification] = None  # type: ignore [name-defined]
@@ -85,42 +149,18 @@ class AirbyteMessage:
     trace: Optional[AirbyteTraceMessage] = None  # type: ignore [name-defined]
     control: Optional[AirbyteControlMessage] = None  # type: ignore [name-defined]
 
+    # override the resolver for AirbyteStreamState to use the custom one
+    _type_resolver = _custom_state_resolver
 
-# Add optimized serdes methods to the protocol data classes:
 
-def _with_serdes(
-    cls,
-    type_resolver: Callable[[type], CustomType[Any, Any] | None] | None = None,
-) -> type:
-    """Decorator to add SerDes (serialize/deserialize) methods to a dataclass."""
-    cls._serializer = Serializer(cls, omit_none=True, custom_type_resolver=type_resolver)
-    cls.to_dict = lambda self: self._serializer.dump(self)
-    cls.to_json = lambda self: orjson.dumps(self._serializer.dump(self)).decode("utf-8")
-    cls.from_json = lambda self, string: self._serializer.load(orjson.loads(string))
-    cls.from_dict = lambda self, dictionary: self._serializer.load(dictionary)
-    return cls
-
-
-def _custom_state_resolver(t: type) -> CustomType[AirbyteStateBlob, dict[str, Any]] | None:
-    class AirbyteStateBlobType(CustomType[AirbyteStateBlob, Dict[str, Any]]):
-        def serialize(self, value: AirbyteStateBlob) -> Dict[str, Any]:
-            # cant use orjson.dumps() directly because private attributes are excluded, e.g. "__ab_full_refresh_sync_complete"
-            return {k: v for k, v in value.__dict__.items()}
-
-        def deserialize(self, value: Dict[str, Any]) -> AirbyteStateBlob:
-            return AirbyteStateBlob(value)
+# These don't need the custom resolver:
+class ConnectorSpecification(ConnectorSpecification, SerDeMixin):
+    pass
 
-        def get_json_schema(self) -> Dict[str, Any]:
-            return {"type": "object"}
 
-    return AirbyteStateBlobType() if t is AirbyteStateBlob else None
+class ConfiguredAirbyteCatalog(ConfiguredAirbyteCatalog, SerDeMixin):
+    pass
 
 
-# Add serdes capabilities to all data classes that need to serialize and deserialize:
-AirbyteMessage = _with_serdes(AirbyteMessage, _custom_state_resolver)
-AirbyteStateMessage = _with_serdes(AirbyteStateMessage, _custom_state_resolver)
-AirbyteStreamState = _with_serdes(AirbyteStreamState, _custom_state_resolver)
-# These don't need the custom resolver:
-ConnectorSpecification = _with_serdes(ConnectorSpecification)
-ConfiguredAirbyteCatalog = _with_serdes(ConfiguredAirbyteCatalog)
-AirbyteStream = _with_serdes(AirbyteStream)
+class AirbyteStream(AirbyteStream, SerDeMixin):
+    pass