Enum generated is not JSON serializable

[deanq]

trafficstars

Describe the bug Refer to this answer on Pydantic's issue by PrettyWood, https://github.com/samuelcolvin/pydantic/issues/2278#issuecomment-764610380

To Reproduce

Example schema:

Pet:
  type: object
  properties:
    name:
      type: string
    pet_type:
      $ref: "#/components/schemas/PetType"

PetType:
  type: string
  enum:
    - dog
    - cat
    - bird
    - fish

Used commandline:

$ datamodel-codegen --input pets.yaml --output pets.py

Expected behavior The generated model for PetType should inherit the specified data type it expects its values to be in: class PetType(str, Enum):

Instead, it currently generates this: class PetType(Enum):

When serialized to JSON, Python would give the error: Object of type PetType is not JSON serializable

Version:

• OS: Mac OSX
• Python version: 3.9.10
• datamodel-code-generator version: 0.11.16

[koxudaxi]

Hi @deanq I'm sorry for my late reply.

Your example will be generated to

from enum import Enum
from typing import Optional

from pydantic import BaseModel


class PetType(Enum):
    dog = 'dog'
    cat = 'cat'
    bird = 'bird'
    fish = 'fish'


class Pet(BaseModel):
    name: Optional[str] = None
    pet_type: Optional[PetType] = None

I can serialize the model to JSON. Could you please share the example code to reproduce the error?

I have checked the link. In the issue, the enum value is used as a dict key. Have you used the generated enum model for the same way?

[pblitz-pg]

We're facing the exact same issue, but I guess there might be some context missing. If this is a different case @deanq , please let me know and i'll open this as a separate one!

This is my sample that produces the above mentioned error:

obj = Pet(pet_type="dog")
marshalled_dict = obj.dict()
json.dumps(marshalled_dict) 
# Fails with TypeError: Object of type PetType is not JSON serializable

In real life applications we often have the need to handle dicts instead of straight up json (testing is our most common case)

if we adapt the generated model

class Pet(BaseModel):
    name: Optional[str] = None
    pet_type: Optional[PetType] = None

    class Config:
        use_enum_values = True

the above code works fine. However, we haven't found a good way of doing this via datamodel-code-generator config.

The other way to make this work, as @deanq already said is

class PetType(str, Enum):
    dog = 'dog'
    cat = 'cat'

which is probably the cleaner way. Again, this would/should/could be done during generation.

[koxudaxi]

@pblitz-pg Thank you for your explanation. OK, I will add a CLI option to support str inherited enum model.

[koxudaxi]

The feature had been released as the --use-subclass-enum option.

  --use-subclass-enum   Define Enum class as subclass with field type when
                        enum has type (int, float, bytes, str)

# generated by datamodel-codegen:
#   filename:  openapi.yaml
#   timestamp: 2022-12-28T16:38:03+00:00

from __future__ import annotations

from enum import Enum


class EnumAbc(int, Enum):
    A = 1
    B = 2
    C = 3


class StringEnumAbc(str, Enum):
    a = 'a'
    b = 'b'
    c = 'c'

oepnapi.yaml

components:
  schemas:
    EnumAbc:
      type: integer
      enum: [1, 2, 3]
      x-enum-varnames: [A, B, C]
    StringEnumAbc:
      type: string
      enum: [ "a", "b", "c" ]