Python枚举可防止无效的属性赋值

Python enum prevent invalid attribute assignment

当我使用函数API创建枚举时,我会返回一个允许任意分配的枚举对象(即它有一个uuu dict_uuuuuuu):

1
2
e = enum.Enum('Things',[('foo',1),('bar',2)])
e.baz = 3

该项不显示在列表中:

1
2
list(e)
[<foo.foo: 1>, <foo.bar: 2>]

但仍可以参考:

1
if thing == e.baz: ...

现在,虽然看起来不太可能发生这种情况,但我想使用枚举的原因之一是为了防止拼写错误和字符串文本,以及在导入模块或尽可能早地捕获这些内容。

有没有一种方法可以动态地构建一个枚举,它的行为更像一个不允许分配任意属性的槽对象?


不一定容易,但可能。我们需要创建一个新的EnumMeta类型1,正常创建Enum,然后在创建Enum后重新分配类型:

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
from enum import Enum, EnumMeta

class FrozenEnum(EnumMeta):
   "prevent creation of new attributes"
    def __getattr__(self, name):
        if name not in self._member_map_:
            raise AttributeError('%s %r has no attribute %r'
                % (self.__class__.__name__, self.__name__, name))
        return super().__getattr__(name)

    def __setattr__(self, name, value):
        if name in self.__dict__ or name in self._member_map_:
            return super().__setattr__(name, value)
        raise AttributeError('%s %r has no attribute %r'
                % (self.__class__.__name__, self.__name__, name))

class Color(Enum):
    red = 1
    green = 2
    blue = 3

Color.__class__ = FrozenEnum

使用中:

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
>>> type(Color)
<class 'FrozenEnum'>

>>> list(Color)
[<Color.red: 1>, <Color.green: 2>, <Color.blue: 3>]

>>> Color.blue
<Color.blue: 3>

>>> Color.baz = 3
Traceback (most recent call last):
  ...
AttributeError: FrozenEnum 'Color' has no attribute 'baz'

>>> Color.baz
Traceback (most recent call last):
  ...
AttributeError: 'FrozenEnum' object has no attribute 'baz'

尝试重新分配成员仍然会产生更友好的错误:

1
2
3
4
>>> Color.blue = 9
Traceback (most recent call last):
  ...
AttributeError: Cannot reassign members.

为了使类重新分配更容易一些,我们可以编写一个修饰器来封装流程:

1
2
3
def freeze(enum_class):
    enum_class.__class__ = FrozenEnum
    return enum_class

使用中:

1
2
3
4
5
@freeze
class Color(Enum):
    red = 1
    green = 2
    blue = 3

请注意,仍然可以覆盖普通属性,例如函数:

1
2
3
4
5
6
7
@freeze
class Color(Enum):
    red = 1
    green = 2
    blue = 3
    def huh(self):
        print("Huh, I am %s!" % self.name)

使用中:

1
2
3
4
5
6
7
8
9
10
11
12
13
>>> Color.huh
<function Color.huh at 0x7f7d54ae96a8>

>>> Color.blue.huh()
Huh, I am blue!

>>> Color.huh = 3
>>> Color.huh
3
>>> Color.blue.huh()
Traceback (most recent call last):
  ...
TypeError: 'int' object is not callable

即使这样也可能被阻止,但我现在还是把它留给别人做练习。

1这只是我见过的第二个需要子类化EnumMeta的情况。另一方面,见this question

公开:我是python stdlib Enumenum34backport和advanced enumeration(aenum库的作者。


要使枚举类完全"只读",只需要使用一个使用__setattr__钩子的元类,它可以防止所有属性赋值。因为元类是在创建后附加到类上的,所以分配正确的枚举值没有问题。

像Ethan的答案一样,我使用EnumMeta类作为自定义元类的基础:

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
from enum import EnumMeta, Enum

class FrozenEnumMeta(EnumMeta):
   "Enum metaclass that freezes an enum entirely"
    def __new__(mcls, name, bases, classdict):
        classdict['__frozenenummeta_creating_class__'] = True
        enum = super().__new__(mcls, name, bases, classdict)
        del enum.__frozenenummeta_creating_class__
        return enum

    def __call__(cls, value, names=None, *, module=None, **kwargs):
        if names is None:  # simple value lookup
            return cls.__new__(cls, value)
        enum = Enum._create_(value, names, module=module, **kwargs)
        enum.__class__ = type(cls)
        return enum

    def __setattr__(cls, name, value):
        members = cls.__dict__.get('_member_map_', {})
        if hasattr(cls, '__frozenenummeta_creating_class__') or name in members:
            return super().__setattr__(name, value)
        if hasattr(cls, name):
            msg ="{!r} object attribute {!r} is read-only"
        else:
            msg ="{!r} object has no attribute {!r}"
        raise AttributeError(msg.format(cls.__name__, name))

    def __delattr__(cls, name):
        members = cls.__dict__.get('_member_map_', {})
        if hasattr(cls, '__frozenenummeta_creating_class__') or name in members:
            return super().__delattr__(name)
        if hasattr(cls, name):
            msg ="{!r} object attribute {!r} is read-only"
        else:
            msg ="{!r} object has no attribute {!r}"
        raise AttributeError(msg.format(cls.__name__, name))

class FrozenEnum(Enum, metaclass=FrozenEnumMeta):
    pass

上面区分了已经可用的属性和新的属性,以便于诊断。它还阻止属性删除,这可能同样重要!

它还为枚举提供元类和FrozenEnum基类;使用它来代替Enum

冻结样本Color枚举:

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
>>> class Color(FrozenEnum):
...     red = 1
...     green = 2
...     blue = 3
...
>>> list(Color)
[<Color.red: 1>, <Color.green: 2>, <Color.blue: 3>]
>>> Color.foo = 'bar'
Traceback (most recent call last):
    # ...
AttributeError: 'Color' object has no attribute 'foo'
>>> Color.red = 42
Traceback (most recent call last):
    # ...
Cannot reassign members.
>>> del Color.red
Traceback (most recent call last):
    # ...
AttributeError: Color: cannot delete Enum member.

请注意,所有属性更改都是不允许的,不允许新属性,删除操作也被阻止。当名称是枚举成员时,我们委托给原始的EnumMeta处理以保持错误消息的稳定性。

如果枚举使用更改枚举类属性的属性,则必须将这些属性白名单,或者允许设置以单个下划线开头的名称;在__setattr__中,确定允许设置哪些名称,并为这些异常使用super().__setattr__(name, value),就像代码现在区分类构造和稍后使用标志属性进行更改。

上面的类可以像Enum()那样用于以编程方式创建枚举:

1
e = FrozenEnum('Things', [('foo',1), ('bar',2)]))

演示:

1
2
3
4
5
6
7
>>> e = FrozenEnum('Things', [('foo',1), ('bar',2)])
>>> e
<enum 'Things'>
>>> e.foo = 'bar'
Traceback (most recent call last):
    # ...
AttributeError: Cannot reassign members.