如何定义同时兼具Protocol与Pydantic BaseModel特性的类型？

嘿，这个问题我熟！你想让一个类型既能像Pydantic BaseModel那样做数据校验和序列化，又能通过Protocol来做结构类型检查对吧？毕竟直接让Protocol继承BaseModel是行不通的，Python不允许这么干，不过咱们有几种靠谱的办法：

方法一：用Protocol定义结构，自动匹配符合条件的BaseModel子类 #

这种方法最直接，利用Python的结构子类型特性——只要一个BaseModel子类有Protocol要求的属性，就会自动被视为符合该Protocol的类型，Pydantic v2也原生支持这种类型注解的校验。

from typing import Protocol
from pydantic import BaseModel

# 先定义Protocol，明确要求必须有id属性
class HasId(Protocol):
    id: int | str  # 你可以根据实际需求调整id的类型，比如只允许int

# 定义几个符合要求的BaseModel子类
class User(BaseModel):
    id: int
    name: str

class Product(BaseModel):
    id: str
    price: float

# 现在定义你的Foo模型，直接用HasId作为data字段的类型
class Foo(BaseModel):
    data: HasId

# 测试正常场景：符合要求的模型都能传入
user = User(id=123, name="Alice")
foo_with_user = Foo(data=user)
print(foo_with_user)  # 输出: data=User(id=123, name='Alice')

product = Product(id="prod-001", price=29.9)
foo_with_product = Foo(data=product)
print(foo_with_product)  # 输出: data=Product(id='prod-001', price=29.9)

# 测试异常场景：没有id字段的模型会触发校验错误
class InvalidModel(BaseModel):
    name: str

try:
    Foo(data=InvalidModel(name="Bob"))
except Exception as e:
    print(f"错误：{e}")  # 会提示缺少id字段

方法二：结合BaseModel基类与Protocol，强制子类继承 #

如果你想让所有符合HasId的模型都基于同一个BaseModel基类（比如统一id的校验规则），可以同时定义一个BaseModel基类和对应的Protocol，让子类继承BaseModel基类，这样既能享受BaseModel的特性，又自动符合Protocol的要求。

from typing import Protocol
from pydantic import BaseModel, field_validator

# 定义BaseModel基类，包含id字段和统一的校验规则
class HasIdBase(BaseModel):
    id: str

    @field_validator('id')
    def id_must_not_be_empty(cls, v: str) -> str:
        if not v.strip():
            raise ValueError("id不能为空")
        return v

# 定义对应的Protocol，和基类结构保持一致
class HasId(Protocol):
    id: str

# 子类继承HasIdBase，自动符合HasId Protocol
class Order(HasIdBase):
    order_number: str

# Foo模型依然用HasId作为类型注解
class Foo(BaseModel):
    data: HasId

# 测试：继承HasIdBase的子类可以正常传入
order = Order(id="order-1001", order_number="20240501001")
foo = Foo(data=order)
print(foo)  # 输出: data=Order(id='order-1001', order_number='20240501001')

方法三：添加校验器，确保传入的是BaseModel实例 #

如果你的需求不仅是要有id属性，还必须是BaseModel的实例（排除那些非BaseModel但有id的普通类），可以给Foo模型的data字段加一个自定义校验器：

from typing import Protocol, Any
from pydantic import BaseModel, field_validator, ValidationError

class HasId(Protocol):
    id: int | str

class Foo(BaseModel):
    data: HasId

    @field_validator('data')
    def ensure_is_basemodel(cls, v: Any) -> Any:
        if not isinstance(v, BaseModel):
            raise ValueError("data必须是Pydantic BaseModel的实例")
        return v

# 测试：非BaseModel类即使有id也会报错
class NonModelWithId:
    def __init__(self, id: int):
        self.id = id

try:
    Foo(data=NonModelWithId(id=789))
except ValidationError as e:
    print(f"错误：{e}")  # 会提示data必须是BaseModel实例

内容来源于stack exchange