Python 中的数据类实例化期间如何将字符串转换为枚举

Question

我有一个数据class DummyClass，其属性 dummy 只能取两个值："foo" 和 "bar"。我想为此属性使用 Enum class，以便在使用无效值时引发错误：

from dataclasses import dataclass, field
from enum import Enum

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

@dataclass
class DummyClass:
    dummy: DummyAttribute

问题是我想使用字符串实例化 DummyClass（不直接使用 DummyAttribute class）。一种方法是使用 __post_init__ 方法：

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

但是我宁愿避免添加 tmp_dummy 属性和 __post_init__ 方法。有没有更简洁的方法将输入字符串转换为 DummyAttribute 枚举？

编辑： @chepner 的回答帮助我找到了一种更简洁的方法：

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)

如果有更好的方法请告诉我！

Answer 1

我认为 __post_init__（如您编辑的那样添加到您的问题中）是您的最佳选择。 dataclass 不支持您想要的那种值转换器。（为此，您可能需要查看 Pydantic 或 attrs 包。）

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)

Answer 2

作为使用 Enum 的替代方法，如果您真正需要的是 str 类型，那么我建议您只使用 str。

请注意，typing.Literal 在这种情况下可能会有帮助：

from dataclasses import dataclass
from typing import Literal


@dataclass
class DummyClass:
    dummy: Literal['foo', 'bar']

    def __post_init__(self):
        # TODO: assert that `self.dummy` is one of valid Literal values
        ...

然后您可以像这样创建您的实例，如果您为字段传递了无效的 Literal 值，您的 IDE 应该立即向您发出警告。

obj1 = DummyClass('bar')
obj2 = DummyClass('blah')  # <-- warning! >:(

如前所述，如果您确实需要验证它是有效的 Literal 值之一，我建议将该逻辑添加到 __post_init__方法。

Python 中的数据类实例化期间如何将字符串转换为枚举

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

python

enums

casting

callable

python-dataclasses