1 """Enumeration metaclass.
2 
3 XXX This is very much a work in progress.
4 
5 """
6 
7 import string
8 
9 class EnumMetaClass:
10     """Metaclass for enumeration.
11 
12     To define your own enumeration, do something like
13 
14     class Color(Enum):
15         red = 1
16         green = 2
17         blue = 3
18 
19     Now, Color.red, Color.green and Color.blue behave totally
20     different: they are enumerated values, not integers.
21 
22     Enumerations cannot be instantiated; however they can be
23     subclassed.
24 
25     """
26 
27     def __init__(self, name, bases, dict):
28         """Constructor -- create an enumeration.
29 
30         Called at the end of the class statement.  The arguments are
31         the name of the new class, a tuple containing the base
32         classes, and a dictionary containing everything that was
33         entered in the class' namespace during execution of the class
34         statement.  In the above example, it would be {'red': 1,
35         'green': 2, 'blue': 3}.
36 
37         """
38         for base in bases:
39             if base.__class__ is not EnumMetaClass:
40                 raise TypeError, "Enumeration base class must be enumeration"
41         bases = filter(lambda x: x is not Enum, bases)
42         self.__name__ = name
43         self.__bases__ = bases
44         self.__dict = {}
45         for key, value in dict.items():
46             self.__dict[key] = EnumInstance(name, key, value)
47 
48     def __getattr__(self, name):
49         """Return an enumeration value.
50 
51         For example, Color.red returns the value corresponding to red.
52 
53         XXX Perhaps the values should be created in the constructor?
54 
55         This looks in the class dictionary and if it is not found
56         there asks the base classes.
57 
58         The special attribute __members__ returns the list of names
59         defined in this class (it does not merge in the names defined
60         in base classes).
61 
62         """
63         if name == '__members__':
64             return self.__dict.keys()
65 
66         try:
67             return self.__dict[name]
68         except KeyError:
69             for base in self.__bases__:
70                 try:
71                     return getattr(base, name)
72                 except AttributeError:
73                     continue
74 
75         raise AttributeError, name
76 
77     def __repr__(self):
78         s = self.__name__
79         if self.__bases__:
80             s = s + '(' + string.join(map(lambda x: x.__name__,
81                                           self.__bases__), ", ") + ')'
82         if self.__dict:
83             list = []
84             for key, value in self.__dict.items():
85                 list.append("%s: %s" % (key, int(value)))
86             s = "%s: {%s}" % (s, string.join(list, ", "))
87         return s
88 
89 
90 class EnumInstance:
91     """Class to represent an enumeration value.
92 
93     EnumInstance('Color', 'red', 12) prints as 'Color.red' and behaves
94     like the integer 12 when compared, but doesn't support arithmetic.
95 
96     XXX Should it record the actual enumeration rather than just its
97     name?
98 
99     """
100 
101     def __init__(self, classname, enumname, value):
102         self.__classname = classname
103         self.__enumname = enumname
104         self.__value = value
105 
106     def __int__(self):
107         return self.__value
108 
109     def __repr__(self):
110         return "EnumInstance(%r, %r, %r)" % (self.__classname,
111                                              self.__enumname,
112                                              self.__value)
113 
114     def __str__(self):
115         return "%s.%s" % (self.__classname, self.__enumname)
116 
117     def __cmp__(self, other):
118         return cmp(self.__value, int(other))
119 
120 
121 # Create the base class for enumerations.
122 # It is an empty enumeration.
123 Enum = EnumMetaClass("Enum", (), {})
124 
125 
126 def _test():
127 
128     class Color(Enum):
129         red = 1
130         green = 2
131         blue = 3
132 
133     print Color.red
134     print dir(Color)
135 
136     print Color.red == Color.red
137     print Color.red == Color.blue
138     print Color.red == 1
139     print Color.red == 2
140 
141     class ExtendedColor(Color):
142         white = 0
143         orange = 4
144         yellow = 5
145         purple = 6
146         black = 7
147 
148     print ExtendedColor.orange
149     print ExtendedColor.red
150 
151     print Color.red == ExtendedColor.red
152 
153     class OtherColor(Enum):
154         white = 4
155         blue = 5
156 
157     class MergedColor(Color, OtherColor):
158         pass
159 
160     print MergedColor.red
161     print MergedColor.white
162 
163     print Color
164     print ExtendedColor
165     print OtherColor
166     print MergedColor
167 
168 if __name__ == '__main__':
169     _test()
170