| """Enumeration metaclass.""" | |
| class EnumMetaclass(type): | |
| """Metaclass for enumeration. | |
| To define your own enumeration, do something like | |
| class Color(Enum): | |
| red = 1 | |
| green = 2 | |
| blue = 3 | |
| Now, Color.red, Color.green and Color.blue behave totally | |
| different: they are enumerated values, not integers. | |
| Enumerations cannot be instantiated; however they can be | |
| subclassed. | |
| """ | |
| def __init__(cls, name, bases, dict): | |
| super(EnumMetaclass, cls).__init__(name, bases, dict) | |
| cls._members = [] | |
| for attr in dict.keys(): | |
| if not (attr.startswith('__') and attr.endswith('__')): | |
| enumval = EnumInstance(name, attr, dict[attr]) | |
| setattr(cls, attr, enumval) | |
| cls._members.append(attr) | |
| def __getattr__(cls, name): | |
| if name == "__members__": | |
| return cls._members | |
| raise AttributeError, name | |
| def __repr__(cls): | |
| s1 = s2 = "" | |
| enumbases = [base.__name__ for base in cls.__bases__ | |
| if isinstance(base, EnumMetaclass) and not base is Enum] | |
| if enumbases: | |
| s1 = "(%s)" % ", ".join(enumbases) | |
| enumvalues = ["%s: %d" % (val, getattr(cls, val)) | |
| for val in cls._members] | |
| if enumvalues: | |
| s2 = ": {%s}" % ", ".join(enumvalues) | |
| return "%s%s%s" % (cls.__name__, s1, s2) | |
| class FullEnumMetaclass(EnumMetaclass): | |
| """Metaclass for full enumerations. | |
| A full enumeration displays all the values defined in base classes. | |
| """ | |
| def __init__(cls, name, bases, dict): | |
| super(FullEnumMetaclass, cls).__init__(name, bases, dict) | |
| for obj in cls.__mro__: | |
| if isinstance(obj, EnumMetaclass): | |
| for attr in obj._members: | |
| # XXX inefficient | |
| if not attr in cls._members: | |
| cls._members.append(attr) | |
| class EnumInstance(int): | |
| """Class to represent an enumeration value. | |
| EnumInstance('Color', 'red', 12) prints as 'Color.red' and behaves | |
| like the integer 12 when compared, but doesn't support arithmetic. | |
| XXX Should it record the actual enumeration rather than just its | |
| name? | |
| """ | |
| def __new__(cls, classname, enumname, value): | |
| return int.__new__(cls, value) | |
| def __init__(self, classname, enumname, value): | |
| self.__classname = classname | |
| self.__enumname = enumname | |
| def __repr__(self): | |
| return "EnumInstance(%s, %s, %d)" % (self.__classname, self.__enumname, | |
| self) | |
| def __str__(self): | |
| return "%s.%s" % (self.__classname, self.__enumname) | |
| class Enum: | |
| __metaclass__ = EnumMetaclass | |
| class FullEnum: | |
| __metaclass__ = FullEnumMetaclass | |
| def _test(): | |
| class Color(Enum): | |
| red = 1 | |
| green = 2 | |
| blue = 3 | |
| print Color.red | |
| print repr(Color.red) | |
| print Color.red == Color.red | |
| print Color.red == Color.blue | |
| print Color.red == 1 | |
| print Color.red == 2 | |
| class ExtendedColor(Color): | |
| white = 0 | |
| orange = 4 | |
| yellow = 5 | |
| purple = 6 | |
| black = 7 | |
| print ExtendedColor.orange | |
| print ExtendedColor.red | |
| print Color.red == ExtendedColor.red | |
| class OtherColor(Enum): | |
| white = 4 | |
| blue = 5 | |
| class MergedColor(Color, OtherColor): | |
| pass | |
| print MergedColor.red | |
| print MergedColor.white | |
| print Color | |
| print ExtendedColor | |
| print OtherColor | |
| print MergedColor | |
| def _test2(): | |
| class Color(FullEnum): | |
| red = 1 | |
| green = 2 | |
| blue = 3 | |
| print Color.red | |
| print repr(Color.red) | |
| print Color.red == Color.red | |
| print Color.red == Color.blue | |
| print Color.red == 1 | |
| print Color.red == 2 | |
| class ExtendedColor(Color): | |
| white = 0 | |
| orange = 4 | |
| yellow = 5 | |
| purple = 6 | |
| black = 7 | |
| print ExtendedColor.orange | |
| print ExtendedColor.red | |
| print Color.red == ExtendedColor.red | |
| class OtherColor(FullEnum): | |
| white = 4 | |
| blue = 5 | |
| class MergedColor(Color, OtherColor): | |
| pass | |
| print MergedColor.red | |
| print MergedColor.white | |
| print Color | |
| print ExtendedColor | |
| print OtherColor | |
| print MergedColor | |
| if __name__ == '__main__': | |
| _test() | |
| _test2() |