andrewshih
andrewshih

Reputation: 527

Output enum name as enum value in a dict showing class attributes in Python

I have an enum and a dataclass as following.

class Mode(Enum):
    MODE_A = auto()
    MODE_B = auto()

@dataclass
class BaseDataClass:
    a: int
    b: int
    c: int
    mode: Mode

    def to_dict(self):
        return vars(self)

the to_dict function displays

b = BaseDataClass(a=1, b=2, c=3, mode=Mode.MODE_A)
b.to_dict()

> {'a': 1, 'b': 2, 'c': 3, 'mode': <Mode.MODE_A: 1>}

I'd like to output the mode name as its value like 'mode': 'MODE_A'. I tried to use dunder methods like __str__ and __repr__ in the enum class but it didn't work. AS it is supposed not to edit __dict__ values, Does anyone suggest a good way to achieve this?

Upvotes: 3

Views: 857

Answers (1)

Ethan Furman
Ethan Furman

Reputation: 69288

Not sure what you tried for your __repr__, but this works:

class Mode(Enum):
    MODE_A = auto()
    MODE_B = auto()
    #
    def __repr__(self):
        return self.name

and in use:

>>> {'mode': Mode.MODE_A}
{'mode': MODE_A}

Upvotes: 1

Related Questions