Created
November 20, 2020 02:34
-
-
Save e2thenegpii/19dd5df64227e27639dac16828098e92 to your computer and use it in GitHub Desktop.
A python enumeration usable as a singledispatch
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
""" | |
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() |
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
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).