Skip to content

Commit

Permalink
refactor: add a mixin for int enum field custom
Browse files Browse the repository at this point in the history
  • Loading branch information
waketzheng committed Nov 22, 2024
1 parent 7f077c1 commit f0884f4
Show file tree
Hide file tree
Showing 3 changed files with 86 additions and 73 deletions.
2 changes: 1 addition & 1 deletion Makefile
Original file line number Diff line number Diff line change
Expand Up @@ -79,7 +79,7 @@ docs: deps
_style:
isort -src $(checkfiles)
black $(checkfiles)
style: _style deps
style: deps _style

build: deps
rm -fR dist/
Expand Down
21 changes: 19 additions & 2 deletions examples/enum_fields.py
Original file line number Diff line number Diff line change
@@ -1,6 +1,7 @@
from enum import Enum, IntEnum

from tortoise import Tortoise, fields, run_async
from tortoise.fields.data import IntEnumFieldMixin, IntField
from tortoise.models import Model


Expand All @@ -16,18 +17,34 @@ class Currency(str, Enum):
USD = "USD"


class Protocol(IntEnum):
A = 10000
B = 80000


class Int32EnumInstance(IntEnumFieldMixin, IntField):
pass


def Int32EnumField(enum_type, **kwargs):
return Int32EnumInstance(enum_type, **kwargs)


class EnumFields(Model):
service: Service = fields.IntEnumField(Service)
currency: Currency = fields.CharEnumField(Currency, default=Currency.HUF)
protocol: Protocol = Int32EnumField(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
Expand Down
Loading

0 comments on commit f0884f4

Please sign in to comment.