Demo/metaclasses/Enum.py

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