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.

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.