Блог

Пытаясь Enum немного лучше понять реализацию Python и разобраться с метапрограммированием, я подумал, что было бы интересно посмотреть, смогу ли я взломать вариант enum с параметризацией его типа по типу его значения. Что-то вроде IntEnum , только больше похоже Enum[int] Enum[str] на , и т.д.

Это оказалось сложнее, чем я думал. Учитывая, как Enum это реализовано, класс не может быть производным от Generic[T] , а также от Enum . Поэтому я попытался включить Enum переменную экземпляра в качестве скрытой переменной экземпляра:

 from enum import Enum
from typing import Generic, TypeVar, Type, Union

T = TypeVar("T")


class TEnumMeta(type):
    def __new__(cls, name, bases, attr_dict):
        instance = super().__new__(cls, name, bases, attr_dict)
        instance._enum = Enum(name, attr_dict)

        for key, attr in attr_dict.items():
            setattr(instance, key, getattr(instance._enum, key))

        return instance


class TEnum(Generic[T], metaclass=TEnumMeta):
    ...
 

К сожалению, добавление типов в приведенный выше код приведет к сбою проверки типа с ошибкой «Тип перечисления как атрибут не поддерживается». Хорошо, не обязательно конец света, но если есть какой-то способ обойти это, я был бы признателен, если бы узнал об этом.

Затем я хотел добавить get() метод, который будет работать с метками или значениями. К сожалению, следующее не прошло проверку типа:

 class TEnum(Generic[T], metaclass=TEnumMeta):
    @classmethod
    def get(cls: "Type[TEnum[T]]", val: Union[T, str, Enum]) -> Enum:
        for key, value in cls._enum.__members__.items():
            if val in [key, value, value.value]:
                return value
        raise ValueError(f"cannot get {val} from {cls}")
 

Это привело к:

metaclass_test.py:21: error: "Type[TEnum[T]]" has no attribute "_enum"

Ага. ХОРОШО. Я полагаю _enum , устанавливается для экземпляра, а не для класса. Следующее обходное решение проверяет тип:

 class TEnum(Generic[T], metaclass=TEnumMeta):
    @classmethod
    def get(cls: "Type[TEnum[T]]", val: Union[T, str, Enum]) -> Enum:
        for nm in dir(cls):
            value = getattr(cls, nm)

            if not isinstance(value, Enum):
                continue

            if val in [nm, value, value.value]:
                return value
    
        raise ValueError(f"cannot get {val} from {cls}")
 

но не работает:

ValueError: cannot get foo from <enum 'TEnum'>

Но если это не удается, то с какой стати это работает?

 class TEnum(Generic[T], metaclass=TEnumMeta):
    ...

def get(cls: Type[TEnum[T]], val: Union[T, str, Enum]) -> Enum:
    for nm in dir(cls):
        value = getattr(cls, nm)

        if not isinstance(value, Enum):
            continue

        if val in [nm, value, value.value]:
            return value

    raise ValueError(f"cannot get {val} from {cls}")


setattr(TEnum, "get", classmethod(get))

class Foo(TEnum[int]):
    foo = 1
    bar = 2
 

Этот тип проверяет и Foo.get("foo") работает. К сожалению, Foo.get("foo") сам по себе не вводит проверку:

metaclass_test.py:46: error: "Type[Foo]" has no attribute "get"

Я не понимаю, почему это работает и почему он не вводит проверку.

Наконец: даже если приведенный выше код сработал, следующий все равно будет вводить проверку:

 class Bar(TEnum[str]):
    baz = 1
    quux = 2
 

Есть ли какой-либо способ сделать эту проверку типа сбоя?