Skip to content

Instantly share code, notes, and snippets.

@leth
Created May 19, 2012 18:24
Show Gist options
  • Save leth/2731825 to your computer and use it in GitHub Desktop.
Save leth/2731825 to your computer and use it in GitHub Desktop.
This is my enum. There are many like it, but this one is mine.
import inspect
class Instance(object):
"""docstring for Instance"""
def __init__(self, *args, **kwargs):
self.args = args
self.kwargs = kwargs
class ValueInstance(object):
def __init__(self, value):
self.value = value
class EnumMetaclass(type):
def __new__(cls, name, bases, dct):
instances = dict()
for k, v in dct.items():
if isinstance(v, Instance) or isinstance(v, ValueInstance):
instances[k] = v
del dct[k]
new_class = super(EnumMetaclass, cls).__new__(cls, name, bases, dct)
for k, v in instances.items():
if isinstance(v, Instance):
instances[k] = instance = new_class(k, *v.args, **v.kwargs)
elif isinstance(v, ValueInstance):
instances[k] = instance = v.value
else:
continue
setattr(new_class, k, instance)
new_class.__instances__ = instances
return new_class
def __iter__(self):
return self.__instances__.itervalues()
def __call__(cls, *args, **kwargs):
if cls is not Enum:
return super(EnumMetaclass, cls).__call__(*args, **kwargs)
else:
args = set(args)
dct = dict(zip(map(str, args), map(ValueInstance, args)))
for k, v in kwargs.iteritems():
if v in dct.values():
raise ValueError('Duplicate item in enum')
dct[k] = ValueInstance(v)
frame = inspect.currentframe().f_back
dct['__module__'] = inspect.getmodule(frame)
dct['__doc__'] = "Anonymous Enum defined at %s:%s." % (frame.f_code.co_filename, frame.f_lineno)
new_class = EnumMetaclass.__new__(EnumMetaclass, 'AnonEnum', (Enum, ), dct)
new_class.__name__ += str(id(new_class))
return new_class
class Enum(object):
"""docstring for Enum"""
__metaclass__ = EnumMetaclass
def __init__(self, name):
self.name = name
from enum import Enum, Instance
import unittest
class Pets(Enum):
cat = Instance()
dog = Instance()
class EnumMixin(object):
def __init__(self, test_class, instance_names):
self.test_class = test_class
self.instance_names = instance_names
self.instance_map = dict([ (i, getattr(test_class, i)) for i in instance_names])
self.instances = self.instance_map.values()
def test_isinstance(self):
for instance in self.instances:
self.assertIsInstance(instance, self.test_class)
def test_name(self):
for name, instance in self.instance_map.items():
self.assertEqual(instance.name, name)
def test_iterable(self):
colors = iter(self.test_class)
expected = set(self.instances)
count = 0
for i in colors:
count += 1
self.assertIn(i, expected)
self.assertEqual(count, len(expected))
class TestSimpleEnum(unittest.TestCase, EnumMixin):
def __init__(self, *args, **kwargs):
unittest.TestCase.__init__(self, *args, **kwargs)
EnumMixin.__init__(self, Pets, ['cat', 'dog'])
class Color(object):
def __init__(self, r, g, b):
self.r = r
self.g = g
self.b = b
class PrimaryColors(Enum, Color):
red = Instance(255, 0, 0)
green = Instance(0, 255, 0)
blue = Instance(0, 0, 255)
def __init__(self, name, r, g, b):
Enum.__init__(self, name)
Color.__init__(self, r, g, b)
class TestChildClassEnum(unittest.TestCase, EnumMixin):
def __init__(self, *args, **kwargs):
unittest.TestCase.__init__(self, *args, **kwargs)
EnumMixin.__init__(self, PrimaryColors, ['red', 'green', 'blue'])
def test_instance_attrs(self):
keys = ['r', 'g', 'b']
expected = {
'red' : [255, 0, 0],
'green': [0, 255, 0],
'blue' : [0, 0, 255],
}
for name, values in expected.items():
instance = getattr(PrimaryColors, name)
for attr, value in zip(keys, values):
self.assertEqual(getattr(instance, attr), value)
class TestAnonEnum(unittest.TestCase):
def test_create(self):
foo = Enum('one', 'two', 'three')
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment