python - 强制正确使用 python Enum

Question

所以我正在使用一个大型 Django 代码库，该代码库自始至终都使用 python 枚举，例如:

from enum import Enum

class Status(Enum):
    active = 'active'

# ... later
assert some_django_model_instance.status == Status.active.value  # so far so good

...但是“.value”部分当然总是被遗忘并被遗漏。现在很难完全放弃枚举，尽管它们的问题多于有用性。有没有办法自动检查像这样的行:

assert some_django_model_instance.status == Status.active  # someone forgot ".value" here!

使用 mypy 或 pylint 或者可能向基本枚举添加一些代码/断言？问题是，Status.active并没有真正调用任何代码，它只是返回一个类，当然该类永远不等于some_django_model_instance.status，这是一个字符串。

Answer 1

您可以通过子类化 enum.EnumMeta来强制执行此操作:

from enum import EnumMeta, Enum as _Enum

class Enum(_Enum, metaclass=EnumMeta):
    def __eq__(self, arg):
        if isinstance(arg, self.__class__):
            return arg is self
        return self.value == arg

现在您不必调用 enum.value 进行比较:

class Method(Enum):
    GET = 'GET'
    POST = 'POST'

>>> get = 'GET'
>>> Method.GET == get
True
>>> get == Method.GET
True
>>> Method.POST == Method.GET
False

这解决了其他人不会忘记调用 .value 进行比较的问题，但会产生一个更大的问题，因为现在人们忘记调用 的可能性成倍增加>.value 插入模型时。

要解决此问题，我建议还对 models.CharField 进行子类化 创建您自己的枚举字段:

class EnumField(models.CharField):
    def __init__(self, enum, **kwargs):
        self.enum = enum

    def from_db_value(self, value, expression, connection):
        if value is not None:
            return self.enum(value)
        return None

    def to_python(self, value):
        if isinstance(value, self.enum):
            return value.value
        return None

    def get_prep_value(self, value):
        if isinstance(value, self.enum):
            value = value.value
        return super().get_prep_value(value)

    def deconstruct(self):
        name, path, args, kwargs = super().deconstruct()
        args.append(self.enum)
        return name, path, args, kwargs

现在您还可以在不调用 .value 的情况下插入模型:

class MyModel(models.Model):
    method = EnumField(enum=Method)

>>> MyModel.objects.create(method=Method.GET)