如何为Peewee模型的create方法添加类型提示？

我在使用Peewee ORM时遇到一个问题：模型继承的Base类中create方法的签名是def create(**query: Any) -> Any，导致调用Group.create()时没有必填字段的类型提示。手动在每个模型中重写create方法虽然能解决问题，但重复代码太多，希望找到更简洁的通用解决方案。

我的模型示例：

class GroupEntity(Base):
    id: UUID = UUIDField(unique=True, primary_key=True, default=uuid4())
    user_id: UUID = ForeignKeyField(User, backref='groups')
    title: str = CharField(max_length=256)
    type: GroupTypeEnum = CharField(choices=GROUP_TYPES)
    parent_id: UUID = UUIDField(null=True)
    price: int = SmallIntegerField(null=True)
    created_at: float = BigIntegerField(default=time.time())

    class Meta:
        table_name = TableNamesEnum.GROUP_ENTITY.value

调用create时无类型提示：

group_entity: Group = Group.create(
    title='best group',
    type=GroupTypeEnum.Channel,
    user_id=user_entity.id
)

方案一：通过元类自动生成带类型提示的create方法 #

利用元类可以在模型类创建时，自动根据字段的类型提示和属性（是否必填、是否有默认值）生成对应的create方法签名，所有继承该元类的模型都会自动获得类型提示。

实现步骤： #

1. 定义自定义元类：

from typing import TypeVar, get_type_hints, Optional
import peewee
from peewee import Model

# 定义模型类型变量
T = TypeVar('T', bound=Model)

class TypedModelMeta(peewee.ModelBase):
    def __new__(cls, name, bases, attrs):
        # 创建模型类
        model_class = super().__new__(cls, name, bases, attrs)
        # 获取模型的字段类型提示
        field_types = get_type_hints(model_class)
        
        required_params = []
        optional_params = []
        
        # 遍历模型字段，区分必填和可选参数
        for field_name, field in model_class._meta.fields.items():
            # 跳过自动生成的主键（有default的主键）
            if field_name == 'id' and field.primary_key and field.default is not None:
                continue
            
            # 获取字段对应的类型
            field_type = field_types.get(field_name, object)
            # 处理Optional类型的字符串表示
            type_str = getattr(field_type, '__name__', repr(field_type))
            
            if field.null or field.default is not None:
                # 可选参数，添加默认值None
                optional_params.append(f"{field_name}: Optional[{type_str}] = None")
            else:
                # 必填参数
                required_params.append(f"{field_name}: {type_str}")
        
        # 构建create方法的代码
        create_code = f"""
@classmethod
def create(cls, {', '.join(required_params + optional_params)}) -> '{name}':
    # 收集参数并调用父类的create方法
    kwargs = {{k: v for k, v in locals().items() if k != 'cls'}}
    return super({name}, cls).create(**kwargs)
"""
        # 执行代码并绑定到模型类
        local_vars = {}
        exec(create_code, globals(), local_vars)
        setattr(model_class, 'create', local_vars['create'])
        
        return model_class

2. 修改Base模型使用该元类：

class Base(Model, metaclass=TypedModelMeta):
    class Meta:
        database = your_database_connection  # 替换为你的数据库连接

3. 模型继承Base即可自动获得类型提示：
   现在GroupEntity.create()会自动提示必填的user_id、title、type参数，以及可选的parent_id、price参数，返回值类型也会正确提示为GroupEntity。

方案二：使用TypedDict结合泛型基类（适用于需要更精细控制的场景） #

如果不想用元类，可以为每个模型定义对应的TypedDict来描述创建参数，再通过泛型基类给create方法添加类型注解。

实现示例： #

from typing import TypedDict, TypeVar, Generic, Optional
from peewee import Model

# 定义类型变量
M = TypeVar('M', bound=Model)
T = TypeVar('T', bound=TypedDict)

class TypedModel(Model, Generic[T]):
    @classmethod
    def create(cls, **kwargs: T) -> M:
        return super().create(**kwargs)

# 定义GroupEntity的创建参数TypedDict
class GroupCreateParams(TypedDict):
    user_id: UUID
    title: str
    type: GroupTypeEnum
    parent_id: Optional[UUID]
    price: Optional[int]

# 模型继承TypedModel并指定参数类型
class GroupEntity(TypedModel[GroupCreateParams], Base):
    # 模型字段定义...
    id: UUID = UUIDField(unique=True, primary_key=True, default=uuid4())
    user_id: UUID = ForeignKeyField(User, backref='groups')
    title: str = CharField(max_length=256)
    type: GroupTypeEnum = CharField(choices=GROUP_TYPES)
    parent_id: UUID = UUIDField(null=True)
    price: int = SmallIntegerField(null=True)
    created_at: float = BigIntegerField(default=time.time())

    class Meta:
        table_name = TableNamesEnum.GROUP_ENTITY.value

这种方式需要为每个模型定义对应的TypedDict，但比手动重写create方法更简洁，且类型提示清晰。

总结 #

• 方案一（元类）是最通用的解决方案，一次定义后所有模型自动受益，无需重复代码。
• 方案二（TypedDict+泛型）适合需要对创建参数有更精细控制的场景，灵活性更高。

内容的提问来源于stack exchange，提问作者Ryzhik Daniel