add enum support to pydantic_model_creator (#1798)

* add pydantic enum support

Author: markus-96

CHANGELOG.rst

@@ -21,6 +21,10 @@ Changed
 - Parametrizes UPDATE, DELETE, bulk update and create operations (#1785)
 - Parametrizes related field queries (#1797)
 
+Added
+^^^^^
+- CharEnumField and IntEnumField is supported by pydantic_model_creator (#1798)
+
 0.22.1
 ------
 Fixed
@@ -52,7 +56,7 @@ Added
 Changed
 ^^^^^^^
 - Change old pydantic docs link to new one (#1775).
-- Refactored pydantic_model_creator, interface not changed  (#1763)
+- Refactored pydantic_model_creator, interface not changed  (#1745)
 - Values are no longer validated to be right type upon loading from database (#1750)
 - Refactored private field names in queryset classes (#1751)
 

CONTRIBUTORS.rst

@@ -62,6 +62,7 @@ Contributors
 * Daniel Szucs ``@Quasar6X``
 * Rui Catarino ``@ruitcatarino``
 * Lance Moe ``@lancemoe``
+* Markus Beckschulte ``@markus-96``
 
 Special Thanks
 ==============

tests/contrib/test_pydantic.py

@@ -15,6 +15,7 @@
     Tournament,
     User,
     json_pydantic_default,
+    EnumFields,
 )
 from tortoise.contrib import test
 from tortoise.contrib.pydantic import (
@@ -1960,3 +1961,122 @@ def test_named_model_with_relations(self):
         Pydantic2 = pydantic_model_creator(self.ModelWithRelations, name="Foo")
 
         self.assertIs(Pydantic1, Pydantic2)
+
+
+class TestPydanticEnum(test.TestCase):
+    def setUp(self) -> None:
+        self.EnumFields_Pydantic = pydantic_model_creator(EnumFields)
+
+    def test_int_enum(self):
+        with self.assertRaises(ValidationError) as cm:
+            self.EnumFields_Pydantic.model_validate({"id": 1, "service": 4, "currency": "HUF"})
+        self.assertEqual(
+            [
+                {
+                    'type': 'enum',
+                    'loc': ('service',),
+                    'msg': 'Input should be 1, 2 or 3',
+                    'input': 4,
+                    'ctx': {'expected': '1, 2 or 3'}
+                }
+            ],
+            cm.exception.errors(include_url=False)
+        )
+        with self.assertRaises(ValidationError) as cm:
+            self.EnumFields_Pydantic.model_validate({"id": 1, "service": "a string, not int", "currency": "HUF"})
+        self.assertEqual(
+            [
+                {
+                    'type': 'enum',
+                    'loc': ('service',),
+                    'msg': 'Input should be 1, 2 or 3',
+                    'input': "a string, not int",
+                    'ctx': {'expected': '1, 2 or 3'}
+                }
+            ],
+            cm.exception.errors(include_url=False)
+        )
+
+    def test_str_enum(self):
+        with self.assertRaises(ValidationError) as cm:
+            self.EnumFields_Pydantic.model_validate({"id": 1, "service": 3, "currency": "GoofyGooberDollar"})
+        self.assertEqual(
+            [
+                {
+                    'type': 'enum',
+                    'loc': ('currency',),
+                    'msg': "Input should be 'HUF', 'EUR' or 'USD'",
+                    'input': 'GoofyGooberDollar',
+                    'ctx': {'expected': "'HUF', 'EUR' or 'USD'"}
+                }
+            ],
+            cm.exception.errors(include_url=False)
+        )
+        with self.assertRaises(ValidationError) as cm:
+            self.EnumFields_Pydantic.model_validate({"id": 1, "service": 3, "currency": 1})
+        self.assertEqual(
+            [
+                {
+                    'type': 'enum',
+                    'loc': ('currency',),
+                    'msg': "Input should be 'HUF', 'EUR' or 'USD'",
+                    'input': 1,
+                    'ctx': {'expected': "'HUF', 'EUR' or 'USD'"}
+                }
+            ],
+            cm.exception.errors(include_url=False)
+        )
+
+    def test_enum(self):
+        with self.assertRaises(ValidationError) as cm:
+            self.EnumFields_Pydantic.model_validate({"id": 1, "service": 4, "currency": 1})
+        self.assertEqual(
+            [
+                {
+                    'type': 'enum',
+                    'loc': ('service',),
+                    'msg': 'Input should be 1, 2 or 3',
+                    'input': 4, 'ctx': {'expected': '1, 2 or 3'}
+                },
+                {
+                    'type': 'enum',
+                    'loc': ('currency',),
+                    'msg': "Input should be 'HUF', 'EUR' or 'USD'",
+                    'input': 1,
+                    'ctx': {'expected': "'HUF', 'EUR' or 'USD'"}
+                }
+            ],
+            cm.exception.errors(include_url=False)
+        )
+
+        # should simply not raise any error:
+        self.EnumFields_Pydantic.model_validate({"id": 1, "service": 3, "currency": "HUF"})
+
+        self.assertEqual(
+            {
+                'additionalProperties': False,
+                'properties': {
+                    'id': {'maximum': 2147483647, 'minimum': -2147483648, 'title': 'Id', 'type': 'integer'},
+                    'service': {
+                        'description': 'python_programming: 1<br/>database_design: 2<br/>system_administration: 3',
+                        'enum': [1, 2, 3],
+                        'ge': -32768,
+                        'le': 32767,
+                        'title': 'Service',
+                        'type': 'integer'
+                    },
+                    'currency': {
+                        'default': 'HUF',
+                        'description': 'HUF: HUF<br/>EUR: EUR<br/>USD: USD',
+                        'enum': ['HUF', 'EUR', 'USD'],
+                        'maxLength': 3,
+                        'title': 'Currency',
+                        'type': 'string'
+                    }
+                },
+                'required': ['id', 'service'],
+                'title': 'EnumFields',
+                'type': 'object'
+            },
+            self.EnumFields_Pydantic.model_json_schema()
+        )

tortoise/contrib/pydantic/creator.py

@@ -1,6 +1,7 @@
 import inspect
 from base64 import b32encode
 from copy import copy
+from enum import IntEnum, Enum
 from hashlib import sha3_224
 from typing import (
     TYPE_CHECKING,
@@ -33,6 +34,7 @@
 )
 from tortoise.contrib.pydantic.utils import get_annotations
 from tortoise.fields import Field, JSONField
+from tortoise.fields.data import IntEnumFieldInstance, CharEnumFieldInstance
 
 if TYPE_CHECKING:  # pragma: nocoverage
     from tortoise.models import Model
@@ -509,7 +511,11 @@ def _process_data_field(
             json_schema_extra["readOnly"] = constraints["readOnly"]
             del constraints["readOnly"]
         fconfig.update(constraints)
-        python_type = getattr(field, "related_model", field.field_type)
+        python_type: Union[Type[Enum], Type[IntEnum], Type]
+        if isinstance(field, (IntEnumFieldInstance, CharEnumFieldInstance)):
+            python_type = field.enum_type
+        else:
+            python_type = getattr(field, "related_model", field.field_type)
         ptype = python_type
         if field.null:
             json_schema_extra["nullable"] = True