EnumField and CharEnumField (tortoise#244)

Author: wallneradam

docs/examples/basic.rst

@@ -93,3 +93,12 @@ Recursive Relations
 ===================
 .. literalinclude::  ../../examples/relations_recursive.py
 
+
+.. rst-class:: html-toggle
+
+.. _example_enum_fields:
+
+===============
+Enumeration Fields
+===============
+.. literalinclude::  ../../examples/enum_fields.py

docs/fields.rst

@@ -70,6 +70,12 @@ Data Fields
 .. autoclass:: tortoise.fields.UUIDField
     :exclude-members: to_db_value, to_python_value
 
+.. autoclass:: tortoise.fields.IntEnumField
+    :exclude-members: to_db_value, to_python_value
+
+.. autoclass:: tortoise.fields.CharEnumField
+    :exclude-members: to_db_value, to_python_value
+
 
 Relational Fields
 -----------------

examples/enum_fields.py

@@ -0,0 +1,22 @@
+import typing
+from enum import Enum, IntEnum
+
+from tortoise import fields
+from tortoise.models import Model
+
+
+class Service(IntEnum):
+    python_programming = 1
+    database_design = 2
+    system_administration = 3
+
+
+class Currency(str, Enum):
+    HUF = "HUF"
+    EUR = "EUR"
+    USD = "USD"
+
+
+class EnumFields(Model):
+    service: Service = typing.cast(Service, fields.IntEnumField(Service))
+    currency: Currency = typing.cast(Currency, fields.CharEnumField(Currency, default=Currency.HUF))

tests/fields/test_enum.py

@@ -0,0 +1,35 @@
+from tests import testmodels
+from tortoise.contrib import test
+from tortoise.exceptions import IntegrityError
+
+
+class TestEnumFields(test.TestCase):
+    async def test_empty(self):
+        with self.assertRaises(IntegrityError):
+            await testmodels.EnumFields.create()
+
+    async def test_create(self):
+        obj0 = await testmodels.EnumFields.create(service=testmodels.Service.system_administration)
+        self.assertIsInstance(obj0.service, testmodels.Service)
+        self.assertIsInstance(obj0.currency, testmodels.Currency)
+        obj = await testmodels.EnumFields.get(id=obj0.id)
+        self.assertIsInstance(obj.service, testmodels.Service)
+        self.assertIsInstance(obj.currency, testmodels.Currency)
+        self.assertEqual(obj.service, testmodels.Service.system_administration)
+        self.assertEqual(obj.currency, testmodels.Currency.HUF)
+        await obj.save()
+        obj2 = await testmodels.EnumFields.get(id=obj.id)
+        self.assertEqual(obj, obj2)
+
+        await obj.delete()
+        obj = await testmodels.EnumFields.filter(id=obj0.id).first()
+        self.assertEqual(obj, None)
+
+    async def test_update(self):
+        obj0 = await testmodels.EnumFields.create(service=testmodels.Service.system_administration)
+        await testmodels.EnumFields.filter(id=obj0.id).update(
+            service=testmodels.Service.database_design, currency=testmodels.Currency.EUR
+        )
+        obj = await testmodels.EnumFields.get(id=obj0.id)
+        self.assertEqual(obj.service, testmodels.Service.database_design)
+        self.assertEqual(obj.currency, testmodels.Currency.EUR)

tests/testmodels.py

@@ -3,7 +3,9 @@
 """
 import binascii
 import os
+import typing
 import uuid
+from enum import Enum, IntEnum
 
 from tortoise import fields
 from tortoise.models import Model
@@ -481,3 +483,20 @@ class Meta:
         table = "sometable"
         unique_together = [["chars", "blip"]]
         table_description = "Source mapped fields"
+
+
+class Service(IntEnum):
+    python_programming = 1
+    database_design = 2
+    system_administration = 3
+
+
+class Currency(str, Enum):
+    HUF = "HUF"
+    EUR = "EUR"
+    USD = "USD"
+
+
+class EnumFields(Model):
+    service: Service = typing.cast(Service, fields.IntEnumField(Service))
+    currency: Currency = typing.cast(Currency, fields.CharEnumField(Currency, default=Currency.HUF))

tortoise/fields.py

@@ -2,6 +2,7 @@
 import functools
 import json
 from decimal import Decimal
+from enum import Enum, IntEnum
 from typing import TYPE_CHECKING, Any, Awaitable, Generic, Optional, Type, TypeVar, Union
 from uuid import UUID, uuid4
 
@@ -407,6 +408,100 @@ def to_python_value(self, value: Any) -> Optional[UUID]:
         return UUID(value)
 
 
+class IntEnumField(SmallIntField):
+    """
+    Enum Field
+
+    A field representing an integer enumeration.
+
+    The description of the field is set automatically if not specified to a multiline list of
+    "name: value" pairs.
+
+    ``enum_type``:
+        The enum class
+    ``description``:
+        The description of the field. It is set automatically if not specified to a multiline list
+        of "name: value" pairs.
+    """
+
+    __slots__ = ("enum_type",)
+
+    def __init__(
+        self, enum_type: Type[IntEnum], description: Optional[str] = None, **kwargs
+    ) -> None:
+        # Validate values
+        for item in enum_type:
+            try:
+                value = int(item.value)
+            except ValueError:
+                raise ConfigurationError("IntEnumField only supports integer enums!")
+            if not 0 <= value < 32768:
+                raise ConfigurationError("The valid range of IntEnumField's values is 0..32767!")
+
+        # 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]:
+        return self.enum_type(value) if value is not None else None
+
+    def to_db_value(self, value: Union[IntEnum, None], instance) -> Union[int, None]:
+        return int(value.value) if value is not None else None
+
+
+class CharEnumField(CharField):
+    """
+    Char Enum Field
+
+    A field representing a character enumeration.
+
+    **Warning**: If ``max_length`` is not specified or equals to zero, the size of represented
+    char fields is automatically detected. So if later you update the enum, you need to update your
+    table schema as well.
+
+    ``enum_type``:
+        The enum class
+    ``description``:
+        The description of the field. It is set automatically if not specified to a multiline list
+        of "name: value" pairs.
+    ``max_length``:
+        The length of the created CharField. If it is zero it is automatically detected from
+        enum_type.
+    """
+
+    __slots__ = ("enum_type",)
+
+    def __init__(
+        self,
+        enum_type: Type[Enum],
+        description: Optional[str] = None,
+        max_length: int = 0,
+        **kwargs,
+    ) -> None:
+        # Automatic description for the field if not specified by the user
+        if description is None:
+            description = "\n".join([f"{e.name}: {str(e.value)}" for e in enum_type])[:2048]
+
+        # Automatic CharField max_length
+        if max_length == 0:
+            for item in enum_type:
+                item_len = len(str(item.value))
+                if item_len > max_length:
+                    max_length = item_len
+
+        super().__init__(description=description, max_length=max_length, **kwargs)
+        self.enum_type = enum_type
+
+    def to_python_value(self, value: Union[str, None]) -> Union[Enum, None]:
+        return self.enum_type(value) if value is not None else None
+
+    def to_db_value(self, value: Union[Enum, None], instance) -> Union[str, None]:
+        return str(value.value) if value is not None else None
+
+
 ForeignKeyNullableRelation = Union[Awaitable[Optional[MODEL]], Optional[MODEL]]
 """
 Type hint for the result of accessing the :class:`.ForeignKeyField` field in the model