1"""Enumeration metaclass.
2
3XXX This is very much a work in progress.
4
5"""
6
7import string
8
9class 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
90class 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.
123Enum = EnumMetaClass("Enum", (), {})
124
125
126def _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
168if __name__ == '__main__':
169    _test()
170