Search code examples
pythonenumscastingcallablepython-dataclasses

How to cast a string to an Enum during instantiation of a dataclass in Python


I have a dataclass DummyClass with an attribute dummy that can only take two values: "foo" and "bar". I would like to use an Enum class for this attribute so that an error is raised when using an invalid value:

from dataclasses import dataclass, field
from enum import Enum

class DummyAttribute(Enum):
    FOO = "foo"
    BAR = "bar"

@dataclass
class DummyClass:
    dummy: DummyAttribute

The problem is that I would like to instantiate DummyClass using a string (not using the DummyAttribute class directly). One way to do it would be using the __post_init__ method:

@dataclass
class DummyClass:
    dummy: DummyAttribute = field(init=False)
    tmp_dummy: str
    def __post_init__(self):
        self.dummy = DummyAttribute(self.tmp_dummy)

However I would prefer to avoid adding a tmp_dummy attribute and a __post_init__ method. Is there any more concise way of casting the input string to the DummyAttribute enum?


Edit: @chepner answer helped me find an even more concise way of doing it:

from dataclasses import dataclass, field
from enum import Enum

class DummyAttribute(Enum):
    FOO = "foo"
    BAR = "bar"

@dataclass
class DummyClass:
    dummy: str | DummyAttribute
    def __post_init__(self):
        self.dummy = DummyAttribute(self.dummy)

Please let me know if there is a better way!


Solution

  • I think __post_init__ (as you edited added to your question) is your best alternative. dataclass doesn't support the kind of value transformer you want. (For that, you might want to look at Pydantic or the attrs package.)

    from dataclasses import dataclass, InitVar
    from enum import Enum
    
    class DummyAttribute(Enum):
        FOO = "foo"
        BAR = "bar"
    
    
    @dataclass
    class DummyClass:
        dummy: DummyAttribute = field(init=False)
        dummyStr: InitVar[str]
    
        def __post_init__(self, dummyStr):
            self.dummy = DummyAttribute(dummyStr)