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