如何向Python枚举类型中添加成员子集?

5

假设我有一个像这样的 Python 枚举:

from enum import Enum

class FRUIT(Enum):
    APPLE = 1
    BANANA = 2
    LEMON = 3
    ORANGE = 4

我希望能以有用的方式“子集”这些内容,比如可以说:

if fruit in FRUIT.CITRUS:
    make_juice()

我在某处定义了:CITRUS = {LEMON, ORANGE}

我希望将子集保留为主枚举的属性,因为这样可以使子集的使用有上下文。

我知道我可以像这样做,但我强烈希望避免方法调用符号。此外,每次需要重建集合似乎是浪费的:

@classmethod
def CITRUS(cls):
    return {cls.LEMON, cls.ORANGE}

有没有一种方法可以在枚举元类完成其工作后添加类属性而不会破坏任何东西?

你可以添加任意数量的类属性。你的类仅仅是从枚举(Enum)中继承而来。问题在于你想让CITRUS变得有用,这很可能会与Enum内置方法冲突,因为这些方法操作该级别下的所有属性。 - Prune
3个回答

5

由于 CITRUS 不是指一种水果本身,而是一种水果类型,因此创建一个包含水果类型成员的单独的 Enum 子类更有意义:

class FRUIT_TYPE(Enum):
    CITRUS = {FRUIT.LEMON, FRUIT.ORANGE}

这样你就可以像这样做:

fruit = FRUIT.LEMON
if fruit in FRUIT_TYPE.CITRUS.value:
     make_juice()

然而,使用FRUIT_TYPE.CITRUS.value进行成员检查看起来很麻烦。为了允许对FRUIT_TYPE.CITRUS本身进行成员检查,您还可以将FRUIT_TYPE作为set的一个子类:

class FRUIT_TYPE(set, Enum):
    CITRUS = {FRUIT.LEMON, FRUIT.ORANGE}

以便您可以执行以下操作:
fruit = FRUIT.LEMON
if fruit in FRUIT_TYPE.CITRUS:
     make_juice()

2
FRUIT_TYPE.CITRUS.value 看起来很糟糕。尝试添加 set,例如 class FRUIT_TYPE(set, Enum): -- 这样,CITRUS 也是一个 set,可以这样说 if fruit in FRUIT_TYPE.CITRUS - Ethan Furman
感谢@EthanFurman提供的绝佳提示。我已按照建议更新了答案。 - blhsing

4

如果你不需要为枚举值指定特定的值,那么使用enum.IntFlag 枚举,而不是简单地使用 Enum,很可能可以获得所需的所有功能。

只需将枚举类声明为 IntFlag,然后就可以自由使用 &, | 和其他按位运算符来获得所需的行为:

In [1]: import enum                                                                                                                  

In [2]: class Fruit(enum.IntFlag): 
   ...:     APPLE = 1 
   ...:     BANANA = 2 
   ...:     LEMON = 4 
   ...:     ORANGE = 8 
In [4]: CITRUS = Fruit.LEMON | Fruit.ORANGE                                                                                          


In [6]: for fruit in Fruit: 
   ...:     if fruit & CITRUS: 
   ...:         print(f"making_juice({fruit.name})") 
   ...:          
   ...:   

这不允许直接在"CITRUS"上进行交互,并且需要一个过滤器模式,就像我上面使用的一样。

然而,就在几个星期前,我正需要这个功能,并且可以在枚举类中实现一个__iter__方法,直接进行筛选:

    def __iter__(self):
        for element in self.__class__:
            if self & element:
                yield element

如果我们简单地将其插入上述枚举中:
In [8]: class Fruit(enum.IntFlag): 
   ...:     APPLE = 1 
   ...:     BANANA = 2 
   ...:     LEMON = 4 
   ...:     ORANGE = 8 
   ...:      
   ...:     def __iter__(self): 
   ...:         for element in self.__class__: 
   ...:             if self & element: 
   ...:                 yield element 
   ...:                                                                                                                              

In [9]: CITRUS = Fruit.LEMON | Fruit.ORANGE                                                                                          

In [10]: for fruit in CITRUS: 
    ...:     print (fruit.name) 
    ...:                                                                                                                             
LEMON
ORANGE
__iter__ 方法并不会与 Fruit 类本身的迭代产生冲突,因为它使用枚举元类 EnumMeta 中的 __iter__ 方法,并且,可以看到,“或”子集的枚举能够正确调用这个方法。也就是说,如果需要,您只需编写适当的 __len____contains__ 方法,就可以拥有从子集中预期的所有功能。我在一个个人项目中使用此代码,它非常好用:https://github.com/jsbueno/terminedia/blob/9714d6890b8336678cd10e0c6275f56392e409ed/terminedia/values.py#L51(尽管现在在枚举下方声明的“unicode_effects”只是一个普通的集合,但现在提到了它,我想我将编写 __contains__ 并使用它来替代这个集合)。

1
你可以在类本身中定义CITRUSCITRUS = LEMON | ORANGE。然后FRUIT.ORANGE in FRUIT.CITRUS就可以工作了。请注意仍需要使用您的__iter__方法(或使用aenum.EnumFRUIT.CITRUS的迭代功能实现。 - Ethan Furman

3

更新:还有一件需要考虑的事情,我认为更好的做法是将子集成员定义为每个枚举成员的属性,例如:

fruit = FRUIT.ORANGE  # ---or whatever, probably in far away code---
...
if fruit.is_citrus:
    make_juice()

这些可以被定义为类上的@property,并且不会受到下面提到的可变性问题的影响。

class FRUIT(Enum):
    APPLE = 1
    BANANA = 2
    LEMON = 3
    ORANGE = 4

    @property
    def is_citrus(self):
        return self in frozenset((FRUIT.LEMON, FRUIT.ORANGE))


感谢其他回答者提供的非常有用的观点。在考虑其他答案后,这是我最终采取的做法,接下来是我的理由:
from enum import Enum

class FRUIT(Enum):
    APPLE = 1
    BANANA = 2
    LEMON = 3
    ORANGE = 4

FRUIT.CITRUS_TYPES = frozenset((FRUIT.LEMON, FRUIT.ORANGE))


这个很好地运作了,也没有破坏其他的“枚举”行为(让我很惊讶)。
# ---CITRUS_TYPES subset has desired behavior---
>>> FRUIT.LEMON in FRUIT.CITRUS_TYPES
True
>>> FRUIT.APPLE in FRUIT.CITRUS_TYPES
False
>>> "foobar" in FRUIT.CITRUS_TYPES
False

# ---CITRUS_TYPES has not become a member of FRUIT enum---
>>> tuple(FRUIT)
(FRUIT.APPLE: 1>, <FRUIT.BANANA: 2>, <FRUIT.LEMON: 3>, <FRUIT.ORANGE: 4>)
>>> FRUIT.APPLE in FRUIT
True
>>> FRUIT.CITRUS_TYPES in FRUIT
DeprecationWarning: using non-Enums in containment checks will raise TypeError in Python 3.8
False

# ---CITRUS_TYPES not reported by dir(FRUIT)---
>>> dir(FRUIT)
['APPLE', 'BANANA', 'LEMON', 'ORANGE', '__class__', '__doc__', '__members__', '__module__']

# ---But it does appear on FRUIT.__dict__---
FRUIT.__dict__ == {
    '_generate_next_value_': <function Enum._generate_next_value_ at 0x1010e9268>, 
    '__module__': '__main__',
    '__doc__': 'An enumeration.',
    '_member_names_': ['APPLE', 'BANANA', 'LEMON', 'ORANGE'],
    '_member_map_': OrderedDict([
        ('APPLE', <FRUIT.APPLE: 1>),
        ('BANANA', <FRUIT.BANANA: 2>),
        ('LEMON', <FRUIT.LEMON: 3>),
        ('ORANGE', <FRUIT.ORANGE: 4>)
    ]),
    '_member_type_': <class 'object'>,
    '_value2member_map_': {
        1: <FRUIT.APPLE: 1>,
        2: <FRUIT.BANANA: 2>,
        3: <FRUIT.LEMON: 3>,
        4: <FRUIT.ORANGE: 4>,
    },
    'APPLE': <FRUIT.APPLE: 1>,
    'BANANA': <FRUIT.BANANA: 2>,
    'LEMON': <FRUIT.LEMON: 3>,
    'ORANGE': <FRUIT.ORANGE: 4>,
    '__new__': <function Enum.__new__ at 0x1010e91e0>,
    'CITRUS_TYPES': frozenset({<FRUIT.LEMON: 3>, <FRUIT.ORANGE: 4>})
}

看起来它将 CITRUS_TYPES 存储在类上,但出于某些原因,将其隐藏在 dir() 中。

然而,它确实存在漏洞,因为添加的属性是可变的,就像任何其他类属性一样;如果客户端代码的某个部分将 FRUIT.CITRUS_TYPES 赋值,FRUIT 不会抱怨,这当然会破坏事情。这种行为与 Enum 成员不同,后者在尝试赋值时引发 AttributeError

我认为这可以通过将其变为 classproperty 来解决,我尝试过,但我的早期尝试没有防止变异。那里描述的更复杂的 classproperty 实现可能有效,但我目前对简单的方法感到满意。

@blhsing提出了一个有趣的问题,即在FRUIT上使用这样的属性是否有意义。我理解他的观点,也许会采纳他的看法,但我目前的观点是将“与水果相关”的特性局限于单个导入名称最适合我。可以考虑FRUIT作为严格的水果类型和子集的集合,因此是一个不同的集合。对于我的当前目的,我发现这种严谨性没有什么回报,更喜欢将FRUIT视为相关常量值的集合,包括成员和子集。当然,你的情况可能不同。就像我说的,我可能会采纳他的观点。


网页内容由stack overflow 提供, 点击上面的
可以查看英文原文,
原文链接