tortoise / tortoise-orm

Familiar asyncio ORM for python, built with relations in mind
https://tortoise.github.io
Apache License 2.0
4.62k stars 383 forks source link

Python Enum.IntEnum datatype do not fit on Tortoise.fields.IntEnumField when not int16 #966

Open igormorgado opened 2 years ago

igormorgado commented 2 years ago

Describe the bug Python Enum.IntEnum datatype do not fit on Tortoise.fields.IntEnumField when any element of enumeration is outside of int16 range.

IMHO, this should not behave like that (or at least configurable), since Python object can handle values larger than that.

To Reproduce

import enum
import tortoise.fields as fields

class PROTOCOL_ID(enum.IntEnum):
     A = 10000
     B = 80000

protocol_id = fields.IntEnumField(enum_type=PROTOCOL_ID)

Fails with

ConfigurationError: The valid range of IntEnumField's values is -32768..32767! 

Expected behavior I expect that I could store values larger or smaller than int16 on that. Or at least a BigintEnumField. and SmallIntEnumField.

IMHO the naming should reflect the regular Int type fields as;

Additional context If not possible, is there a workaround? I do not want store that protocol static data into a database to increase the query complexity.

Python 3.9.7 Tortoise 0.17.8

MovisLi commented 4 days ago

Me too. Or it should called SmallIntEnumField instead of IntEnumField

MovisLi commented 4 days ago

To replace tortoise.field.data, maybe this is better:

from enum import IntEnum
from typing import TYPE_CHECKING, Any, Optional, Type, TypeVar, Union

from tortoise.exceptions import ConfigurationError
from tortoise.fields.data import IntField, SmallIntField

if TYPE_CHECKING:  # pragma: nocoverage
    from tortoise.models import Model

IntEnumType = TypeVar("IntEnumType", bound=IntEnum)

class SmallIntEnumFieldInstance(SmallIntField):
    def __init__(
        self,
        enum_type: Type[IntEnum],
        description: Optional[str] = None,
        generated: bool = False,
        **kwargs: Any,
    ) -> None:
        # Validate values
        minimum = 1 if generated else -32768
        for item in enum_type:
            try:
                value = int(item.value)
            except ValueError:
                raise ConfigurationError("IntEnumField only supports integer enums!")
            if not minimum <= value < 32768:
                raise ConfigurationError("The valid range of SmallIntEnumField's values is {}..32767!".format(minimum))

        # Automatic description for the field if not specified by the user
        if description is None:
            description = "\n".join([f"{e.name}: {int(e.value)}" for e in enum_type])[:2048]

        super().__init__(description=description, **kwargs)
        self.enum_type = enum_type

    def to_python_value(self, value: Union[int, None]) -> Union[IntEnum, None]:
        value = self.enum_type(value) if value is not None else None
        self.validate(value)
        return value

    def to_db_value(self, value: Union[IntEnum, None, int], instance: "Union[Type[Model], Model]") -> Union[int, None]:
        if isinstance(value, IntEnum):
            value = int(value.value)
        if isinstance(value, int):
            value = int(self.enum_type(value))
        self.validate(value)
        return value

def SmallIntEnumField(
    enum_type: Type[IntEnumType],
    description: Optional[str] = None,
    **kwargs: Any,
) -> IntEnumType:
    return SmallIntEnumFieldInstance(enum_type, description, **kwargs)  # type: ignore

class IntEnumFieldInstance(IntField):
    def __init__(
        self,
        enum_type: Type[IntEnum],
        description: Optional[str] = None,
        generated: bool = False,
        **kwargs: Any,
    ) -> None:
        # Validate values
        minimum = 1 if generated else -2147483648
        for item in enum_type:
            try:
                value = int(item.value)
            except ValueError:
                raise ConfigurationError("IntEnumField only supports integer enums!")
            if not minimum <= value < 2147483648:
                raise ConfigurationError("The valid range of IntEnumField's values is {}..2147483647!".format(minimum))

        # Automatic description for the field if not specified by the user
        if description is None:
            description = "\n".join([f"{e.name}: {int(e.value)}" for e in enum_type])[:2048]

        super().__init__(description=description, **kwargs)
        self.enum_type = enum_type

    def to_python_value(self, value: Union[int, None]) -> Union[IntEnum, None]:
        value = self.enum_type(value) if value is not None else None
        self.validate(value)
        return value

    def to_db_value(self, value: Union[IntEnum, None, int], instance: "Union[Type[Model], Model]") -> Union[int, None]:
        if isinstance(value, IntEnum):
            value = int(value.value)
        if isinstance(value, int):
            value = int(self.enum_type(value))
        self.validate(value)
        return value

def IntEnumField(
    enum_type: Type[IntEnumType],
    description: Optional[str] = None,
    **kwargs: Any,
) -> IntEnumType:
    return IntEnumFieldInstance(enum_type, description, **kwargs)  # type: ignore