Logo Questions Linux Laravel Mysql Ubuntu Git Menu
 

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'}
like image 469
Nepo Znat Avatar asked Apr 21 '20 07:04

Nepo Znat


3 Answers

Actually you can do it. asdict has keyword argument dict_factory which allows you to handle your data there:

from dataclasses import dataclass, asdict
from enum import Enum


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


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


def custom_asdict_factory(data):

    def convert_value(obj):
        if isinstance(obj, Enum):
            return obj.value
        return obj

    return dict((k, convert_value(v)) for k, v in data)


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

print(asdict(foobar, dict_factory=custom_asdict_factory))
# {'name': 'John', 'template': 'foobar'}
like image 189
doomatel Avatar answered Oct 19 '22 08:10

doomatel


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'}

like image 45
bad_coder Avatar answered Oct 19 '22 08:10

bad_coder


from dataclasses import dataclass, asdict
from enum import Enum


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


@dataclass
class Foobar:
    name: str
    template: FoobarEnum


def my_dict(data):

    return {
        field: value.value if isinstance(value, Enum) else value
        for field, value in data
    }


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

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

assert asdict(foobar, dict_factory=my_dict) == data
like image 6
Evgeniy_Burdin Avatar answered Oct 19 '22 06:10

Evgeniy_Burdin