autogen_core中的DataclassJsonMessageSerializer类

源代码

import json
from dataclasses import asdict, dataclass, fields
from typing import Any, ClassVar, Dict, List, Protocol, Sequence, TypeVar, cast, get_args, get_origin, runtime_checkable, Union
from pydantic import BaseModelfrom types import NoneType, UnionTypedef _type_name(cls: type[Any] | Any) -> str:if isinstance(cls, type):return cls.__name__else:return cast(str, cls.__class__.__name__)def is_union(t: object) -> bool:origin = get_origin(t)return origin is Union or origin is UnionTypeT = TypeVar("T")class MessageSerializer(Protocol[T]):@propertydef data_content_type(self) -> str: ...@propertydef type_name(self) -> str: ...def deserialize(self, payload: bytes) -> T: ...def serialize(self, message: T) -> bytes: ...@runtime_checkable
class IsDataclass(Protocol):# as already noted in comments, checking for this attribute is currently# the most reliable way to ascertain that something is a dataclass__dataclass_fields__: ClassVar[Dict[str, Any]]def is_dataclass(cls: type[Any]) -> bool:return hasattr(cls, "__dataclass_fields__")def has_nested_dataclass(cls: type[IsDataclass]) -> bool:# iterate fields and check if any of them are dataclassesreturn any(is_dataclass(f.type) for f in cls.__dataclass_fields__.values())def contains_a_union(cls: type[IsDataclass]) -> bool:return any(is_union(f.type) for f in cls.__dataclass_fields__.values())def has_nested_base_model(cls: type[IsDataclass]) -> bool:for f in fields(cls):field_type = f.type# Resolve forward references and other annotationsorigin = get_origin(field_type)args = get_args(field_type)# If the field type is directly a subclass of BaseModelif isinstance(field_type, type) and issubclass(field_type, BaseModel):return True# If the field type is a generic type like List[BaseModel], Tuple[BaseModel, ...], etc.if origin is not None and args:for arg in args:# Recursively check the argument typesif isinstance(arg, type) and issubclass(arg, BaseModel):return Trueelif get_origin(arg) is not None:# Handle nested generics like List[List[BaseModel]]if has_nested_base_model_in_type(arg):return True# Handle Union typeselif args:for arg in args:if isinstance(arg, type) and issubclass(arg, BaseModel):return Trueelif get_origin(arg) is not None:if has_nested_base_model_in_type(arg):return Truereturn Falsedef has_nested_base_model_in_type(tp: Any) -> bool:"""Helper function to check if a type or its arguments is a BaseModel subclass."""origin = get_origin(tp)args = get_args(tp)if isinstance(tp, type) and issubclass(tp, BaseModel):return Trueif origin is not None and args:for arg in args:if has_nested_base_model_in_type(arg):return Truereturn FalseDataclassT = TypeVar("DataclassT", bound=IsDataclass)JSON_DATA_CONTENT_TYPE = "application/json"
"""JSON data content type"""# TODO: what's the correct content type? There seems to be some disagreement over what it should be
PROTOBUF_DATA_CONTENT_TYPE = "application/x-protobuf"
"""Protobuf data content type"""class DataclassJsonMessageSerializer(MessageSerializer[DataclassT]):def __init__(self, cls: type[DataclassT]) -> None:if contains_a_union(cls):raise ValueError("Dataclass has a union type, which is not supported. To use a union, use a Pydantic model")if has_nested_dataclass(cls) or has_nested_base_model(cls):raise ValueError("Dataclass has nested dataclasses or base models, which are not supported. To use nested types, use a Pydantic model")self.cls = cls@propertydef data_content_type(self) -> str:return JSON_DATA_CONTENT_TYPE@propertydef type_name(self) -> str:return _type_name(self.cls)def deserialize(self, payload: bytes) -> DataclassT:message_str = payload.decode("utf-8")return self.cls(**json.loads(message_str))def serialize(self, message: DataclassT) -> bytes:return json.dumps(asdict(message)).encode("utf-8")

代码解释

这段代码定义了一个用于序列化和反序列化数据类（dataclass）和 Pydantic 模型的序列化器，特别是针对 JSON 格式。它的目标是提供一种结构化的方式来序列化和反序列化消息，并检查不支持的特性，例如嵌套的数据类、嵌套的 Pydantic 模型以及数据类中联合类型（union type）。

1. 导入

import json
from dataclasses import asdict, dataclass, fields
from typing import Any, ClassVar, Dict, List, Protocol, Sequence, TypeVar, cast, get_args, get_origin, runtime_checkable, Union
from pydantic import BaseModelfrom types import NoneType, UnionType

这部分导入了必要的模块：

• json: 用于 JSON 序列化/反序列化。
• dataclasses: 用于处理数据类 (asdict, dataclass, fields)。
• typing: 用于类型提示和处理泛型 (Any, ClassVar, Dict, List, Protocol, Sequence, TypeVar, cast, get_args, get_origin, runtime_checkable, Union)。
• pydantic: 用于处理 Pydantic 模型 (BaseModel)。
• types: 用于处理类型提示，包括 NoneType 和 UnionType。

2. 辅助函数

• _type_name(cls): 返回类或对象的名称。
• is_union(t): 检查类型 t 是否为 Union 或 UnionType。
• is_dataclass(cls): 检查类 cls 是否为数据类。
• has_nested_dataclass(cls): 检查数据类 cls 是否包含任何嵌套的数据类作为字段。
• contains_a_union(cls): 检查数据类 cls 是否包含任何 Union 类型作为字段类型。
• has_nested_base_model(cls): 检查数据类 cls 是否包含任何嵌套的 Pydantic BaseModel 子类作为字段，包括处理 List[BaseModel] 等泛型和嵌套泛型。
• has_nested_base_model_in_type(tp): has_nested_base_model 使用的辅助函数，用于递归检查类型或其参数是否为 BaseModel 子类。

3. 类型变量和常量

• T: 一个泛型类型变量。
• DataclassT: 绑定到 IsDataclass 的类型变量。
• JSON_DATA_CONTENT_TYPE: JSON 内容类型的常量。
• PROTOBUF_DATA_CONTENT_TYPE: Protobuf 内容类型的常量（虽然在 JSON 序列化器中未使用）。

4. MessageSerializer 协议

class MessageSerializer(Protocol[T]):@propertydef data_content_type(self) -> str: ...@propertydef type_name(self) -> str: ...def deserialize(self, payload: bytes) -> T: ...def serialize(self, message: T) -> bytes: ...

定义了消息序列化器的协议（接口）。任何实现此协议的类都必须定义指定的属性和方法：

• data_content_type: 序列化数据的 MIME 类型。
• type_name: 被序列化类型的名称。
• deserialize(payload): 将字节负载反序列化为 T 类型的对象。
• serialize(message): 将 T 类型的对象序列化为字节负载。

5. IsDataclass 协议

@runtime_checkable
class IsDataclass(Protocol):__dataclass_fields__: ClassVar[Dict[str, Any]]

一个运行时可检查的协议，用于确定类是否为数据类。它检查 __dataclass_fields__ 属性是否存在。

6. DataclassJsonMessageSerializer 类

class DataclassJsonMessageSerializer(MessageSerializer[DataclassT]):# ...

此类实现了 MessageSerializer 协议，用于数据类的 JSON 序列化。

• __init__(self, cls): 构造函数，它接受数据类类型 cls 作为参数。它执行关键检查：
  • 如果数据类包含 Union 类型，则引发 ValueError（使用 Pydantic 处理联合类型）。
  • 如果数据类包含嵌套的数据类或 Pydantic 模型，则引发 ValueError（使用 Pydantic 处理嵌套类型）。
• data_content_type 属性：返回 JSON_DATA_CONTENT_TYPE。
• type_name 属性：返回数据类的名称。
• deserialize(payload): 将 JSON 负载反序列化为数据类的实例。
• serialize(message): 将数据类实例序列化为 JSON 负载。

总结

这段代码提供了一个专门用于简单数据类（没有联合或嵌套的数据类/模型）的 JSON 序列化器。它强调使用 Pydantic 来处理更复杂的场景，从而实现了清晰的关注点分离，并使序列化过程更具可预测性。这是一种很好的方法，可以避免直接在数据类中泛型序列化嵌套结构和联合类型的复杂性。

几个例子

@dataclass
class Point:x: inty: intserializer = DataclassJsonMessageSerializer(Point)point = Point(x=10, y=20)
serialized_data = serializer.serialize(point)
print(f"Serialized data: {serialized_data}")  # Output: b'{"x": 10, "y": 20}'deserialized_point = serializer.deserialize(serialized_data)
print(f"Deserialized point: {deserialized_point}")  # Output: Point(x=10, y=20)print(isinstance(deserialized_point, Point)) # Output: True

Serialized data: b'{"x": 10, "y": 20}'
Deserialized point: Point(x=10, y=20)
True

@dataclass
class Line:points: list[Point]line = Line(points=[Point(x=1, y=2), Point(x=3, y=4)])
serializer_line = DataclassJsonMessageSerializer(Line)serialized_line = serializer_line.serialize(line)
print(f"Serialized line: {serialized_line}") # Output: b'{"points": [{"x": 1, "y": 2}, {"x": 3, "y": 4}]}'deserialized_line = serializer_line.deserialize(serialized_line)
print(f"Deserialized line: {deserialized_line}") # Output: Line(points=[Point(x=1, y=2), Point(x=3, y=4)])print(isinstance(deserialized_line, Line)) # Output: True

Serialized line: b'{"points": [{"x": 1, "y": 2}, {"x": 3, "y": 4}]}'
Deserialized line: Line(points=[{'x': 1, 'y': 2}, {'x': 3, 'y': 4}])
True

@dataclass
class Rectangle:top_left: Pointbottom_right: Point# 这会抛出 ValueError，因为 Rectangle 包含嵌套的数据类 Point
try:serializer_rectangle = DataclassJsonMessageSerializer(Rectangle)
except ValueError as e:print(f"Error: {e}")  # Output: Error: Dataclass has nested dataclasses or base models, which are not supported. To use nested types, use a Pydantic model

Error: Dataclass has nested dataclasses or base models, which are not supported. To use nested types, use a Pydantic model

from typing import Union@dataclass
class Shape:shape_type: Union[str, int]try:serializer_shape = DataclassJsonMessageSerializer(Shape)
except ValueError as e:print(f"Error: {e}") # Output: Error: Dataclass has a union type, which is not supported. To use a union, use a Pydantic model

Error: Dataclass has a union type, which is not supported. To use a union, use a Pydantic model