Обеспечение правильного использования Python Enum

Поэтому я работаю с большой кодовой базой 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" постоянно забывается и исчезает. К настоящему времени было бы трудно полностью отказаться от Enums, хотя они были скорее проблематичными, чем полезными. Есть ли способ автоматической проверки строк, подобных этим:

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

скажем, с Mypy или Pylint или, возможно, добавив некоторый код / утверждений в базовый Enum? Проблема в том, что Status.active самом деле не вызывает никакого кода, он просто возвращает класс, и, конечно, этот класс никогда не равен some_django_model_instance.status , который является строкой.

python django pylint mypy

Всего 2 ответа

Вы можете применить это, enum.EnumMeta подкласс 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 для создания собственного поля enum:

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

Вы можете заставить mypy обнаруживать эти типы проблемных сравнений, используя --strict-equality флага / config командной строки --strict-equality . Если этот флаг включен, выполнение some_str == Status.active приведет к ошибке, подобной следующей:

error: Non-overlapping equality check (left operand type: "str", right operand type: "Literal[Status.active]")

Примечание: этот флаг будет проверять все сравнения ложных равенств, а не только перечисления.

Однако он будет отключен для любых сравнений на равенство, где любой из операндов определил пользовательский метод __eq__ , поскольку пользовательский метод действительно может делать что угодно.