/Demo/newmetaclasses/Enum.py

  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()