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

Question

Так что я работаю с большой Django кодовой базой, которая использует python Enums везде, например:

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, который является строкой.

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)

Answer 2

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

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

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

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