Python项目中5个Enum枚举巧妙操作详解

目录

• 用Enum替换字符串
• 值为整数的Enum
• 在数据库模型里使用Enum
• 给Enum添加方法
• 自动赋值的Enum
• 总结一下

刚开始在别人的代码里看到枚举，我其实觉得有点多余。用字符串或者数字常量不也一样吗？为啥非要整个Enum出来，感觉是把简单问题复杂化了。但后来我接手了一个后端项目，里面有大量的常量定义，我才慢慢体会到Enum的好处。它不仅能让代码更整洁，还能强制保证代码的一致性，避免了很多低级错误。

在开始之前，先提一嘴我的开发环境。我现在用的是ServBay，堪称神器。

以前要部署Python环境，版本切来切去特别麻烦，尤其是有时候要维护一个Python 2的老项目，同时又要写Python 3的新功能。ServBay能一键部署好Python环境，而且可以自由切换不同的Python版本，甚至让多个版本同时运行。不管是Python 2.x的老古董还是最新的Python 3.x，它都能轻松搞定。

环境干净又卫生，我才能把心思花在琢磨代码本身上，比如怎么用好Enum这种细节。

下面我就分享几个我在实际项目里用过的Enum模式，没什么高深理论，都是些实实在在帮我解决过问题的小技巧。

这是Enum最基础，也是最常见的用法。以前，我的代码里可能有很多这样的硬编码字符串：

# 不推荐的写法
def check_ticket_status(status_string):
    if status_string == "OPEN":
        print("工单开启")
    elif status_string == "IN_PROGRESS":
        print("工单处理中")
    elif status_string == "CLOSED":
        print("工单已关闭")

这么写也能跑，但问题很大。status参数可以传入任何字符串，一不小心手滑，把CLOSED写成了closed，这个条件就匹配不上了，bug就这么悄无声息地出现了。

换成Enum之后，代码就严谨多了：

from enum import Enum

class TicketState(Enum):
    OPEN = "OPEN"
    IN_PROGRESS = "IN_PROGRESS"
    CLOSED = "CLOSED"

def handle_ticket(state: TicketState):
    if state == TicketState.OPEN:
        print("工单当前状态：开启")
    elif state == TicketState.IN_PROGRESS:
        print("工单当前状态：处理中")
    elif state == TicketState.CLOSED:
        print("工单当前状态：已关闭")

# 调用时必须使用Enum成员，非常安全
handle_ticket(TicketState.IN_PROGRESS)

现在，如果有人想传入一个TicketState里没有定义的值，静态类型检查工具（比如mypy）直接就会报警，代码在运行前就能发现问题。这就从根源上杜绝了拼写错误导致的bug。在处理后端服务各种复杂状态流转时，这个模式特别有用。

我之前维护过一个很老的系统，对接的API返回的状态码全是数字，比如1代表成功，2代表失败。代码里到处是if status_code == 1:这种数字，过段时间别说别人，我自己都看不懂了。

如果用Enum改造了这部分代码，可读性大大提高：

from enum import IntEnum

# 使用IntEnum，它继承自int，行为更像整数
class ApiErrorCode(IntEnum):
    OK = 0
    INVALID_PARAMS = 1001
    SERVER_ERROR = 2001
    AUTH_FAILURE = 3001

def process_api_response(code: int):
    try:
        error_code = ApiErrorCode(code)
    except ValueError:
        print(f"收到未知的错误码: 
")
        return

    if error_code == ApiErrorCode.OK:
        print("请求成功")
    elif error_code == ApiErrorCode.SERVER_ERROR:
        print("服务端内部错误")
    else:
        print(f"业务错误: {error_code.name}")

# 示例调用
process_api_response(0)
process_api_response(3001)
process_api_response(9999) # 测试未知错误码

这里用了IntEnum，它是Enum的一个变体，成员的值都是整数，并且可以直接和整数进行比较。通过ApiErrorCode(code)可以方便地把整数转换为Enum成员，如果转换失败则说明收到了一个未定义的错误码，逻辑清晰了很多。

在数据库里存一些固定选项的字段时，Enum也特别好用。比如，文章的发布状态，可能只有“草稿”、“已发布”、“已归档”三种。直接在数据库里存字符串当然可以，但在代码的模型层用Enum来约束会更安全。

from enum import Enum
from pydantic import BaseModel
from typing import List

class ArticleStatus(str, Enum):
    DRAFT = "draft"
    PUBLISHED = "published"
    ARCHIVED = "archived"

class Article(BaseModel):
    id: int
    title: str
    status: ArticleStatus
    tags: List[str]

# 创建一个文章实例
new_article = Article(
    id=101, 
    title="Python Enum使用技巧", 
    status=ArticleStatus.PUBLISHED,
    tags=["python", "enum"]
)

# ArticleStatus.PUBLISHED可以直接作为字符串参与JSON序列化
print(new_article.model_dump_json())
# 输出: {"id":101,"title":"Python Enum使用技巧","status":"published","tags":["python","enum"]}

注意这里的写法：class ArticleStatus(str, Enum):。通过多继承一个str，这个Enum的成员就拥有了字符串的所有特性。这样做的好处是，它能和Pydantic、SQLAlchemy这类库无缝集成，序列化成JSON或者存入数据库时，它表现得就像一个普通的字符串，但在代码层面，它又保留了Enum的类型安全和代码补全等优点。

这个玩法是我后来才发现的，在Enum内部直接定义方法。一开始我觉得挺奇怪的，Enum不是用来放常量的吗？但如果某些逻辑和这个Enum强绑定，那把方法写在里面就非常合适。

from enum import Enum

class Environment(Enum):
    DEV = "development"
    TEST = "testing"
    PROD = "production"

    @property
    def is_production(self) -> bool:
        return self == Environment.PROD

    def get_db_url(self) -> str:
        # 实际项目中，这里可能会从配置中读取
        db_urls = {
            "development": "sqlite:///dev.db",
            "testing": "sqlite:///test.db",
            "production": "postgresql://user:pass@host/prod_db"
        }
        return db_urls[self.value]

# 使用起来非常直观
current_env = Environment.PROD
print(f"当前环境: {current_env.value}")
print(f"是否为生产环境: {current_env.is_production}")
print(f"数据库连接: {current_env.get_db_url()}")

每个Environment的成员自己就知道对应的数据库连接，还能判断自己是不是生产环境。这比在外面写一堆if current_env == "production":要优雅和内聚得多。

有时候，Enum成员的名字和值是一样的，一个个手动敲赋值会显得很啰嗦。比如，我们要定义一组操作权限。

enum模块里的auto()可以帮我们偷懒：

from enum import Enum, auto

class Permission(Enum):
    # 重写这个特殊方法，让auto()返回成员的大写字符串名
    def _generate_next_value_(name, start, count, last_values):
        return name.upper()

    READ = auto()
    WRITE = auto()
    EXECUTE = auto()
    DELETE = auto()

# 验证一下
print(f"Read permission value: {Permission.READ.value}")  # 输出: Read permission value: READ
print(f"Delete permission value: {Permission.DELETE.value}") # 输出: Delete permission value: DELETE

通过重写内部的_generate_next_value_方法，我们可以自定义auto()的行为。这里我让它自动将成员名转为大写作为值，这样定义Permission时就非常简洁，完全不用写等号右边的部分。

以上就是我实际在项目中用得比较多的五种Enum模式。刚开始我觉得Enum是形式主义，但随着项目越来越复杂，我越来越喜欢用它。它让代码变得更严格，但也更清晰。

当然，如果代码很简单，或者只有一两个常量，我可能还是会直接用字符串。但对于中大型项目，特别是多人协作的时候，Enum真的能帮我们避免很多因为常量写错、不一致导致的傻瓜式错误。

总的来说，Enum不是什么高深的技巧，但用在对的地方，确实能让代码质量上一个台阶。就像一个好用的开发工具 ServBay 能帮我们摆平环境配置的烦恼一样，好的代码习惯能让我们在几个月后回头看自己的代码时，少一些头疼，多一份舒心。

以上就是Python项目中5个Enum枚举巧妙操作详解的详细内容，更多关于Python Enum枚举的资料请关注风君子博客其它相关文章！

您可能感兴趣的文章:

• Python中模拟enum枚举类型的5种方法分享
• python模拟enum枚举类型的方法小结
• 浅谈Python 的枚举 Enum
• Python枚举之Enum模块的原理与应用详解
• Python Enum枚举类的定义及使用场景最佳实践
• Python枚举(enum)模块使用的全面指南