基本用法

在 Python 中,enum 模块提供了一种定义枚举的方法。枚举是一组绑定到唯一常量值的符号名称(成员)。Enum 类是创建枚举的基类。

以下是如何使用 Enum 模块的基本示例:

  1. 定义枚举:
from enum import Enum

class Color(Enum):
    RED = 1
    GREEN = 2
    BLUE = 3
  1. 访问枚举成员:
print(Color.RED)          # 输出: Color.RED
print(Color.RED.name)     # 输出: RED
print(Color.RED.value)    # 输出: 1
  1. 枚举成员的迭代:
for color in Color:
    print(color)
# 输出:
# Color.RED
# Color.GREEN
# Color.BLUE
  1. 枚举成员的比较:
print(Color.RED == Color.RED)     # 输出: True
print(Color.RED == Color.BLUE)    # 输出: False
  1. 通过名称或值获取枚举成员:
print(Color['RED'])       # 输出: Color.RED
print(Color(1))           # 输出: Color.RED
  1. 在条件语句中使用枚举:

    def is_primary_color(color):
       return color in {Color.RED, Color.GREEN, Color.BLUE}
    
    print(is_primary_color(Color.RED))   # 输出: True
    print(is_primary_color(Color(4)))    # 引发 ValueError
    

高级用法

  1. 使用 auto() 自动编号:
from enum import Enum, auto

class Animal(Enum):
    DOG = auto()
    CAT = auto()
    HORSE = auto()

for animal in Animal:
    print(f'{animal.name} = {animal.value}')
# 输出:
# DOG = 1
# CAT = 2
# HORSE = 3
  1. 自定义枚举行为:
from enum import Enum

class Weekday(Enum):
    MONDAY = 1
    TUESDAY = 2
    WEDNESDAY = 3
    THURSDAY = 4
    FRIDAY = 5
    SATURDAY = 6
    SUNDAY = 7

    def is_weekend(self):
        return self in (Weekday.SATURDAY, Weekday.SUNDAY)

print(Weekday.MONDAY.is_weekend())  # 输出: False
print(Weekday.SUNDAY.is_weekend())  # 输出: True
  1. 枚举的子类:
from enum import Enum

class IntEnum(Enum):
    def __int__(self):
        return self.value

class HttpStatus(IntEnum):
    OK = 200
    NOT_FOUND = 404

print(int(HttpStatus.OK))      # 输出: 200
print(HttpStatus.NOT_FOUND)    # 输出: HttpStatus.NOT_FOUND

扩展

注释

通常在业务中会有注释的习惯,即为每个枚举值编写注释,以增加代码的可读性,可以用以下方式实现:

from enum import Enum


class HttpStatus(Enum):
    Not_Method = 405
    Not_Found = 404
    OK = 200
    Bad_Request = 400


class HttpStatusComment(Enum):
    Not_Method = "Method Not Allowed"
    Not_Found = "Not Found"
    OK = "OK"
    Bad_Request = "Bad Request"


print(HttpStatusComment[HttpStatus.OK.name].value)  # 输出:OK
print(HttpStatus[HttpStatusComment.Not_Method.name].value)  # 输出:405

但是这种方式会使代码的后期维护成本增加,可以定义一个枚举基类,增加注释成员,通过继承的方式达到这种效果:

from enum import Enum


class TupleEnum(Enum):
    def __new__(cls, value, description):
        obj = object.__new__(cls)
        obj._value_ = value
        obj.description = description
        return obj

    def __init__(self, value, description):
        self._value_ = value
        self.description = description

    @property
    def value(self):
        return self._value_

    # 通过注释获取枚举值
    @classmethod
    def get_value_by_description(cls, description):
        for member in cls:
            if member.description == description:
                return member.value
        raise ValueError(f"No member with description '{description}' found")


class HttpStatus(TupleEnum):
    Not_Method = (405, "Method Not Allowed")
    Not_Found = (404, "Not Found")
    OK = (200, "OK")
    Bad_Request = (400, "Bad Request")


print(HttpStatus.get_value_by_description("OK"))  # 输出:200
print(HttpStatus.Not_Method.value)  # 输出:405
print(HttpStatus.Not_Found.description)  # 输出:Not Found
print(HttpStatus.Bad_Request.name)  # 输出:Bad_Request
print(HttpStatus["OK"].name)  # 输出:OK
print(HttpStatus(400).value)  # 输出:400

上述示例中,可以在保持原有的方法和功能的基础上新增一个成员:注释,不需要再定义两个枚举类,提高了代码的可读性和可维护性。后续如果有其他需求,如:使用其他的数据类型(字典或其他)、增加其他成员、增加自定义方法等都可以在此基础上修改

总结

使用枚举可以通过为常量提供有意义的名称并将相关的常量分组在一起,提高代码的可读性和可维护性。