详解enum模块(Python)

Last Updated: 2024-05-08 07:56:45 Wednesday

-- TOC --

一直很少在Python编码中用到enum模块,基本都是自己随便定义一组全大写的变量,直到阅读别人的代码,才发现很有必要专门学习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)

继承自Enum的类xyz,就成了一个枚举类型。

这种定义方式,与自己定义一组全大写的变量,有如下好处:

  1. 不能重复出现枚举成员name,比如上例中的CAT;
  2. 如果使用了@unique装饰器,成员不能具有相同的值value,比如上例注释中的BIRD;
  3. 不会出现不小心修改一个枚举成员的值value的情况。

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

<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

Enum枚举类型可以哈希,Enum is hashable:

>>> 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))

运行效果如下,auto从1开始:

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

auto接口的返回值可以被自定义:

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

Enum对象可以hashable,IntEnum对象还可以被当作int来使用。

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对象。

在上一个case,直接在定义的时候,做bit操作:

>>> 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来对待。

本文链接:https://cs.pynote.net/sf/python/202209212/

-- EOF --

-- MORE --