Welcome, guest | Sign In | My Account | Store | Cart

Inspired by various threads[1] on the python-ideas list, here's a simple enum type. Enums are generated either via the class syntax or via Enum.make().

class Breakfast(Enum):
    SPAM, HAM, EGGS
    BACON, SAUSAGE

Breakfast = Enum.make("SPAM HAM EGGS BACON SAUSAGE")

Here are some of the features:

Enum:

  • inheriting from an enum inherits copies of its values.
  • the export() method allows for exposing an enum's values in another namespace.

Enum Values:

  • the underlying values within an enum are essentially useless, diminishing the temptation to rely on them.
  • identity is equality, like with None, True, and False.
  • enum values support bitwise operations within the same enum.
  • the result of a bitwise operation is always the same object given the same inputs.

[1] see http://mail.python.org/pipermail/python-ideas/2013-January/019003.html

Python, 184 lines
  1
  2
  3
  4
  5
  6
  7
  8
  9
 10
 11
 12
 13
 14
 15
 16
 17
 18
 19
 20
 21
 22
 23
 24
 25
 26
 27
 28
 29
 30
 31
 32
 33
 34
 35
 36
 37
 38
 39
 40
 41
 42
 43
 44
 45
 46
 47
 48
 49
 50
 51
 52
 53
 54
 55
 56
 57
 58
 59
 60
 61
 62
 63
 64
 65
 66
 67
 68
 69
 70
 71
 72
 73
 74
 75
 76
 77
 78
 79
 80
 81
 82
 83
 84
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
"""enum.py

This module provides a simple enum type, Enum, which may be subclassed
to create a new enum type:

>>> class Breakfast(Enum):
...     SPAM, HAM, EGGS
...     BACON, SAUSAGE
>>> Breakfast.SPAM
EnumValue(<Breakfast.SPAM>)

Enum values implement the bitwise operators:

>>> Breakfast.SPAM | Breakfast.HAM
EnumValue(<Breakfast.(SPAM|HAM)>)

The idea of using __prepare__() to make this work came from Michael Foord:

http://mail.python.org/pipermail/python-ideas/2013-January/019108.html

"""
__all__ = ['Enum']


class EnumValue:
    """An enum value, associated with an Enum.

    EnumValue instances behave much like frozen sets, particularly with
    regard to bitwise operations (&, ^, |) as well implementing
    subtraction and inversion (~).

    """

    def __init__(self, enum, name, values=None):
        # XXX restrict name to all caps?
        self._enum = enum
        self._name = name
        self._values = values

    def __repr__(self):
        if self._name is not None:
            values = self._name
        else:
            values = "({})".format("|".join(v._name for v in self.values))
        return "{}(<{}.{}>)".format(self.__class__.__name__,
                                    self._enum.__name__, values)

    def __eq__(self, other):
        return self is other

    def __hash__(self):
        return hash(self._name or self.values)

    def _merge(self, op, other):
        try:
            if self._enum != other._enum:
                return NotImplemented
            values = op(other.values)
        except AttributeError:
            return NotImplemented
        else:
            if not values:
                raise TypeError("An EnumValue must not be empty")
            return self._enum._join(values)

    def __and__(self, other):
        return self._merge(self.values.__and__, other)

    def __xor__(self, other):
        return self._merge(self.values.__xor__, other)

    def __or__(self, other):
        return self._merge(self.values.__or__, other)

    def __sub__(self, other):
        return self._merge(self.values.__sub__, other)

    def __invert__(self):
        inverted = self._enum.values - self.values
        return self._enum.join(inverted)

    @property
    def values(self):
        if self._name is not None:
            return frozenset([self])
        else:
            return self._values


class _EnumNamespace(dict):
    def __init__(self):
        self._names = set()
    def __getitem__(self, key):
        try:
            return dict.__getitem__(self, key)
        except KeyError:
            if key.startswith("_"):
                raise
            self._names.add(key)
    def __setitem__(self, key, value):
        if key.startswith("_"):
            dict.__setitem__(self, key, value)
        else:
            raise KeyError(key)
    def __delitem__(self, key):
        raise KeyError(key)


class _EnumMeta(type):

    @classmethod
    def make(meta, names, name='SimpleEnum'):
        """Return a new Enum subclass based on the passed names."""
        if isinstance(names, str):
            names = names.replace(',', ' ').split()
        else:
            names = list(map(str, names))

        ns = _EnumNamespace()
        ns._names = names
        return meta(name, (Enum,), ns)

    @classmethod
    def __prepare__(meta, name, bases):
        return _EnumNamespace()

    def __init__(cls, name, bases, namespace):
        # validation
        if len(bases) > 1:
            raise TypeError("expected 1 base class, got {}".format(len(bases)))
        elif not bases or bases == (object,):
            pass
        elif not issubclass(bases[0], Enum):
            raise TypeError("enums can only inherit from Enum")
        else:
            for value in bases[0].values:
                namespace._names.add(value._name)


        # set the values
        values = frozenset(EnumValue(cls, name) for name in namespace._names)
        super().__setattr__('values', values)
        for value in values:
            super().__setattr__(value._name, value)

        super().__setattr__('_joined_values', {})

    def __setattr__(cls, name, value):
        raise TypeError("enums are immutable")

    def __delattr__(cls, name):
        raise TypeError("enums are immutable")

    def __len__(cls):
        return len(cls._values)

    def __contains__(cls, value):
        return value in cls._values

    def __iter__(cls):
        return iter(cls._values)

    def _join(cls, values):
        # XXX make sure the values are in the enum?
        if values in cls._joined_values:
            return cls._joined_values[values]
        elif len(values) == 1:
            value = next(iter(values))
            cls._joined_values[values] = value
            return value
        else:
            joined = EnumValue(cls, None, values)
            cls._joined_values[values] = joined
            return joined

    def export(cls, namespace):
        """Copy the enums values into the namespace."""
        namespace.update((v._name, v) for v in cls.values)


class Enum(metaclass=_EnumMeta):
    """A base enum type, strictly for inheriting."""
    def __new__(cls, *args, **kwargs):
        raise TypeError("enums do not have instances")
Created by Eric Snow on Wed, 13 Feb 2013 (BSD)
Python recipes (4591)
Eric Snow's recipes (39)

Required Modules

  • (none specified)

Other Information and Tasks