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?