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 BaseModel

from types import NoneType, UnionType


def _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 UnionType

T = TypeVar("T")


class MessageSerializer(Protocol[T]):
    @property
    def data_content_type(self) -> str: ...

    @property
    def 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 dataclasses
    return 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 annotations
        origin = get_origin(field_type)
        args = get_args(field_type)

        # If the field type is directly a subclass of BaseModel
        if 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 types
                if isinstance(arg, type) and issubclass(arg, BaseModel):
                    return True
                elif 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 types
        elif args:
            for arg in args:
                if isinstance(arg, type) and issubclass(arg, BaseModel):
                    return True
                elif get_origin(arg) is not None:
                    if has_nested_base_model_in_type(arg):
                        return True
    return False


def 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 True
    if origin is not None and args:
        for arg in args:
            if has_nested_base_model_in_type(arg):
                return True
    return False


DataclassT = 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

    @property
    def data_content_type(self) -> str:
        return JSON_DATA_CONTENT_TYPE

    @property
    def 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 BaseModel

from 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]):
    @property
    def data_content_type(self) -> str: ...

    @property
    def 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: int
    y: int

serializer = 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: Point
    bottom_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