如何修改枚举属性，实现dataclass内容的直接访问？

你定义了如下Dataclass和枚举类：

from dataclasses import dataclass
from abc import ABC
from typing import Union, ClassVar
from enum import Enum

@dataclass
class BaseObject(ABC):
    obj: int
    data: Union[str, int, float, None] = None
    datatype: ClassVar[type]
    datalen: ClassVar[int]
    readonly: bool = True

@dataclass
class StrObject(BaseObject):
    datatype = str
    datalen = 16

# 生成唯一ID的示例函数
def next_id():
    global _id_counter
    _id_counter += 1
    return _id_counter

_id_counter = 0

class ObjectsPSI9000(Enum):
    DEVICE_TYPE = StrObject(next_id())
    DEVICE_SN = StrObject(next_id())
    # 假设已定义FloatObject
    # NOMINAL_VOLT = FloatObject(next_id())
    # NOMINAL_CURR = FloatObject(next_id())

当前访问属性时需要写DEVICE_TYPE.value.obj，操作繁琐，希望能“解包”Dataclass内容，简化访问方式。

解决方案 #

方法一：自定义代理枚举基类 #

创建一个继承自Enum的基类，通过重写__getattr__方法，让枚举成员直接代理其value的属性：

class ProxyEnum(Enum):
    def __getattr__(self, name):
        try:
            # 优先访问枚举成员自身的属性
            return super().__getattr__(name)
        except AttributeError:
            # 自身没有该属性时，从value中获取
            return getattr(self.value, name)

修改你的枚举类继承这个基类：

class ObjectsPSI9000(ProxyEnum):
    DEVICE_TYPE = StrObject(next_id())
    DEVICE_SN = StrObject(next_id())

之后就可以直接通过枚举成员访问Dataclass的属性：

print(ObjectsPSI9000.DEVICE_TYPE.obj)
print(ObjectsPSI9000.DEVICE_TYPE.data)

方法二：让BaseObject实现描述符协议 #

让BaseObject实现__get__方法，使枚举成员直接等同于对应的Dataclass实例：

@dataclass
class BaseObject(ABC):
    obj: int
    data: Union[str, int, float, None] = None
    datatype: ClassVar[type]
    datalen: ClassVar[int]
    readonly: bool = True

    def __get__(self, instance, owner):
        return self

此时枚举成员本身就是BaseObject的实例，直接访问属性即可：

print(ObjectsPSI9000.DEVICE_TYPE.obj)

方法三：优化IDE自动补全（可选） #

如果使用方法一，为了让IDE识别到Dataclass的属性并提供自动补全，可以重写__dir__方法，将value的属性加入枚举成员的属性列表：

class ProxyEnum(Enum):
    def __getattr__(self, name):
        try:
            return super().__getattr__(name)
        except AttributeError:
            return getattr(self.value, name)
    
    def __dir__(self):
        # 合并枚举自身属性和value的属性
        return super().__dir__() + dir(self.value)

内容的提问来源于stack exchange，提问作者Paweł Pietraszko