Skip to content

Instantly share code, notes, and snippets.

@e2thenegpii
Created November 20, 2020 02:34
Show Gist options
  • Save e2thenegpii/19dd5df64227e27639dac16828098e92 to your computer and use it in GitHub Desktop.
Save e2thenegpii/19dd5df64227e27639dac16828098e92 to your computer and use it in GitHub Desktop.
A python enumeration usable as a singledispatch
"""
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()
@e2thenegpii
Copy link
Author

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).

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment