Merge pull request #29 from jordaneremieff/enum-choices

Add tests for TextChoices and IntegerChoices and change default schema representation

djantic/fields.py

@@ -91,6 +91,8 @@ def ModelSchemaField(field: Any) -> tuple:
                 enum_choices,
                 module=__name__,
             )
+            if field.has_default() and isinstance(field.default, Enum):
+                default = field.default.value
         else:
             internal_type = field.get_internal_type()
             if internal_type in STR_TYPES:
@@ -118,7 +120,7 @@ def ModelSchemaField(field: Any) -> tuple:
         blank = field_options.pop("blank", False)
         null = field_options.pop("null", False)
 
-        if field.has_default():
+        if default is Required and field.has_default():
             if callable(field.default):
                 default_factory = field.default
                 default = Undefined

djantic/main.py

@@ -49,6 +49,7 @@ def __new__(
                     )
 
                 annotations = namespace.get("__annotations__", {})
+
                 try:
                     fields = config.model._meta.get_fields()
                 except AttributeError as exc:

tests/test_fields.py

@@ -1,6 +1,6 @@
 import pytest
 
-from testapp.models import Record, Configuration
+from testapp.models import Record, Configuration, Preference, FoodChoices, GroupChoices
 from djantic import ModelSchema
 
 
@@ -130,3 +130,59 @@ class Config:
             },
         },
     }
+
+
+@pytest.mark.django_db
+def test_enum_choices():
+    class PreferenceSchema(ModelSchema):
+        class Config:
+            model = Preference
+            use_enum_values = True
+
+    assert PreferenceSchema.schema() == {
+        "title": "PreferenceSchema",
+        "description": "Preference(id, name, preferred_food, preferred_group)",
+        "type": "object",
+        "properties": {
+            "id": {"title": "Id", "description": "id", "type": "integer"},
+            "name": {
+                "title": "Name",
+                "description": "name",
+                "maxLength": 128,
+                "type": "string",
+            },
+            "preferred_food": {
+                "title": "Preferred Food",
+                "description": "preferred_food",
+                "default": "ba",
+                "allOf": [{"$ref": "#/definitions/PreferredFoodEnum"}],
+            },
+            "preferred_group": {
+                "title": "Preferred Group",
+                "description": "preferred_group",
+                "default": 1,
+                "allOf": [{"$ref": "#/definitions/PreferredGroupEnum"}],
+            },
+        },
+        "required": ["name"],
+        "definitions": {
+            "PreferredFoodEnum": {
+                "title": "PreferredFoodEnum",
+                "description": "An enumeration.",
+                "enum": ["ba", "ap"],
+            },
+            "PreferredGroupEnum": {
+                "title": "PreferredGroupEnum",
+                "description": "An enumeration.",
+                "enum": [1, 2],
+            },
+        },
+    }
+
+    preference = Preference.objects.create(name="Jordan")
+    assert PreferenceSchema.from_django(preference).dict() == {
+        "id": 1,
+        "name": "Jordan",
+        "preferred_food": "ba",
+        "preferred_group": 1,
+    }

tests/testapp/models.py

@@ -201,3 +201,23 @@ def upload_image_handler(instance, filename):
 class Attachment(models.Model):
     description = models.CharField(max_length=255)
     image = models.ImageField(blank=True, null=True, upload_to=upload_image_handler)
+
+
+class FoodChoices(models.TextChoices):
+    BANANA = "ba", "A delicious yellow Banana"
+    APPLE = "ap", "A delicious red Apple"
+
+
+class GroupChoices(models.IntegerChoices):
+    GROUP_1 = 1, "First group"
+    GROUP_2 = 2, "Second group"
+
+
+class Preference(models.Model):
+    name = models.CharField(max_length=128)
+    preferred_food = models.CharField(
+        max_length=2, choices=FoodChoices.choices, default=FoodChoices.BANANA
+    )
+    preferred_group = models.IntegerField(
+        choices=GroupChoices.choices, default=GroupChoices.GROUP_1
+    )