ghetto_enum.py

from __future__ import absolute_import

import six

try:
    import enum
except ImportError:
    enum = None

def make_enum(cls_name, enums):
    # Basic validations.
    if not enums:
        raise ValueError("At least one enum"
                         " name must be provided")
    if not cls_name:
        raise ValueError("A class name must be provided")
    for e in enums:
        if not isinstance(e, six.string_types):
            raise TypeError("Enum name '%s' must"
                            " be a string type" % e)
    # Just use the built-in one (if we can).
    if six.PY3:
        return enum.Enum(cls_name, enums)
    else:
        # Otherwise make a good enough equivalent.
        class EnumValue(object):
            __slots__ = ('name', 'value')
            def __init__(self, name, value):
                self.name = name
                self.value = value
            def __repr__(self):
                return "<%s.%s: %s>" % (cls_name,
                                        self.name,
                                        self.value)
        enums_vals = [EnumValue(e, i + 1)
                      for i, e in enumerate(enums)]
        class EnumMeta(type):
            def __new__(cls, name, parents, dct):
                return super(EnumMeta, cls).__new__(
                    cls, name, parents, dct)
            def __iter__(self):
                for e in enums_vals:
                    yield e
        return EnumMeta(cls_name, (object,), 
                        dict(zip(enums, enums_vals)))


Colors = make_enum('Colors', ['RED', 'GREEN', 'BLUE'])
print(dir(Colors.RED))

print(Colors.RED == Colors.RED)

print(list(Colors))