1"""Enumeration metaclass.""" 2 3class EnumMetaclass(type): 4 """Metaclass for enumeration. 5 6 To define your own enumeration, do something like 7 8 class Color(Enum): 9 red = 1 10 green = 2 11 blue = 3 12 13 Now, Color.red, Color.green and Color.blue behave totally 14 different: they are enumerated values, not integers. 15 16 Enumerations cannot be instantiated; however they can be 17 subclassed. 18 """ 19 20 def __init__(cls, name, bases, dict): 21 super(EnumMetaclass, cls).__init__(name, bases, dict) 22 cls._members = [] 23 for attr in dict.keys(): 24 if not (attr.startswith('__') and attr.endswith('__')): 25 enumval = EnumInstance(name, attr, dict[attr]) 26 setattr(cls, attr, enumval) 27 cls._members.append(attr) 28 29 def __getattr__(cls, name): 30 if name == "__members__": 31 return cls._members 32 raise AttributeError, name 33 34 def __repr__(cls): 35 s1 = s2 = "" 36 enumbases = [base.__name__ for base in cls.__bases__ 37 if isinstance(base, EnumMetaclass) and not base is Enum] 38 if enumbases: 39 s1 = "(%s)" % ", ".join(enumbases) 40 enumvalues = ["%s: %d" % (val, getattr(cls, val)) 41 for val in cls._members] 42 if enumvalues: 43 s2 = ": {%s}" % ", ".join(enumvalues) 44 return "%s%s%s" % (cls.__name__, s1, s2) 45 46class FullEnumMetaclass(EnumMetaclass): 47 """Metaclass for full enumerations. 48 49 A full enumeration displays all the values defined in base classes. 50 """ 51 52 def __init__(cls, name, bases, dict): 53 super(FullEnumMetaclass, cls).__init__(name, bases, dict) 54 for obj in cls.__mro__: 55 if isinstance(obj, EnumMetaclass): 56 for attr in obj._members: 57 # XXX inefficient 58 if not attr in cls._members: 59 cls._members.append(attr) 60 61class EnumInstance(int): 62 """Class to represent an enumeration value. 63 64 EnumInstance('Color', 'red', 12) prints as 'Color.red' and behaves 65 like the integer 12 when compared, but doesn't support arithmetic. 66 67 XXX Should it record the actual enumeration rather than just its 68 name? 69 """ 70 71 def __new__(cls, classname, enumname, value): 72 return int.__new__(cls, value) 73 74 def __init__(self, classname, enumname, value): 75 self.__classname = classname 76 self.__enumname = enumname 77 78 def __repr__(self): 79 return "EnumInstance(%s, %s, %d)" % (self.__classname, self.__enumname, 80 self) 81 82 def __str__(self): 83 return "%s.%s" % (self.__classname, self.__enumname) 84 85class Enum: 86 __metaclass__ = EnumMetaclass 87 88class FullEnum: 89 __metaclass__ = FullEnumMetaclass 90 91def _test(): 92 93 class Color(Enum): 94 red = 1 95 green = 2 96 blue = 3 97 98 print Color.red 99 100 print repr(Color.red) 101 print Color.red == Color.red 102 print Color.red == Color.blue 103 print Color.red == 1 104 print Color.red == 2 105 106 class ExtendedColor(Color): 107 white = 0 108 orange = 4 109 yellow = 5 110 purple = 6 111 black = 7 112 113 print ExtendedColor.orange 114 print ExtendedColor.red 115 116 print Color.red == ExtendedColor.red 117 118 class OtherColor(Enum): 119 white = 4 120 blue = 5 121 122 class MergedColor(Color, OtherColor): 123 pass 124 125 print MergedColor.red 126 print MergedColor.white 127 128 print Color 129 print ExtendedColor 130 print OtherColor 131 print MergedColor 132 133def _test2(): 134 135 class Color(FullEnum): 136 red = 1 137 green = 2 138 blue = 3 139 140 print Color.red 141 142 print repr(Color.red) 143 print Color.red == Color.red 144 print Color.red == Color.blue 145 print Color.red == 1 146 print Color.red == 2 147 148 class ExtendedColor(Color): 149 white = 0 150 orange = 4 151 yellow = 5 152 purple = 6 153 black = 7 154 155 print ExtendedColor.orange 156 print ExtendedColor.red 157 158 print Color.red == ExtendedColor.red 159 160 class OtherColor(FullEnum): 161 white = 4 162 blue = 5 163 164 class MergedColor(Color, OtherColor): 165 pass 166 167 print MergedColor.red 168 print MergedColor.white 169 170 print Color 171 print ExtendedColor 172 print OtherColor 173 print MergedColor 174 175if __name__ == '__main__': 176 _test() 177 _test2() 178