Enum union in Python

Question

This is a snippet of code written in Python 3.9 to implement the concept of "enum unions": an enum made up of several sub-enums whose set of members is the union of the sub-enums' members. This aims to be fully compatible with the standard library enum. Up-to-date code with some documentation can be found here.

Usage and desired behaviour:

import enum
class EnumA(enum.Enum):
 A = 1
class EnumB(enum.Enum):
 B = 2
 ALIAS = 1
class EnumC(enum.Enum):
 C = 3
 

>>> UnionAB = enum_union(EnumA, EnumB)
>>> UnionAB.__members__
mappingproxy({'A': <EnumA.A: 1>, 'B': <EnumB.B: 2>, 'ALIAS': <EnumA.A: 1>})
>>> list(UnionAB)
[<EnumA.A: 1>, <EnumB.B: 2>]
>>> EnumA.A in UnionAB
True
>>> EnumB.ALIAS in UnionAB
True
>>> isinstance(EnumB.B, UnionAB)
True
>>> issubclass(UnionAB, enum.Enum)
True
>>> UnionABC = enum_union(UnionAB, EnumC)
>>> UnionABC.__members__
mappingproxy({'A': <EnumA.A: 1>, 'B': <EnumB.B: 2>, 'ALIAS': <EnumA.A: 1>, 'C': <EnumC.C: 3>})
>>> set(UnionAB).issubset(UnionABC)
True

The code is below. I'm mainly concerned with design, maintainability, compatibility with enum and intuitive (unsurprising) behaviour for the user. The implementation uses the internals of enum.EnumMeta; I'm aware this might be brittle or might not be reliable in the long run but it's the only way of reusing most of enum.EnumMeta's code without rewriting the whole thing.

import enum
import itertools as itt
from functools import reduce
import operator
from typing import Literal, Union
import more_itertools as mitt
AUTO = object()
class UnionEnumMeta(enum.EnumMeta):
 """
 The metaclass for enums which are the union of several sub-enums.
 Union enums have the _subenums_ attribute which is a tuple of the enums forming the
 union.
 """
 @classmethod
 def make_union(
 mcs, *subenums: enum.EnumMeta, name: Union[str, Literal[AUTO], None] = AUTO
 ) -> enum.EnumMeta:
 """
 Create an enum whose set of members is the union of members of several enums.
 Order matters: where two members in the union have the same value, they will
 be considered as aliases of each other, and the one appearing in the first
 enum in the sequence will be used as the canonical members (the aliases will
 be associated to this enum member).
 :param subenums: Sequence of sub-enums to make a union of.
 :param name: Name to use for the enum class. AUTO will result in a combination
 of the names of all subenums, None will result in "UnionEnum".
 :return: An enum class which is the union of the given subenums.
 """
 subenums = mcs._normalize_subenums(subenums)
 class UnionEnum(enum.Enum, metaclass=mcs):
 pass
 union_enum = UnionEnum
 union_enum._subenums_ = subenums
 if duplicate_names := reduce(
 set.intersection, (set(subenum.__members__) for subenum in subenums)
 ):
 raise ValueError(
 f"Found duplicate member names in enum union: {duplicate_names}"
 )
 # If aliases are defined, the canonical member will be the one that appears
 # first in the sequence of subenums.
 # dict union keeps last key so we have to do it in reverse:
 union_enum._value2member_map_ = value2member_map = reduce(
 operator.or_, (subenum._value2member_map_ for subenum in reversed(subenums))
 )
 # union of the _member_map_'s but using the canonical member always:
 union_enum._member_map_ = member_map = {
 name: value2member_map[member.value]
 for name, member in itt.chain.from_iterable(
 subenum._member_map_.items() for subenum in subenums
 )
 }
 # only include canonical aliases in _member_names_
 union_enum._member_names_ = list(
 mitt.unique_everseen(
 itt.chain.from_iterable(subenum._member_names_ for subenum in subenums),
 key=member_map.__getitem__,
 )
 )
 if name is AUTO:
 name = (
 "".join(subenum.__name__.removesuffix("Enum") for subenum in subenums)
 + "UnionEnum"
 )
 UnionEnum.__name__ = name
 elif name is not None:
 UnionEnum.__name__ = name
 return union_enum
 def __repr__(cls):
 return f"<union of {', '.join(map(str, cls._subenums_))}>"
 def __instancecheck__(cls, instance):
 return any(isinstance(instance, subenum) for subenum in cls._subenums_)
 @classmethod
 def _normalize_subenums(mcs, subenums):
 """Remove duplicate subenums and flatten nested unions"""
 # we will need to collapse at most one level of nesting, with the inductive
 # hypothesis that any previous unions are already flat
 subenums = mitt.collapse(
 (e._subenums_ if isinstance(e, mcs) else e for e in subenums),
 base_type=enum.EnumMeta,
 )
 subenums = mitt.unique_everseen(subenums)
 return tuple(subenums)
def enum_union(*enums, **kwargs):
 return UnionEnumMeta.make_union(*enums, **kwargs)

Answer 1

For the benefit of anyone else that, like me, does not read Typing nor Black, here is the code reformatted. Comments are interspersed.

AUTO = object()
class UnionEnumMeta(enum.EnumMeta):
 """
 The metaclass for enums which are the union of several sub-enums.
 Union enums have the _subenums_ attribute which is a tuple of the enums
 forming the union.
 """
 @classmethod
 def make_union(mcs, *subenums, name=AUTO):
 """
 Create an enum whose set of members is the union of members of
 several enums.
 Order matters: where two members in the union have the same value,
 they will be considered as aliases of each other, and the one appearing
 in the first enum in the sequence will be used as the canonical members
 (the aliases will be associated to this enum member).
 :param subenums: Sequence of sub-enums to make a union of.
 :param name: Name to use for the enum class. AUTO will result in a
 combination of the names of all subenums, None will result
 in "UnionEnum".
 :return: An enum class which is the union of the given subenums.
 """
 subenums = mcs._normalize_subenums(subenums)
 class UnionEnum(enum.Enum, metaclass=mcs):
 pass
 union_enum = UnionEnum

Did you mean to instantiate UnionEnum here? Looks like you forgot the parenthesis.

 union_enum._subenums_ = subenums
 if duplicate_names := reduce(
 set.intersection,
 (set(subenum.__members__) for subenum in subenums),
 ):
 raise ValueError(
 f"Found duplicate member names in enum union: {duplicate_names}"
 )
 # If aliases are defined, the canonical member will be the one that
 # appears first in the sequence of subenums.
 # dict union keeps last key so we have to do it in reverse:
 union_enum._value2member_map_ = value2member_map = reduce(
 operator.or_,
 (subenum._value2member_map_ for subenum in reversed(subenums)),
 )
 # union of the _member_map_'s but using the canonical member always:
 union_enum._member_map_ = member_map = {
 name: value2member_map[member.value]
 for name, member in itt.chain.from_iterable(
 subenum._member_map_.items() for subenum in subenums
 )}
 # only include canonical aliases in _member_names_
 union_enum._member_names_ = list(
 mitt.unique_everseen(
 itt.chain.from_iterable(
 subenum._member_names_ for subenum in subenums
 ),
 key=member_map.__getitem__,
 ))
 if name is AUTO:
 name = (
 "".join(
 subenum.__name__.removesuffix("Enum")
 for subenum in subenums
 )
 + "UnionEnum"
 )
 UnionEnum.__name__ = name
 elif name is not None:
 UnionEnum.__name__ = name

What happens if name is None?

 return union_enum
 def __repr__(cls):
 return f"<union of {', '.join(map(str, cls._subenums_))}>"
 def __instancecheck__(cls, instance):
 return any(
 isinstance(instance, subenum)
 for subenum in cls._subenums_
 )
 @classmethod
 def _normalize_subenums(mcs, subenums):
 """Remove duplicate subenums and flatten nested unions"""
 # we will need to collapse at most one level of nesting, with the
 # inductive hypothesis that any previous unions are already flat
 subenums = mitt.collapse(
 (e._subenums_ if isinstance(e, mcs) else e for e in subenums),
 base_type=enum.EnumMeta,
 )
 subenums = mitt.unique_everseen(subenums)
 return tuple(subenums)
def enum_union(*enums, **kwargs):
 return UnionEnumMeta.make_union(*enums, **kwargs)

All in all, this looks very interesting. While the pieces you are referencing from EnumMeta are implementation details, they are unlikely to change. I am curious how UnionEnum would be used to solve the original Stackoverflow question?

---

Disclosure: I am the author of the Python stdlib Enum, the enum34 backport, and the Advanced Enumeration (aenum) library.

• Anakhand: Thanks for the review! "Did you mean to instantiate UnionEnum here?" Not really; enums are classes, so what I want to return is the class itself (which is dynamically created every time that method is run); one can think of the UnionEnum class definition as the instantiation of UnionEnumMeta (or perhaps "meta-instantiation"). Then again, you could ask me why make the union_enum alias, and the answer would just be "It felt weird setting attributes of a dynamically created class and so I used a lowercase alias to mentally hide that" :) [continued]
• Anakhand: [continuation] Regarding the name=None case, that just leaves the default UnionEnum name. Finally, you are right that I didn't really specify how this can be used in the original SO question; I've added an edit there. The idea is computing the union of the "base" enum and the "extension" enum to achieve exactly the desired behaviour (and more, since unions support additional non-alias members too).

Answer 2

I have made some observations and improvements myself, and I thought it could be useful to add them here as well.

• What happens with the empty union? Currently, the reduce calls will raise because they will receive an empty iterator.

  This can be fixed by simply adding the appropriate identity element to each reduce operation (as the initial argument), which will serve as the default value if the iterable is empty.

  Also, a minor aesthetic detail, an empty union would have a repr of <union of >.

• The check for duplicates is implemented as checking that the intersection of all names is empty:

  if duplicate_names := reduce(
   set.intersection, (set(subenum.__members__) for subenum in subenums)
   ):
   raise ValueError(...)
  

  But this is not quite what we want to check: we don't want to check whether some names appear in all subenums; if a name appears in more than one subenum it's enough to mark it as a duplicate. This is most easily solved with a good old-fashioned for loop:

  names, duplicates = set(), set()
  for subenum in subenums:
   for name in subenum.__members__:
   (duplicates if name in names else names).add(name)
  if duplicates:
   raise ValueError(f"Found duplicate member names: {duplicates}")
  

  I guess I was too eager to make this look "mathematical" with reduce and set intersection.

• The merging of the _value2member_map_s is done via PEP 584 dict union:

   union_enum._value2member_map_ = value2member_map = reduce(
   operator.or_, (subenum._value2member_map_ for subenum in reversed(subenums))
   )
  

  Something to keep in mind is that this makes a copy for every intermediate operand, and so this will be inefficient for moderate to large numbers of subenums. If, however, this is intended for just a small size of subenums that are specified manually in the varargs list, it should be fine. An advantage of using the union operator is that it is a pure function and hence there's no risk of accidentally modifying one of the subenums.

• Since union enums are not really a type in themselves (semantically speaking), they just aggregate a bunch of other enums, it might make sense to implement equality in terms of subenums:

  def __eq__(cls, other):
   """Equality based on the tuple of subenums (order-sensitive)."""
   if not isinstance(other, UnionEnumMeta):
   return NotImplemented
   return cls._subenums_ == other._subenums_
  

  Using the example above, this would result in

  >>> enum_union(UnionAB, EnumC) == enum_union(EnumA, EnumB, EnumC)
  True
  

  which makes sense intuitively.

• hjpotter92: You can also declare the function as an alias: enum_union = UnionEnumMeta.make_union without need to push another function call in the stack