Demo/newmetaclasses/Enum.py

   1 """Enumeration metaclass."""
   2
   3 class 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
  46 class 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
  61 class 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
  85 class Enum:
  86     __metaclass__ = EnumMetaclass
  87
  88 class FullEnum:
  89     __metaclass__ = FullEnumMetaclass
  90
  91 def _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
 133 def _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
 175 if __name__ == '__main__':
 176     _test()
 177     _test2()