feat: add mixin for int enum field custom

CHANGELOG.rst

@@ -16,6 +16,10 @@ Fixed
 - Fix unable to use ManyToManyField if OneToOneField passed as Primary Key (#1783)
 - Fix sorting by Term (e.g. RawSQL) (#1788)
 
+Added
+^^^^^
+- Add IntEnumFieldMixin to make it easy to Custom IntEnumField (#1781)
+
 0.22.0
 ------
 Fixed

Makefile

@@ -79,7 +79,7 @@ docs: deps
 _style:
 	isort -src $(checkfiles)
 	black $(checkfiles)
-style: _style deps
+style: deps _style
 
 build: deps
 	rm -fR dist/

examples/enum_fields.py

@@ -1,6 +1,8 @@
 from enum import Enum, IntEnum
+from typing import Type, cast
 
 from tortoise import Tortoise, fields, run_async
+from tortoise.fields.data import IntEnumFieldMixin, IntEnumType, IntField
 from tortoise.models import Model
 
 
@@ -16,18 +18,36 @@ class Currency(str, Enum):
     USD = "USD"
 
 
+class Protocol(IntEnum):
+    A = 10000
+    B = 80000  # >32767, beyond the 'le' value in fields.IntEnumFieldInstance.constraints
+
+
+class IntEnumFieldInstance(IntEnumFieldMixin, IntField):
+    pass
+
+
+def IntEnumField(enum_type: Type[IntEnumType], **kwargs) -> IntEnumType:
+    return cast(IntEnumType, IntEnumFieldInstance(enum_type, **kwargs))
+
+
 class EnumFields(Model):
-    service: Service = fields.IntEnumField(Service)
     currency: Currency = fields.CharEnumField(Currency, default=Currency.HUF)
+    # When each value of the enum_type is between [-32768, 32767], use the fields.IntEnumField
+    service: Service = fields.IntEnumField(Service)
+    # Else, you can use a custom Field
+    protocol: Protocol = IntEnumField(Protocol)
 
 
 async def run():
     await Tortoise.init(db_url="sqlite://:memory:", modules={"models": ["__main__"]})
     await Tortoise.generate_schemas()
 
-    obj0 = await EnumFields.create(service=Service.python_programming, currency=Currency.USD)
+    obj0 = await EnumFields.create(
+        service=Service.python_programming, currency=Currency.USD, protocol=Protocol.A
+    )
     # also you can use valid int and str value directly
-    await EnumFields.create(service=1, currency="USD")
+    await EnumFields.create(service=1, currency="USD", protocol=Protocol.B.value)
 
     try:
         # invalid enum value will raise ValueError

tests/fields/test_enum.py

@@ -1,5 +1,6 @@
 from enum import IntEnum
 
+from examples.enum_fields import IntEnumField as CustomIntEnumField
 from tests import testmodels
 from tortoise.contrib import test
 from tortoise.exceptions import ConfigurationError, IntegrityError
@@ -24,6 +25,12 @@ class BadIntEnumIfGenerated(IntEnum):
     system_administration = 3
 
 
+class BadCustomIntEnum(IntEnum):
+    python_programming = 2147483648
+    database_design = 2
+    system_administration = 3
+
+
 class TestIntEnumFields(test.TestCase):
     async def test_empty(self):
         with self.assertRaises(IntegrityError):
@@ -118,6 +125,11 @@ def test_manual_description(self):
         fld = IntEnumField(testmodels.Service, description="foo")
         self.assertEqual(fld.description, "foo")
 
+    def test_custom_int_enum_field(self):
+        assert CustomIntEnumField(BadIntEnum1)
+        with self.assertRaises(ConfigurationError):
+            CustomIntEnumField(BadCustomIntEnum)
+
 
 class TestCharEnumFields(test.TestCase):
     async def test_create(self):

tortoise/fields/data.py

@@ -1,10 +1,12 @@
+from __future__ import annotations
+
 import datetime
 import functools
 import json
 import warnings
 from decimal import Decimal
 from enum import Enum, IntEnum
-from typing import TYPE_CHECKING, Any, Callable, Optional, Type, TypeVar, Union
+from typing import TYPE_CHECKING, Any, Callable, Optional, Type, TypeVar, Union, cast
 from uuid import UUID, uuid4
 
 from pypika import functions
@@ -668,7 +670,10 @@ class _db_mssql:
         SQL_TYPE = "VARBINARY(MAX)"
 
 
-class IntEnumFieldInstance(SmallIntField):
+class IntEnumFieldMixin:
+    validate: Callable
+    constraints: dict[str, int]
+
     def __init__(
         self,
         enum_type: Type[IntEnum],
@@ -677,27 +682,30 @@ def __init__(
         **kwargs: Any,
     ) -> None:
         # Validate values
-        minimum = 1 if generated else -32768
+        minimum = 1 if generated else self.constraints["ge"]
+        maximum = self.constraints["le"]
         for item in enum_type:
             try:
                 value = int(item.value)
             except ValueError:
                 raise ConfigurationError("IntEnumField only supports integer enums!")
-            if not minimum <= value < 32768:
+            if not minimum <= value <= maximum:
+                # To extend value range, see: https://tortoise.github.io/examples/basic.html#enumeration-fields
                 raise ConfigurationError(
-                    "The valid range of IntEnumField's values is {}..32767!".format(minimum)
+                    f"The valid range of IntEnumField's values is {minimum}..{maximum}!"
                 )
 
         # 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)
+        super().__init__(
+            description=description, generated=generated, **kwargs
+        )  # type:ignore[call-arg]
         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
-        return value
+        return self.enum_type(value) if value is not None else None
 
     def to_db_value(
         self, value: Union[IntEnum, None, int], instance: "Union[Type[Model], Model]"
@@ -713,6 +721,10 @@ def to_db_value(
 IntEnumType = TypeVar("IntEnumType", bound=IntEnum)
 
 
+class IntEnumFieldInstance(IntEnumFieldMixin, SmallIntField):
+    pass
+
+
 def IntEnumField(
     enum_type: Type[IntEnumType],
     description: Optional[str] = None,
@@ -735,7 +747,7 @@ def IntEnumField(
         of "name: value" pairs.
 
     """
-    return IntEnumFieldInstance(enum_type, description, **kwargs)  # type: ignore
+    return cast("IntEnumType", IntEnumFieldInstance(enum_type, description, **kwargs))
 
 
 class CharEnumFieldInstance(CharField):
@@ -805,4 +817,4 @@ def CharEnumField(
 
     """
 
-    return CharEnumFieldInstance(enum_type, description, max_length, **kwargs)  # type: ignore
+    return cast("CharEnumType", CharEnumFieldInstance(enum_type, description, max_length, **kwargs))