serialization/validation support django-enumfields(2) (#5)

Author: provinzkraut

README.md

@@ -154,6 +154,14 @@ This can be mitigated by:
    to ensure relationships are fully loaded
 
 
+## Serialization / validation of 3rd party field types
+
+Additionally, the following 3rd party fields / types are supported if the
+`DjangoModelPlugin` is installed:
+
+- `django-enumfields`
+- `django-enumfields2`
+
 
 ## Contributing
 

litestar_django/plugin.py

@@ -1,13 +1,15 @@
+import contextlib
 from typing import Any
 
 from django.db import models  # type: ignore[import-untyped]
-from litestar.plugins.base import SerializationPlugin
+from litestar.config.app import AppConfig
+from litestar.plugins.base import SerializationPlugin, InitPlugin
 from litestar.typing import FieldDefinition
 
 from litestar_django.dto import DjangoModelDTO
 
 
-class DjangoModelPlugin(SerializationPlugin):
+class DjangoModelPlugin(InitPlugin, SerializationPlugin):
     def __init__(self) -> None:
         self._type_dto_map: dict[type[models.Model], type[DjangoModelDTO[Any]]] = {}
 
@@ -35,3 +37,29 @@ def create_dto_for_type(
         self._type_dto_map[annotation] = dto_type = DjangoModelDTO[annotation]  # type:ignore[valid-type]
 
         return dto_type
+
+    def on_app_init(self, app_config: AppConfig) -> AppConfig:
+        type_encoders = (
+            dict(app_config.type_encoders) if app_config.type_encoders else {}
+        )
+        type_decoders = (
+            list(app_config.type_decoders) if app_config.type_decoders else []
+        )
+        with contextlib.suppress(ImportError):
+            import enumfields  # type: ignore[import-untyped]
+
+            def _is_enumfields_enum(v: Any) -> bool:
+                return issubclass(v, (enumfields.Enum, enumfields.IntEnum))
+
+            def _decode_enumfields_enum(type_: Any, value: Any) -> Any:
+                return type_(value)
+
+            type_encoders[enumfields.Enum] = lambda v: v.value
+            type_encoders[enumfields.IntEnum] = lambda v: v.value
+
+            type_decoders.append((_is_enumfields_enum, _decode_enumfields_enum))
+
+        app_config.type_encoders = type_encoders
+        app_config.type_decoders = type_decoders
+
+        return app_config

pyproject.toml

@@ -1,6 +1,6 @@
 [project]
 name = "litestar-django"
-version = "0.1.1"
+version = "0.2.0"
 description = "Django model support for Litestar"
 readme = "README.md"
 license = { text = "MIT" }

tests/some_app/app/models.py

@@ -131,6 +131,11 @@ class ModelInvalidRegexValidator(models.Model):
     )
 
 
+class ModelWithCustomFields(models.Model):
+    enum_field = enumfields.EnumField(StdEnum)
+    enumfields_enum = enumfields.EnumField(LabelledEnum)
+
+
 class Author(models.Model):
     name = models.CharField(max_length=100)
 

tests/test_dto_integration.py

@@ -1,5 +1,5 @@
 import secrets
-from typing import Annotated
+from typing import Annotated, Any
 
 import pytest
 from litestar import get, Response, Litestar, post
@@ -8,7 +8,13 @@
 
 from litestar_django.dto import DjangoModelDTO
 from litestar_django.plugin import DjangoModelPlugin
-from tests.some_app.app.models import Author, Book, Genre, ModelWithFields
+from tests.some_app.app.models import (
+    Author,
+    Book,
+    Genre,
+    ModelWithFields,
+    ModelWithCustomFields,
+)
 
 
 @pytest.mark.django_db(transaction=True)
@@ -156,6 +162,38 @@ def handler(data: Author) -> Author:
         assert res.json() == {"id": author.id, "name": author_name, "books": []}
 
 
+@pytest.mark.django_db(transaction=True)
+def test_enumfields() -> None:
+    @post(
+        "/",
+        dto=DjangoModelDTO[Annotated[ModelWithCustomFields, DTOConfig(exclude={"id"})]],
+    )
+    async def post_handler(data: ModelWithCustomFields) -> dict[str, Any]:
+        await data.asave()
+        return {
+            "id": data.id,
+            "enum_field": data.enum_field.value,
+            "enumfields_enum": data.enumfields_enum.value,
+        }
+
+    @get("/{obj_id:int}", dto=DjangoModelDTO[ModelWithCustomFields])
+    async def get_handler(obj_id: int) -> ModelWithCustomFields:
+        return await ModelWithCustomFields.objects.aget(id=obj_id)
+
+    with create_test_client(
+        [get_handler, post_handler], plugins=[DjangoModelPlugin()]
+    ) as client:
+        res = client.post("/", json={"enum_field": "ONE", "enumfields_enum": "TWO"})
+        assert res.status_code == 201
+        data = res.json()
+        model_id = data.pop("id")
+        assert data == {"enum_field": "ONE", "enumfields_enum": "TWO"}
+
+        res = client.get(f"/{model_id}")
+        assert res.status_code == 200
+        assert res.json() == {"id": model_id, **data}
+
+
 def test_schema() -> None:
     @get("/", dto=DjangoModelDTO[ModelWithFields], sync_to_thread=True)
     def handler() -> ModelWithFields: