Search code examples
pythonenumspython-dataclasses

How to use enum value in asdict function from dataclasses module


I have a dataclass with a field template of type Enum. When using the asdict function it converts my dataclass to a dictionary. Is it possible to use the value attribute of FoobarEnum to return the string value instead of the Enum object?

My initial idea was to use the dict_factory=dict parameter of the asdict function and provide my own factory but I couldn't figure out how to do this.

from dataclasses import dataclass, asdict
from enum import Enum


@dataclass
class Foobar:
  name: str
  template: "FoobarEnum"


class FoobarEnum(Enum):
  FIRST = "foobar"
  SECOND = "baz"


foobar = Foobar(name="John", template=FoobarEnum.FIRST)

print(asdict(foobar))

Current output:

{'name': 'John', 'template': <FoobarEnum.FIRST: 'foobar'>}

Goal:

{'name': 'John', 'template': 'foobar'}

Solution

  • This can't be done with standard library except maybe by some metaclass enum hack I'm not aware of. Enum.name and Enum.value are builtin and not supposed to be changed.

    The approach of using the dataclass default_factory isn't going to work either. Because default_factory is called to produce default values for the dataclass members, not to customize access to members.

    You can either have the Enum member or the Enum.value as a dataclass member, and that's what asdict() will return.

    If you want to keep an Enum member -not just the Enum.value- as a dataclass member, and have a function converting it to dictionary that returns the Enum.value instead of the Enum member, the correct way to do it is implementing your own method to return the dataclass as a dictionary.

    from dataclasses import dataclass
    from enum import Enum
    
    
    class FoobarEnum(Enum):
        FIRST = "foobar"
        SECOND = "baz"
    
    
    @dataclass
    class Foobar:
        name: str
        template: FoobarEnum
    
        def as_dict(self):
            return {
                'name': self.name,
                'template': self.template.value
            }
    
    
    # Testing.
    print(Foobar(name="John", template=FoobarEnum.FIRST).as_dict())
    # {'name': 'John', 'template': 'foobar'}