Typing dataclass that can only take enum values

user_10468528 0
Viewed 81

I have a dataclass that can take values that are part of an enum.

class MyEnum(Enum):
    A = "valueA"
    B = "valueB"

@dataclass
class MyDataclass:
    value: MyEnum

When creating my dataclass, the types don't match as it is considering str != MyEnum.A.

param = MyDataclass(value="valueA")

Any suggestion on how should MyDataclass be typed instead?

EDIT: The dataclass is being initialized with the string values I receive from another API

python-typing typing enums python
1 Answers

Your class is typed correctly if you really want MyDataclass().value to be a MyEnum value. The problem is that you, as the person instantiating MyDataclass, are responsible for actually passing a value of type MyEnum to __init__.

If your use case requires a string to be passed, you'll have to handle the conversion yourself in __post_init__, which does mean altering the class definition slightly.

from dataclasses import dataclass, InitVar, field


@dataclass
class MyDataclass:
    value: MyEnum = field(init=False)
    value_str: InitVar[str]

    def __post_init__(self, value_str):
        self.value = MyEnum(value_str)

Automatic conversion will require a different library, for example Pydantic or (I believe) attrs; it's not a feature that dataclasses supplies.

The amount of static type protection this gives you is limited, because you are stating that any string can be passed as an argument to __init__/__post_init__, but really only the ones that are valid arguments to MyEnum are valid. This isn't really something that can be expressed using the typing module, as Enum itself isn't designed with static type checking in mind.

user_11268410
Related