typed Enum implementation that uses Enum; either doesn't work or doesn't type check

Viewed 248

Trying to understand Python's Enum implementation a bit better, and come to grips with meta-programming, I thought it might be fun to see if I could hack up an enum variant with its type parametrized on the type of its value. Something like IntEnum, only more like Enum[int], Enum[str], etc.

This has proven more difficult than I thought. Given how Enum is implemented there's no way for a class to derive from Generic[T] as well as from Enum. So I tried containing an Enum as a behind-the-scenes instance variable:

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

Unfortunately, adding types to the above code will fail type checking with an "Enum type as attribute is not supported" error. OK, not the end the world, necessarily, but if there's some way around this I'd appreciate learning about it.

Then I wanted to add a get() method that would work on labels or values. Unfortunately, the following didn't pass its type check:

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}")

It resulted in:

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

Huh. OK. I suppose _enum is set on the instance, not the class. The following workaround type checks:

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}")

but doesn't work:

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

But if that fails, then why on earth does this work?

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

This type checks, and Foo.get("foo") works. Unfortunately, Foo.get("foo") itself doesn't type check:

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

I'm confused both why it works and why it doesn't type check.

Finally: even if the above code worked, the following would still type check:

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

Is there any way to make that fail type checking?

0 Answers
Related