详解enum模块 | CS笔记

其实一直很少在Python编码中用到enum模块，一般都是自己随便定义一组全大写的变量，直到阅读别人的代码，才发现很有必要专门学习enum模块的知识。

from enum import Enum, unique

@unique
class xyz(Enum):
    CAT = 0
    DOG = 1
    FISH = 2
    # CAT = 3 # error, attempt to reuse key
    # BIRD = 2 # error, duplicate value with FISH (by @unique)
    BIRD = '123'

# xyz.CAT = 9 # error, cannot reassign member
print(type(xyz))
print(xyz.CAT)
print(repr(xyz.CAT))
print(xyz.CAT.name)
print(xyz.CAT.value)
print(isinstance(xyz.CAT, xyz))  # true ^OO^
print(type(xyz.CAT))

for it in xyz:
    print(it, it.value)

这种定义方式，与自己定义一组全大写的变量，有好几个区别：

我觉得很有必要全部收下上面的3个好处。下面是代码运行效果：

$ python3 test.py
<class 'enum.EnumMeta'>
xyz.CAT
<xyz.CAT: 0>
CAT
0
True
<enum 'xyz'>
xyz.CAT 0
xyz.DOG 1
xyz.FISH 2
xyz.BIRD 123

>>> food = {}
>>> food[xyz.CAT] = 'fish'
>>> food[xyz.DOG] = 'bone'
>>> food
{<xyz.CAT: 0>: 'fish', <xyz.DOG: 1>: 'bone'}

如果枚举成员的具体值并不重要，可以优雅地使用auto接口自动赋值：

from enum import Enum, auto


class xyz(Enum):
    CAT = auto()
    DOG = auto()
    FISH = auto()
    BIRD = auto()


print(list(xyz))

[<xyz.CAT: 1>, <xyz.DOG: 2>, <xyz.FISH: 3>, <xyz.BIRD: 4>]

from enum import Enum, auto
from pprint import pprint


class AutoName(Enum):
    def _generate_next_value_(name, start, count, last_values):
        """ must be the first function in class """
        return '##'+name
        #return count # 0,1,2,3.... start from 0


class Ordinal(AutoName):
    NORTH = auto()
    SOUTH = auto()
    EAST = auto()
    WEST = auto()

pprint(list(Ordinal))

[<Ordinal.NORTH: '##NORTH'>,
 <Ordinal.SOUTH: '##SOUTH'>,
 <Ordinal.EAST: '##EAST'>,
 <Ordinal.WEST: '##WEST'>]

下面是函数式的Enum类型定义方式，很适合追求代码量少的同学：

>>> animal = Enum('animal', 'CAT DOG BEE FISH')
>>> list(animal)
[<animal.CAT: 1>, <animal.DOG: 2>, <animal.BEE: 3>, <animal.FISH: 4>]
>>> animal.CAT
<animal.CAT: 1>
>>> animal.DOG
<animal.DOG: 2>
>>> animal.BEE.name
'BEE'
>>> animal.FISH.value
4

请注意：每个枚举成员都是枚举类型的一个实例，这与单纯的class variable完全不一样！每个枚举类型的成员，都是一个类型实例对象。请看下面的代码：

from enum import Enum

class PixelFormat(Enum):
    """ Pixel formats.
    """

    RGB = 'raw '
    """ Shape: ``(h,w,3)`` """

    BGR = '24BG'
    """ Shape: ``(h,w,3)`` """

    RGBA = 'ABGR'
    """ Shape: ``(h,w,4)`` """

    GRAY = 'J400'
    """ Shape: ``(h,w)`` """

    I420 = 'I420'
    """ Shape: any of size ``w * h * 3/2`` """

    NV12 = 'NV12'
    """ Shape: any of size ``w * h * 3/2`` """

    YUYV = 'YUY2'
    """ Shape: any of size ``w * h * 2`` """

    UYVY = 'UYVY'
    """ Shape: any of size ``w * h * 2`` """

    def __str__(self):
        return self.name

    def __repr__(self):
        return f'{type(self).__name__}.{self.name}'


print(str(PixelFormat.RGB))
print(repr(PixelFormat.RGB))

上面这段代码，并没有实例化PixelFormat，但实际上，这里面的每个成员已经是对象了，每个对象有name和value这样的属性。上面的代码运行结果如下：

RGB
PixelFormat.RGB

Enum枚举成员支持是否相同的比较，不支持大小比较，与其它类型数据比较一律False：

>>> animal.CAT is animal.CAT
True
>>> animal.CAT is not animal.FISH
True
>>> animal.CAT == animal.CAT
True
>>> animal.CAT != animal.FISH
True
>>> animal.CAT == 1
False
>>> animal.CAT == '1'
False

IntEnum

IntEnum是Enum类型的一个变体，前面介绍了Enum成员直接相互之间比较是否相同，与其它类型对象比较一律False，相互之间也不能比较大小，IntEnum这个变体就是用来解救这些场景的。

>>> from enum import IntEnum
>>> color = IntEnum('color', 'RED GREEN BLUE')
>>> list(color)
[<color.RED: 1>, <color.GREEN: 2>, <color.BLUE: 3>]
>>> color.RED == 1
True
>>> color.RED > color.GREEN
False
>>> color.RED < color.GREEN
True
>>> ['r','g','b'][color.BLUE-1]
'b'
>>> lang = IntEnum('lang', 'C CPP')
>>> list(lang)
[<lang.C: 1>, <lang.CPP: 2>]
>>> lang.C == color.RED
True

IntFlag

IntFlag的IntEnum的基础上，增加了bit位操作，与或非异或（&|~^）。

>>> from enum import IntFlag
>>> opt = IntFlag('opt','R W N')
>>> list(opt)
[<opt.R: 1>, <opt.W: 2>, <opt.N: 4>]
>>> RW = opt.R | opt.W
>>> RW
<opt.W|R: 3>
>>> isinstance(RW, opt)
True

IntFlag对象创建的时候，自动按bit位进行赋值，如上面的opt.N。bit操作得到的结果，依然是IntFlag对象。

>>> from enum import IntFlag, auto
>>> class opt(IntFlag):
...   R = auto()
...   W = auto()
...   N = auto()
...   RW = R|W
...
>>> list(opt)
[<opt.R: 1>, <opt.W: 2>, <opt.N: 4>, <opt.RW: 3>]

Flag

Flag与IntFlag相似，不同之处也与Enum与IntEnum不同之处相似，不同的Flag对象之间不能相互操作，Flag也不能被当作int来对待。