A python enumeration usable as a singledispatch

enumfun.py

"""
I often find myself parsing buffers that contain typed buffers of data.
Python enumerations are great. I love to use them to map symantic names
onto a collection of values. However they often are lacking in performing
actions based on the enumerant.  Typically I have to do something like
the following

class foo(enum.Enum):
    a = 1
    b = 2
    c = 3

    def do_foo(self):
        if self is foo.a:
            pass # do a foo
        elif self is foo.b:
            pass # do b foo
        elif self is foo.c:
            pass # do c foo
        else:
            raise ValueError() # how did we even get here?

The foo.do_foo function is essentially a dispatch method typically calling
constructors of various types.

Python 3.7 added support for singledispatch methods which are fantastic
for implementing a function like foo.do_foo, however out of the box they
don't work with enums.  This is because type(foo.a) == foo for standard
python enums.

The code below is my kluge/elegant attempt to make python enums work with
the singledispatch decorator.

Python enumerations allow us to define a custom type for the enumerants
by inheriting a class that implements the __new__ method.  This is the case
for enum.IntEnum.  It's definition is basicaly the following:

class IntEnum(int, Enum):
    pass

In order to support singledispatch each enumerant must have a unique type.

Taking a page from IntEnum I created a class UniqueEnumType that dynamically
generates a type in its __new__ method based on the enumerant value and
class name so the type of each enumerant is unique.

Most of the implementation of UniqueEnumType as found below is related to
making our enumerants look and feel like standard python enumerants except
for the fact that foo.enumerant.__class__ is unique to the enumerant.

We use a python metaclass in UniqueEnumType so that the types will
compare and hash equal.

"""
import enum
from types import DynamicClassAttribute
from functools import singledispatch

class UniqueEnumMeta(type):
    """
    Metaclass that makes all produced classes compare equal and hash equal
    if their respective names match.
    """
    def __eq__(self, other):
        return self.__name__ == other.__name__

    def __hash__(self):
        return hash(self.__name__)

class UniqueEnumType:

    @classmethod
    def create_type(cls, name, value, **kwargs):
        """
        Produce a new type whose name is a concatenation of the name and value
        Unfortunately because we're using metaclasses pickling these enumerants
        becomes impossible.
        """
        cls_name = "%s_%r" % (name, value)

        bases = (UniqueEnumType,)
        attributes = {
            **kwargs,
            "_orig_class_name": name,
            "__module__": None, # makes printing nicer
        }

        return UniqueEnumMeta(cls_name, bases, attributes)

    def __new__(cls, *args, **kwargs):
        """
        Dynamically create an enumerant specific class
        """
        if not issubclass(cls, enum.Enum):
            cls._value_ = args[0]
            return object.__new__(cls)
        custom_type = UniqueEnumType.create_type(cls.__name__, args[0], **cls.__dict__)
        return custom_type.__new__(custom_type, *args, **kwargs)

    @DynamicClassAttribute
    def name(self):
        """
        Give ourselves a name attribute like normal python enumerants
        """
        return self._name_

    @DynamicClassAttribute
    def value(self):
        """
        Give ourselves a value attribute like normal python enumerants
        """
        return self._value_

    def __repr__(self):
        """
        Make our repr just like a normal python enumerant
        """
        return "<%s.%s: %r>" % (self._orig_class_name, self.name, self.value)


class foo(UniqueEnumType, enum.Enum):
    """
    Example enumeration

    Notice type(foo.a) == <class 'foo_1'> instead of <class 'foo'>
    """
    a = 1
    b = 2
    c = 3

    @classmethod
    def parse(cls, view):
        """
        Normal enum class method parses a new type and calls the registerd
        enumerant handler
        """
        return cls(view[0]).parse_type(view[1:])

    @singledispatch
    def parse_type(self, view):
        """
        Default value handler will be called if no enumerant handler is registered
        """
        raise NotImplementedError()

    @parse_type.register(UniqueEnumType.create_type("foo", 2))
    def _(self, view):
        """
        Parser function for parsing foo.b
        The syntax for registering a enum method is kind of a kludge.
        We essentially create an enum that hashes the same type
        that will be produced for the enumerant.

        Unfortunately calling @parse_type.register(b) wouldn't work
        because at this point b == int(2) and hasn't been mangled
        by the enum.EnumMeta metaclass.
        """
        print("Parsing type related to b")
        return two(), view

    @parse_type.register(UniqueEnumType.create_type("foo", 3))
    def _(self, view):
        """
        Same as above different type
        """
        print("parsing type related to c")
        return three(), view


class one:
    @foo.parse_type.register(type(foo.a))
    def _(enumerant, view):
        """
        Parsing methods don't have to be members of the enumeration class

        Notice the registration syntax changes a little bit.

        The registration syntax for in enum functions is a little cleaner
        type(foo.a) returns the type specific to foo.a

        The parameters are a little querky in that the first parameter is the
        enumerant that is responsable for calling the dispatcher
        """
        print(enumerant)
        print(view)
        print("Parsing type related to a")
        return one(), view

class two:
    pass

class three:
    pass

def main():
    view = memoryview(b"\x01\x02\x03")

    while view:
        val, view = foo.parse(view)
        print(val)

if __name__ == "__main__":
    main()

Licensed CC BY-SA
https://creativecommons.org/licenses/by-sa/4.0/

Hi, nice use of the singledispatch decorator. Someone at StackOverflow shared your gist in a comment to a question that's just about identical to the problem you're trying to solve here, and I thought you'd be interested in checking out my solution using a decorator class instead: https://stackoverflow.com/a/73105260/6890912

When I wrote this I mostly wanted to play around with the singledispatch decorator, and I feel like a custom __new__ method on enums is a vastly underused feature.

I definitely like that your solution doesn't involve metaclasses and dynamic typing, but it doesn't register a default (easily rectified by using a defaultdict in the bind decorator though adding complexity). Also when you forget to register a callable for a type you'll get a key error which I feel like is a little bit confusing as I'd expect a NotImplemented exception would be better (also easily fixed).