用好 Django 为你提供的枚举类

2021-09-27 Python Web

假设我们定义了一个付款单（Payment）模型，该模型定义了一个status字段，用不同的数字代表该付款单不同的状态：

0: 代表付款单已被创建
1: 代表付款单已由创建人提交
2: 代表付款单已被财务部门审核完毕
3: 代表付款单已被财务部门完成付款
4: 代表付款单已被财务部门驳回

那么可以这样编写：

from django.db import models

class Payment(models.Model):

    status_choices = [
        (0, '已创建'),
        (1, '已提交'),
        (2, '已审核'),
        (3, '已支付'),
        (4, '已驳回')
    ]

    status = models.SmallIntegerField("状态", choices=status_choices, default=0)

这样编写的代码有很大的问题：如果你只是简单地用数字代表了付款单状态的话，那么你的项目代码中会充满大量不明所以的数字。

比如：查询所有“已支付”的付款单，Payment.objects.filter(status=3)，你在阅读代码时可能已经知道“3”代表“已支付”，但其他人阅读代码时只会是一头雾水。

而且，如果以后需要修改status_choices，那重构代码将会非常困难。

正确的做法是给这些数字“取个名字”，最简单的方法是使用Python标准库中的枚举类：

from enum import Enum

from django.db import models

class Payment(models.Model):

    class Statuses(Enum):
        Created = 0
        Submitted = 1
        Reviewed = 2
        Paid = 3
        Rejected = 4

    status_choices = [
        (Statuses.Created.value, '已创建'),
        (Statuses.Submitted.value, '已提交'),
        (Statuses.Reviewed.value, '已审核'),
        (Statuses.Paid.value, '已支付'),
        (Statuses.Rejected, '已驳回')
    ]

    status = models.SmallIntegerField("状态", choices=status_choices, default=Statuses.Created.value)

我建议将Statuses直接定义在模型对象中，这样Statuses作为模型类的类属性，方便使用。

现在你可以将Payment.objects.filter(status=3)改成Payment.objects.filter(status=Payment.Statuses.Paid.value)，代码的可读性增加了。

但是使用enum.Enum仍然不是最完美的解决方案，还是有一些缺点：

如果要获取Payment.Statuses.Paid的描述文本，那还是需要从status_choices中获取。
Payment.Statuses.Paid并不是数字，Payment.Statuses.Paid.value才是数字。

最好的解决方法，是使用Django为你提供好的枚举类：

from django.db import models

class Payment(models.Model):

    class Statuses(models.IntegerChoices):
        Created = 0, "已创建"
        Submitted = 1, "已提交"
        Reviewed = 2, "已审核"
        Paid = 3, "已支付"
        Rejected = 4, "已驳回"

    status = models.SmallIntegerField("状态", choices=Statuses.choices, default=Statuses.Created)

现在Statuses类的每个成员的值为一个二元元组：(真实值, 描述文本)，描述文本将作为Statuses成员的label属性。 (如果没有提供描述文本，那么描述文本默认为成员的名字)

使用IntegerChoices会带来很多便利：

print(Payment.Statuses.choices)
print(Payment.Statuses.values)
print(Payment.Statuses.names)
print(Payment.Statuses.labels)
# Output:
# [(0, '已创建'), (1, '已提交'), (2, '已审核'), (3, '已支付'), (4, '已驳回')]
# [0, 1, 2, 3, 4]
# ['Created', 'Submitted', 'Reviewed', 'Paid', 'Rejected']
# ['已创建', '已提交', '已审核', '已支付', '已驳回']

print(Payment.Statuses.Paid)
print(Payment.Statuses.Paid.value)
print(Payment.Statuses.Paid.name)
# Output:
# <Statuses.Paid: 3>
# 3
# Paid

# 要获取Payment.Statuses.Paid的描述文本, 直接取label属性就可以了
print(Payment.Statuses.Paid.label)
# Output:
# 已支付

# 真实值可以直接与Payment.Statuses的成员进行比较, 不用再取value属性了
assert 3 == Payment.Statuses.Paid
assert 4 > Payment.Statuses.Paid
# 查询时也不需要
Payment.objects.filter(status=Payment.Statuses.Paid)
# 通过真实值获取描述文本
Payment.Statuses(3).label == "已支付"
# 直接获取某Payment对象的status描述文本
payment = Payment.objects.first()
payment.get_status_display()

IntegerChoices适用于值为整型时的情况，如果值为字符串，可以使用TextChoices。

Pzqqt's Blog

Welcome to Pzqqt's Blog

用好 Django 为你提供的枚举类

用好 Django 为你提供的枚举类