Add a mixin to python enum after it's created?

Question

The Problem

Suppose I have an enum defined:

from enum import Enum, auto

class Foo(Enum):
   DAVE_GROHL = auto()
   MR_T = auto()

and I want to extend this class with specific methods only (ie, no new enum members):

class MixinFoo(Foo):
   def catch_phrase(self):
      if self is Foo.MR_T:
         return "I pity da foo!"
      else:
         return "There goes my hero!"

This will fail with TypeError, since by design enum's can't be extended once members have been defined. Python enum allows you to do this in the other direction (ie: define mixin with no members and subclass the mixin to add members), but in my specific use case, defining MixinFoo before Foo to use as a base class is not an option (eg: Foo is defined in another module and I don't want to recreate it or modify that module's code).

What I've considered so far:

Tweaking the EnumMeta. EnumMeta checks for existing members in the base class in the __prepare__ method, so overriding this method and deferring the check in a new metaclass could work. Unfortunately, EnumMeta checks for existing members in several other functions, and redefining all those functions seems like bad practice in general (lots of duplicated code).

Adding functions after the fact. I know it's possible to do something like:

def catch_phrase(self):
   ...

Foo.catch_phrase = catch_phrase

and while this might work, there are some significant disadvantages that I would like to avoid (eg: very cumbersome for a large number of functions, static/class methods, properties of normal class definition might be difficult to implement)

Solving the problem with composition over inheritance:

class FooWrapper():
   def __init__(self, foo):
      self._foo = foo

   def catch_phrase(self):
      ...

While this is possible, I'm not a huge fan of this method.

Some Related Questions:

Why does EnumMeta have checks for existing members outside of __prepare__? From a design perspective, this seems to not only be redundant, but also makes it unnecessarily difficult to extend the EnumMeta as in my case. I understand the rationale behind not allowing extension of enums with more members, but mixins were obviously considered in the design of the enum library. Was this an oversight in the design of the enum? Was it considered and rejected? Is it intentional for some reason?

Furthermore, one of the checks for existing members is buried in the _EnumDict class, which is not accessible, so I'm not sure what I'm trying to do is possible without recreating a fairly substantial portion of the enum library.

The overall intent is that I have a common enum where the members are known and fixed, but the use case and methods are application specific (ie, different applications that use Foo would have different greetings and possibly other functions to add, but no application would need to add other Foos).

Answer 1

Option 1: assign new methods to existing Enum

Does simply assigning a method to Foo works?

Like in Foo.catch = lambda self=None: print(Foo if self is None else "MR_T" if self is Foo.MR_T else "Ka-bong")

If not, why? It does behave, for all purposes as a method would, and "self" is filled with the "Foo" member if the method is called from the member.

If you want to benefit from Python's mechanism for creating classes and the "class" statement block syntax, this is easily feasible with a "pseudo" metaclass: a callable that will take in the contents of the class body, along with the name and bases - it can them just add the newly created methods and attributes to the existing Enum class:

def EnumMixinExtender(name, bases, ns):
    original = bases[0]
    for k, v in ns.items():
        setattr(original, k, v)
    return original

This will allow you to extend Foo in place with

class MixinFoo(Foo, metaclass=EnumMixinExtender):
   def catch_phrase(self=None):
        ...

With this option, MixinFoo is Foo and MixinFoo.member1 is Foo.member1: the Foo class is monkey patched in place - no need to care about import order: everyone using "Foo" can use "catch_phrase".

Option 2 - MixinFoo != Foo:

If one must preserve a distinct Foo class and need that Foo.member1 is not MixinFoo.member1 asserts True - the way out, still using a pseudo-metaclass is to invert things, and create a copy of Foo that would derive from the new class-members defined in the Mixin. The metaclass code just does that in the reversed order, so that the Enum class comes later into the mix.

def ReverseMixinAdder(name, bases, ns):
    pre_bases = tuple(base for base in bases if not isinstance(base, enum.EnumMeta))
    pos_bases = tuple(base for base in bases if isinstance(base, enum.EnumMeta))

    new_mixin = type(f"{name}_mixin", pre_bases, ns)
    member_space = {}
    for base in pos_bases:
        member_space |= {k: v.value for k, v in base.__members__.items()}
    new_enum = Enum(name, member_space, type=new_mixin )
    return new_enum

Just use that as the metaclas, like above, and you get an independent "MixinFoo" class, that can be passed around and used independent from Foo. The drawback is that not only its members "are not" Foo members, just copies, but MixinFoo won't be a subclass or otherwise related to Foo at all.

(btw, just ignore the mechanisms allowing one to merge more than one Enum if you won't use them. I had not tested that part either)

Option 3 - MixinFoo != Foo, but issubclass(MixinFoo, Foo):

If one must have MixinFoo as a distinct class from Foo and still have MixinFoo.member1 is Foo.member1 and issubclass(MixinFoo, Foo) to assert True,

Well... looking into the traceback when one tries to extend an Enum with mixin methods we have:

    565                 if issubclass(base, Enum) and base._member_names_:
--> 566                     raise TypeError(

Sure this is an implementation detail - but all EnumMeta checks at this point is... member_names - which happens to be assignable.

If one simply assigns a falsey value to "Foo.member_names", it can be subclassed, with extra mixin methods and ordinary attributes at will. The ._member_names_ can be restored afterwards.

Doing this, these assertions hold: FooMixin.member1 is Foo.member1, FooMixin is not Foo, issubclass(FooMixin, Foo). However, the new methods can not be called directly on the members - that is, in the previous two options one can do FooMixin.member1.catch_phrase(), while under this, as the member enum keeps the class Foo which does not have the new method, that does not work. (The workaround is FooMixin.catch_phrase(FooMixin.member1) .

If one wants this last part to work, Foo.members can have their __class__ attribute updated to FooMixin - and it will work, but then the original Foo.members are updated inplace as well.

I have the intuition this final form is what you are really asking for here -

class MetaMixinEnum(enum.EnumMeta):
    registry = {}

    @classmethod
    def _get_enum(mcls, bases):
        enum_index, member_names = next((i, base._member_names_) for i, base in enumerate(bases) if issubclass(base, Enum))
        return bases[enum_index], member_names

    @classmethod
    def __prepare__(mcls, name, bases):
        base_enum, member_names = mcls._get_enum(bases)
        mcls.registry[base_enum] = member_names
        base_enum._member_names_ = []
        return super().__prepare__(name, bases)

    def __new__(mcls, name, bases, ns):
        base_enum, _  = mcls._get_enum(bases)
        try:
            new_cls = super().__new__(mcls, name, bases, ns)
        finally:
            member_names = base_enum._member_names_ = mcls.registry[base_enum]
        new_cls._member_names_ = member_names[:]
        for name in member_names:
            setattr(getattr(new_cls, name), "__class__", new_cls)
        return new_cls


class FooMixin(Foo, metaclass=MetaMixinEnum):
    def catch_phrase(self):
        ...

Answer 2

Option 4 - issubclass(Foo, MixinFoo)

Building on @jsbueno's with inspiration from an answer from @l4mpi in this question. Here we monkeypatch Foo such that it is a subclass of MixinFoo, as would be the case for a typical Enum mixin. We simply add MixinFoo to Foo's bases:

class Foo(Enum):
   ...

class MixinFoo(Enum):
   def catch_phrase(self):
      ...

Foo.__bases__ = (MixinFoo, ) + Foo.__bases__

for foo in Foo:
   foo.catch_phrase()

Answer 3

Why does EnumMeta have checks for existing members outside of __prepare__?

Good question. It no longer does (as of 3.11).

Furthermore, one of the checks for existing members is buried in the _EnumDict class

No, the check in _EnumDict is for members already defined in that enum:

class Foo(Enum):
    BAR = 1
    BAR = 1 # this errors out because of _EnumDict

---

I like your Option 4 -- but of course, it should be a decorator:

def add_methods_to(cls):
    def insert_class(scls):
        cls.__bases__ = (scls, ) + cls.__bases__
        return scls
    return insert_class

and in action:

@add_methods_to(Foo)
class MixinFoo:             # mix-in class does not need to be an enum
    def catch_phrase(self):
        return 'blah blah'