3

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?

andrewshih
  • 477
  • 6
  • 12
  • Does it work as you want if you just do `print(Mode.MODE_A)`? – Barmar Mar 18 '21 at 00:52
  • @Barmar Yes, by print, it shows like this `>>> print(Mode.MODE_A)` then `Mode.BUDGET` – andrewshih Mar 18 '21 at 00:55
  • 2
    If you're worrying about how this prints in the command line interpreter, then you are worrying about nothing. For your users, you will do formatting, and it will look correct. – Tim Roberts Mar 18 '21 at 01:01
  • What I want to have is a way to output a dict that has `{'mode': 'MODE_A'}` instead of {`'mode': }` by that by that `to_dict` function. – andrewshih Mar 18 '21 at 01:05

1 Answers1

1

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}
Ethan Furman
  • 63,992
  • 20
  • 159
  • 237