Tortoise ORM 让IntEnumField字段支持负数值转换 Python

784 阅读1分钟

介绍

Tortoise ORM 自带的IntEnumField字段,不支持负数,那实际开发中会用-1表示某种不存在的状态,就会引出超出0~32768范围报错;同时在接口返回中,并不是返回状态值,而是枚举的属性--!导致无法正常序列化。

解决思路

1、重载数字枚举字段 2、覆盖字段的to_python_value方法

代码

from enum import IntEnum
from typing import Type

from tortoise import fields, ConfigurationError
from tortoise.fields.data import IntEnumFieldInstance

class IntEnumField(IntEnumFieldInstance):
    """重载数字枚举模型字段"""

    def __init__(self, enum_type: Type[IntEnum], description=None, **kwargs):
        try:
            super(IntEnumField, self).__init__(enum_type, description, **kwargs)
        except ConfigurationError as cex:
            # 忽略超出0 <= value < 32768范围报错
            for item in enum_type:
                try:
                    int(item.value)
                except ValueError:
                    raise ConfigurationError("IntEnumField only supports integer enums!")
            if description is None:
                description = "\n".join(
                    [f"{e.name}: {int(e.value)}" for e in enum_type])[:2048]
            super(IntEnumFieldInstance, self).__init__(description=description, **kwargs)
            self.enum_type = enum_type

    def to_python_value(self, value: int) -> [int, None]:
        # value = self.enum_type(value) if value is not None else None
        # 直接返回枚举数值,覆盖原来的枚举属性,否则无法json序列化处理。
        value = value if value is not None else None
        self.validate(value)
        return value

fields.IntEnumField = IntEnumField

操作

从自定义的这个代码中引入fields模块,即可在定义模型时使用重载后的IntEnumField字段。 from 你的路径 import fields