blob: 73dd613887d22b55d15bf7be18a927db446506b2 [file] [log] [blame]
Ethan Furman6b3d64a2013-06-14 16:55:46 -07001import sys
Ethan Furmane03ea372013-09-25 07:14:41 -07002from types import MappingProxyType, DynamicClassAttribute
Ethan Furmanee47e5c2016-08-31 00:12:15 -07003from functools import reduce
Raymond Hettinger381dc6c2016-12-29 23:57:12 -07004from operator import or_ as _or_
Ethan Furman6b3d64a2013-06-14 16:55:46 -07005
Ethan Furmanc7915072015-09-17 22:55:40 -07006# try _collections first to reduce startup cost
Ethan Furmane5754ab2015-09-17 22:03:52 -07007try:
8 from _collections import OrderedDict
9except ImportError:
10 from collections import OrderedDict
11
12
Ethan Furmanc16595e2016-09-10 23:36:59 -070013__all__ = [
14 'EnumMeta',
15 'Enum', 'IntEnum', 'Flag', 'IntFlag',
16 'auto', 'unique',
17 ]
Ethan Furman6b3d64a2013-06-14 16:55:46 -070018
19
Ethan Furman101e0742013-09-15 12:34:36 -070020def _is_descriptor(obj):
21 """Returns True if obj is a descriptor, False otherwise."""
22 return (
23 hasattr(obj, '__get__') or
24 hasattr(obj, '__set__') or
25 hasattr(obj, '__delete__'))
26
27
Ethan Furman6b3d64a2013-06-14 16:55:46 -070028def _is_dunder(name):
29 """Returns True if a __dunder__ name, False otherwise."""
30 return (name[:2] == name[-2:] == '__' and
31 name[2:3] != '_' and
Ethan Furman648f8602013-10-06 17:19:54 -070032 name[-3:-2] != '_' and
33 len(name) > 4)
Ethan Furman6b3d64a2013-06-14 16:55:46 -070034
35
36def _is_sunder(name):
37 """Returns True if a _sunder_ name, False otherwise."""
38 return (name[0] == name[-1] == '_' and
39 name[1:2] != '_' and
Ethan Furman648f8602013-10-06 17:19:54 -070040 name[-2:-1] != '_' and
41 len(name) > 2)
Ethan Furman6b3d64a2013-06-14 16:55:46 -070042
Ethan Furman6b3d64a2013-06-14 16:55:46 -070043def _make_class_unpicklable(cls):
44 """Make the given class un-picklable."""
Ethan Furmanca1b7942014-02-08 11:36:27 -080045 def _break_on_call_reduce(self, proto):
Ethan Furman6b3d64a2013-06-14 16:55:46 -070046 raise TypeError('%r cannot be pickled' % self)
Ethan Furmanca1b7942014-02-08 11:36:27 -080047 cls.__reduce_ex__ = _break_on_call_reduce
Ethan Furman6b3d64a2013-06-14 16:55:46 -070048 cls.__module__ = '<unknown>'
49
Ethan Furman3515dcc2016-09-18 13:15:41 -070050_auto_null = object()
Ethan Furmanc16595e2016-09-10 23:36:59 -070051class auto:
52 """
53 Instances are replaced with an appropriate value in Enum class suites.
54 """
Ethan Furman3515dcc2016-09-18 13:15:41 -070055 value = _auto_null
Ethan Furmanc16595e2016-09-10 23:36:59 -070056
Ethan Furman101e0742013-09-15 12:34:36 -070057
Ethan Furman6b3d64a2013-06-14 16:55:46 -070058class _EnumDict(dict):
Ethan Furman101e0742013-09-15 12:34:36 -070059 """Track enum member order and ensure member names are not reused.
Ethan Furman6b3d64a2013-06-14 16:55:46 -070060
61 EnumMeta will use the names found in self._member_names as the
62 enumeration member names.
63
64 """
65 def __init__(self):
66 super().__init__()
67 self._member_names = []
Ethan Furmanc16595e2016-09-10 23:36:59 -070068 self._last_values = []
Ethan Furman6b3d64a2013-06-14 16:55:46 -070069
70 def __setitem__(self, key, value):
Ethan Furman101e0742013-09-15 12:34:36 -070071 """Changes anything not dundered or not a descriptor.
Ethan Furman6b3d64a2013-06-14 16:55:46 -070072
73 If an enum member name is used twice, an error is raised; duplicate
74 values are not checked for.
75
76 Single underscore (sunder) names are reserved.
77
78 """
79 if _is_sunder(key):
Ethan Furmanee47e5c2016-08-31 00:12:15 -070080 if key not in (
Ethan Furman3515dcc2016-09-18 13:15:41 -070081 '_order_', '_create_pseudo_member_',
Ethan Furmanee47e5c2016-08-31 00:12:15 -070082 '_generate_next_value_', '_missing_',
83 ):
Ethan Furmane8e61272016-08-20 07:19:31 -070084 raise ValueError('_names_ are reserved for future Enum use')
Ethan Furmanc16595e2016-09-10 23:36:59 -070085 if key == '_generate_next_value_':
86 setattr(self, '_generate_next_value', value)
Ethan Furman101e0742013-09-15 12:34:36 -070087 elif _is_dunder(key):
Ethan Furmane8e61272016-08-20 07:19:31 -070088 if key == '__order__':
89 key = '_order_'
Ethan Furman101e0742013-09-15 12:34:36 -070090 elif key in self._member_names:
91 # descriptor overwriting an enum?
92 raise TypeError('Attempted to reuse key: %r' % key)
93 elif not _is_descriptor(value):
94 if key in self:
95 # enum overwriting a descriptor?
Ethan Furmanee47e5c2016-08-31 00:12:15 -070096 raise TypeError('%r already defined as: %r' % (key, self[key]))
Ethan Furmanc16595e2016-09-10 23:36:59 -070097 if isinstance(value, auto):
Ethan Furman3515dcc2016-09-18 13:15:41 -070098 if value.value == _auto_null:
99 value.value = self._generate_next_value(key, 1, len(self._member_names), self._last_values[:])
100 value = value.value
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700101 self._member_names.append(key)
Ethan Furmanc16595e2016-09-10 23:36:59 -0700102 self._last_values.append(value)
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700103 super().__setitem__(key, value)
104
105
Ezio Melotti9a3777e2013-08-17 15:53:55 +0300106# Dummy value for Enum as EnumMeta explicitly checks for it, but of course
107# until EnumMeta finishes running the first time the Enum class doesn't exist.
108# This is also why there are checks in EnumMeta like `if Enum is not None`
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700109Enum = None
110
Ethan Furman332dbc72016-08-20 00:00:52 -0700111
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700112class EnumMeta(type):
113 """Metaclass for Enum"""
114 @classmethod
Ethan Furman332dbc72016-08-20 00:00:52 -0700115 def __prepare__(metacls, cls, bases):
Ethan Furmanee47e5c2016-08-31 00:12:15 -0700116 # create the namespace dict
117 enum_dict = _EnumDict()
118 # inherit previous flags and _generate_next_value_ function
119 member_type, first_enum = metacls._get_mixins_(bases)
120 if first_enum is not None:
121 enum_dict['_generate_next_value_'] = getattr(first_enum, '_generate_next_value_', None)
122 return enum_dict
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700123
Ethan Furman65a5a472016-09-01 23:55:19 -0700124 def __new__(metacls, cls, bases, classdict):
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700125 # an Enum class is final once enumeration items have been defined; it
126 # cannot be mixed with other types (int, float, etc.) if it has an
127 # inherited __new__ unless a new __new__ is defined (or the resulting
128 # class will fail).
129 member_type, first_enum = metacls._get_mixins_(bases)
130 __new__, save_new, use_args = metacls._find_new_(classdict, member_type,
131 first_enum)
132
133 # save enum items into separate mapping so they don't get baked into
134 # the new class
Ethan Furmanee47e5c2016-08-31 00:12:15 -0700135 enum_members = {k: classdict[k] for k in classdict._member_names}
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700136 for name in classdict._member_names:
137 del classdict[name]
138
Ethan Furmane8e61272016-08-20 07:19:31 -0700139 # adjust the sunders
140 _order_ = classdict.pop('_order_', None)
141
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700142 # check for illegal enum names (any others?)
Ethan Furmanee47e5c2016-08-31 00:12:15 -0700143 invalid_names = set(enum_members) & {'mro', }
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700144 if invalid_names:
145 raise ValueError('Invalid enum member name: {0}'.format(
146 ','.join(invalid_names)))
147
Ethan Furman48a724f2015-04-11 23:23:06 -0700148 # create a default docstring if one has not been provided
149 if '__doc__' not in classdict:
150 classdict['__doc__'] = 'An enumeration.'
151
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700152 # create our new Enum type
153 enum_class = super().__new__(metacls, cls, bases, classdict)
Ethan Furman520ad572013-07-19 19:47:21 -0700154 enum_class._member_names_ = [] # names in definition order
155 enum_class._member_map_ = OrderedDict() # name->value map
Ethan Furman5e5a8232013-08-04 08:42:23 -0700156 enum_class._member_type_ = member_type
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700157
Ethan Furman354ecf12015-03-11 08:43:12 -0700158 # save attributes from super classes so we know if we can take
159 # the shortcut of storing members in the class dict
Ethan Furman3803ad42016-05-01 10:03:53 -0700160 base_attributes = {a for b in enum_class.mro() for a in b.__dict__}
Ethan Furman354ecf12015-03-11 08:43:12 -0700161
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700162 # Reverse value->name map for hashable values.
Ethan Furman520ad572013-07-19 19:47:21 -0700163 enum_class._value2member_map_ = {}
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700164
Ethan Furman2da95042014-03-03 12:42:52 -0800165 # If a custom type is mixed into the Enum, and it does not know how
166 # to pickle itself, pickle.dumps will succeed but pickle.loads will
167 # fail. Rather than have the error show up later and possibly far
168 # from the source, sabotage the pickle protocol for this class so
169 # that pickle.dumps also fails.
170 #
171 # However, if the new class implements its own __reduce_ex__, do not
172 # sabotage -- it's on them to make sure it works correctly. We use
173 # __reduce_ex__ instead of any of the others as it is preferred by
174 # pickle over __reduce__, and it handles all pickle protocols.
175 if '__reduce_ex__' not in classdict:
Ethan Furmandc870522014-02-18 12:37:12 -0800176 if member_type is not object:
177 methods = ('__getnewargs_ex__', '__getnewargs__',
178 '__reduce_ex__', '__reduce__')
Ethan Furman2da95042014-03-03 12:42:52 -0800179 if not any(m in member_type.__dict__ for m in methods):
Ethan Furmandc870522014-02-18 12:37:12 -0800180 _make_class_unpicklable(enum_class)
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700181
182 # instantiate them, checking for duplicates as we go
183 # we instantiate first instead of checking for duplicates first in case
184 # a custom __new__ is doing something funky with the values -- such as
185 # auto-numbering ;)
186 for member_name in classdict._member_names:
Ethan Furmanee47e5c2016-08-31 00:12:15 -0700187 value = enum_members[member_name]
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700188 if not isinstance(value, tuple):
189 args = (value, )
190 else:
191 args = value
192 if member_type is tuple: # special case for tuple enums
193 args = (args, ) # wrap it one more time
194 if not use_args:
195 enum_member = __new__(enum_class)
Ethan Furmanb41803e2013-07-25 13:50:45 -0700196 if not hasattr(enum_member, '_value_'):
197 enum_member._value_ = value
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700198 else:
199 enum_member = __new__(enum_class, *args)
Ethan Furmanb41803e2013-07-25 13:50:45 -0700200 if not hasattr(enum_member, '_value_'):
Ethan Furmanee47e5c2016-08-31 00:12:15 -0700201 if member_type is object:
202 enum_member._value_ = value
203 else:
204 enum_member._value_ = member_type(*args)
Ethan Furman520ad572013-07-19 19:47:21 -0700205 value = enum_member._value_
Ethan Furman520ad572013-07-19 19:47:21 -0700206 enum_member._name_ = member_name
Ethan Furmanc850f342013-09-15 16:59:35 -0700207 enum_member.__objclass__ = enum_class
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700208 enum_member.__init__(*args)
209 # If another member with the same value was already defined, the
210 # new member becomes an alias to the existing one.
Ethan Furman520ad572013-07-19 19:47:21 -0700211 for name, canonical_member in enum_class._member_map_.items():
Ethan Furman0081f232014-09-16 17:31:23 -0700212 if canonical_member._value_ == enum_member._value_:
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700213 enum_member = canonical_member
214 break
215 else:
216 # Aliases don't appear in member names (only in __members__).
Ethan Furman520ad572013-07-19 19:47:21 -0700217 enum_class._member_names_.append(member_name)
Ethan Furman354ecf12015-03-11 08:43:12 -0700218 # performance boost for any member that would not shadow
219 # a DynamicClassAttribute
220 if member_name not in base_attributes:
221 setattr(enum_class, member_name, enum_member)
222 # now add to _member_map_
Ethan Furman520ad572013-07-19 19:47:21 -0700223 enum_class._member_map_[member_name] = enum_member
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700224 try:
225 # This may fail if value is not hashable. We can't add the value
226 # to the map, and by-value lookups for this value will be
227 # linear.
Ethan Furman520ad572013-07-19 19:47:21 -0700228 enum_class._value2member_map_[value] = enum_member
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700229 except TypeError:
230 pass
231
232 # double check that repr and friends are not the mixin's or various
233 # things break (such as pickle)
Ethan Furmandc870522014-02-18 12:37:12 -0800234 for name in ('__repr__', '__str__', '__format__', '__reduce_ex__'):
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700235 class_method = getattr(enum_class, name)
236 obj_method = getattr(member_type, name, None)
237 enum_method = getattr(first_enum, name, None)
238 if obj_method is not None and obj_method is class_method:
239 setattr(enum_class, name, enum_method)
240
241 # replace any other __new__ with our own (as long as Enum is not None,
242 # anyway) -- again, this is to support pickle
243 if Enum is not None:
244 # if the user defined their own __new__, save it before it gets
245 # clobbered in case they subclass later
246 if save_new:
247 enum_class.__new_member__ = __new__
248 enum_class.__new__ = Enum.__new__
Ethan Furmane8e61272016-08-20 07:19:31 -0700249
250 # py3 support for definition order (helps keep py2/py3 code in sync)
251 if _order_ is not None:
252 if isinstance(_order_, str):
253 _order_ = _order_.replace(',', ' ').split()
254 if _order_ != enum_class._member_names_:
255 raise TypeError('member order does not match _order_')
256
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700257 return enum_class
258
Ethan Furman5de67b12016-04-13 23:52:09 -0700259 def __bool__(self):
260 """
261 classes/types should always be True.
262 """
263 return True
264
Ethan Furmand9925a12014-09-16 20:35:55 -0700265 def __call__(cls, value, names=None, *, module=None, qualname=None, type=None, start=1):
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700266 """Either returns an existing member, or creates a new enum class.
267
268 This method is used both when an enum class is given a value to match
269 to an enumeration member (i.e. Color(3)) and for the functional API
Ethan Furman23bb6f42016-11-21 09:22:05 -0800270 (i.e. Color = Enum('Color', names='RED GREEN BLUE')).
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700271
Ethan Furman2da95042014-03-03 12:42:52 -0800272 When used for the functional API:
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700273
Ethan Furman2da95042014-03-03 12:42:52 -0800274 `value` will be the name of the new class.
275
276 `names` should be either a string of white-space/comma delimited names
Ethan Furmand9925a12014-09-16 20:35:55 -0700277 (values will start at `start`), or an iterator/mapping of name, value pairs.
Ethan Furman2da95042014-03-03 12:42:52 -0800278
279 `module` should be set to the module this class is being created in;
280 if it is not set, an attempt to find that module will be made, but if
281 it fails the class will not be picklable.
282
283 `qualname` should be set to the actual location this class can be found
284 at in its module; by default it is set to the global scope. If this is
285 not correct, unpickling will fail in some circumstances.
286
287 `type`, if set, will be mixed in as the first base class.
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700288
289 """
290 if names is None: # simple value lookup
291 return cls.__new__(cls, value)
292 # otherwise, functional API: we're creating a new Enum type
Ethan Furmand9925a12014-09-16 20:35:55 -0700293 return cls._create_(value, names, module=module, qualname=qualname, type=type, start=start)
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700294
295 def __contains__(cls, member):
Ethan Furman0081f232014-09-16 17:31:23 -0700296 return isinstance(member, cls) and member._name_ in cls._member_map_
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700297
Ethan Furman64a99722013-09-22 16:18:19 -0700298 def __delattr__(cls, attr):
299 # nicer error message when someone tries to delete an attribute
300 # (see issue19025).
301 if attr in cls._member_map_:
302 raise AttributeError(
303 "%s: cannot delete Enum member." % cls.__name__)
304 super().__delattr__(attr)
305
Ethan Furman388a3922013-08-12 06:51:41 -0700306 def __dir__(self):
Ethan Furman64a99722013-09-22 16:18:19 -0700307 return (['__class__', '__doc__', '__members__', '__module__'] +
308 self._member_names_)
Ethan Furman388a3922013-08-12 06:51:41 -0700309
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700310 def __getattr__(cls, name):
311 """Return the enum member matching `name`
312
313 We use __getattr__ instead of descriptors or inserting into the enum
314 class' __dict__ in order to support `name` and `value` being both
315 properties for enum members (which live in the class' __dict__) and
316 enum members themselves.
317
318 """
319 if _is_dunder(name):
320 raise AttributeError(name)
321 try:
Ethan Furman520ad572013-07-19 19:47:21 -0700322 return cls._member_map_[name]
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700323 except KeyError:
324 raise AttributeError(name) from None
325
326 def __getitem__(cls, name):
Ethan Furman520ad572013-07-19 19:47:21 -0700327 return cls._member_map_[name]
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700328
329 def __iter__(cls):
Ethan Furman520ad572013-07-19 19:47:21 -0700330 return (cls._member_map_[name] for name in cls._member_names_)
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700331
332 def __len__(cls):
Ethan Furman520ad572013-07-19 19:47:21 -0700333 return len(cls._member_names_)
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700334
Ethan Furman2131a4a2013-09-14 18:11:24 -0700335 @property
336 def __members__(cls):
337 """Returns a mapping of member name->value.
338
339 This mapping lists all enum members, including aliases. Note that this
340 is a read-only view of the internal mapping.
341
342 """
343 return MappingProxyType(cls._member_map_)
344
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700345 def __repr__(cls):
346 return "<enum %r>" % cls.__name__
347
Ethan Furman2131a4a2013-09-14 18:11:24 -0700348 def __reversed__(cls):
349 return (cls._member_map_[name] for name in reversed(cls._member_names_))
350
Ethan Furmanf203f2d2013-09-06 07:16:48 -0700351 def __setattr__(cls, name, value):
352 """Block attempts to reassign Enum members.
353
354 A simple assignment to the class namespace only changes one of the
355 several possible ways to get an Enum member from the Enum class,
356 resulting in an inconsistent Enumeration.
357
358 """
359 member_map = cls.__dict__.get('_member_map_', {})
360 if name in member_map:
361 raise AttributeError('Cannot reassign members.')
362 super().__setattr__(name, value)
363
Ethan Furmand9925a12014-09-16 20:35:55 -0700364 def _create_(cls, class_name, names=None, *, module=None, qualname=None, type=None, start=1):
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700365 """Convenience method to create a new Enum class.
366
367 `names` can be:
368
369 * A string containing member names, separated either with spaces or
Ethan Furmand9925a12014-09-16 20:35:55 -0700370 commas. Values are incremented by 1 from `start`.
371 * An iterable of member names. Values are incremented by 1 from `start`.
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700372 * An iterable of (member name, value) pairs.
Ethan Furmand9925a12014-09-16 20:35:55 -0700373 * A mapping of member name -> value pairs.
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700374
375 """
376 metacls = cls.__class__
377 bases = (cls, ) if type is None else (type, cls)
Ethan Furmanee47e5c2016-08-31 00:12:15 -0700378 _, first_enum = cls._get_mixins_(bases)
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700379 classdict = metacls.__prepare__(class_name, bases)
380
381 # special processing needed for names?
382 if isinstance(names, str):
383 names = names.replace(',', ' ').split()
Dong-hee Nadcc8ce42017-06-22 01:52:32 +0900384 if isinstance(names, (tuple, list)) and names and isinstance(names[0], str):
Ethan Furmanee47e5c2016-08-31 00:12:15 -0700385 original_names, names = names, []
Ethan Furmanc16595e2016-09-10 23:36:59 -0700386 last_values = []
Ethan Furmanee47e5c2016-08-31 00:12:15 -0700387 for count, name in enumerate(original_names):
Ethan Furmanc16595e2016-09-10 23:36:59 -0700388 value = first_enum._generate_next_value_(name, start, count, last_values[:])
389 last_values.append(value)
390 names.append((name, value))
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700391
392 # Here, names is either an iterable of (name, value) or a mapping.
393 for item in names:
394 if isinstance(item, str):
395 member_name, member_value = item, names[item]
396 else:
397 member_name, member_value = item
398 classdict[member_name] = member_value
399 enum_class = metacls.__new__(metacls, class_name, bases, classdict)
400
401 # TODO: replace the frame hack if a blessed way to know the calling
402 # module is ever developed
403 if module is None:
404 try:
405 module = sys._getframe(2).f_globals['__name__']
406 except (AttributeError, ValueError) as exc:
407 pass
408 if module is None:
409 _make_class_unpicklable(enum_class)
410 else:
411 enum_class.__module__ = module
Ethan Furmanca1b7942014-02-08 11:36:27 -0800412 if qualname is not None:
413 enum_class.__qualname__ = qualname
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700414
415 return enum_class
416
417 @staticmethod
418 def _get_mixins_(bases):
419 """Returns the type for creating enum members, and the first inherited
420 enum class.
421
422 bases: the tuple of bases that was given to __new__
423
424 """
425 if not bases:
426 return object, Enum
427
428 # double check that we are not subclassing a class with existing
429 # enumeration members; while we're at it, see if any other data
430 # type has been mixed in so we can use the correct __new__
431 member_type = first_enum = None
432 for base in bases:
433 if (base is not Enum and
434 issubclass(base, Enum) and
Ethan Furman520ad572013-07-19 19:47:21 -0700435 base._member_names_):
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700436 raise TypeError("Cannot extend enumerations")
437 # base is now the last base in bases
438 if not issubclass(base, Enum):
439 raise TypeError("new enumerations must be created as "
440 "`ClassName([mixin_type,] enum_type)`")
441
442 # get correct mix-in type (either mix-in type of Enum subclass, or
443 # first base if last base is Enum)
444 if not issubclass(bases[0], Enum):
445 member_type = bases[0] # first data type
446 first_enum = bases[-1] # enum type
447 else:
448 for base in bases[0].__mro__:
449 # most common: (IntEnum, int, Enum, object)
450 # possible: (<Enum 'AutoIntEnum'>, <Enum 'IntEnum'>,
451 # <class 'int'>, <Enum 'Enum'>,
452 # <class 'object'>)
453 if issubclass(base, Enum):
454 if first_enum is None:
455 first_enum = base
456 else:
457 if member_type is None:
458 member_type = base
459
460 return member_type, first_enum
461
462 @staticmethod
463 def _find_new_(classdict, member_type, first_enum):
464 """Returns the __new__ to be used for creating the enum members.
465
466 classdict: the class dictionary given to __new__
467 member_type: the data type whose __new__ will be used by default
468 first_enum: enumeration to check for an overriding __new__
469
470 """
471 # now find the correct __new__, checking to see of one was defined
472 # by the user; also check earlier enum classes in case a __new__ was
473 # saved as __new_member__
474 __new__ = classdict.get('__new__', None)
475
476 # should __new__ be saved as __new_member__ later?
477 save_new = __new__ is not None
478
479 if __new__ is None:
480 # check all possibles for __new_member__ before falling back to
481 # __new__
482 for method in ('__new_member__', '__new__'):
483 for possible in (member_type, first_enum):
484 target = getattr(possible, method, None)
485 if target not in {
486 None,
487 None.__new__,
488 object.__new__,
489 Enum.__new__,
490 }:
491 __new__ = target
492 break
493 if __new__ is not None:
494 break
495 else:
496 __new__ = object.__new__
497
498 # if a non-object.__new__ is used then whatever value/tuple was
499 # assigned to the enum member name will be passed to __new__ and to the
500 # new enum member's __init__
501 if __new__ is object.__new__:
502 use_args = False
503 else:
504 use_args = True
505
506 return __new__, save_new, use_args
507
508
509class Enum(metaclass=EnumMeta):
510 """Generic enumeration.
511
512 Derive from this class to define new enumerations.
513
514 """
515 def __new__(cls, value):
516 # all enum instances are actually created during class construction
517 # without calling this method; this method is called by the metaclass'
518 # __call__ (i.e. Color(3) ), and by pickle
519 if type(value) is cls:
Ethan Furman23bb6f42016-11-21 09:22:05 -0800520 # For lookups like Color(Color.RED)
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700521 return value
522 # by-value search for a matching enum member
523 # see if it's in the reverse mapping (for hashable values)
Ethan Furman2aa27322013-07-19 19:35:56 -0700524 try:
Ethan Furman520ad572013-07-19 19:47:21 -0700525 if value in cls._value2member_map_:
526 return cls._value2member_map_[value]
Ethan Furman2aa27322013-07-19 19:35:56 -0700527 except TypeError:
528 # not there, now do long search -- O(n) behavior
Ethan Furman520ad572013-07-19 19:47:21 -0700529 for member in cls._member_map_.values():
Ethan Furman0081f232014-09-16 17:31:23 -0700530 if member._value_ == value:
Ethan Furman2aa27322013-07-19 19:35:56 -0700531 return member
Ethan Furmanee47e5c2016-08-31 00:12:15 -0700532 # still not found -- try _missing_ hook
533 return cls._missing_(value)
534
Ethan Furmanc16595e2016-09-10 23:36:59 -0700535 def _generate_next_value_(name, start, count, last_values):
536 for last_value in reversed(last_values):
537 try:
538 return last_value + 1
539 except TypeError:
540 pass
541 else:
Ethan Furmanee47e5c2016-08-31 00:12:15 -0700542 return start
Ethan Furmanc16595e2016-09-10 23:36:59 -0700543
Ethan Furmanee47e5c2016-08-31 00:12:15 -0700544 @classmethod
545 def _missing_(cls, value):
Ethan Furman0081f232014-09-16 17:31:23 -0700546 raise ValueError("%r is not a valid %s" % (value, cls.__name__))
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700547
548 def __repr__(self):
549 return "<%s.%s: %r>" % (
Ethan Furman520ad572013-07-19 19:47:21 -0700550 self.__class__.__name__, self._name_, self._value_)
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700551
552 def __str__(self):
Ethan Furman520ad572013-07-19 19:47:21 -0700553 return "%s.%s" % (self.__class__.__name__, self._name_)
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700554
Ethan Furman388a3922013-08-12 06:51:41 -0700555 def __dir__(self):
Ethan Furman0ae550b2014-10-14 08:58:32 -0700556 added_behavior = [
557 m
558 for cls in self.__class__.mro()
559 for m in cls.__dict__
Ethan Furman354ecf12015-03-11 08:43:12 -0700560 if m[0] != '_' and m not in self._member_map_
Ethan Furman0ae550b2014-10-14 08:58:32 -0700561 ]
Ethan Furmanec5f8eb2014-10-21 13:40:35 -0700562 return (['__class__', '__doc__', '__module__'] + added_behavior)
Ethan Furman388a3922013-08-12 06:51:41 -0700563
Ethan Furmanec15a822013-08-31 19:17:41 -0700564 def __format__(self, format_spec):
565 # mixed-in Enums should use the mixed-in type's __format__, otherwise
566 # we can get strange results with the Enum name showing up instead of
567 # the value
568
569 # pure Enum branch
570 if self._member_type_ is object:
571 cls = str
572 val = str(self)
573 # mix-in branch
574 else:
575 cls = self._member_type_
Ethan Furman0081f232014-09-16 17:31:23 -0700576 val = self._value_
Ethan Furmanec15a822013-08-31 19:17:41 -0700577 return cls.__format__(val, format_spec)
578
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700579 def __hash__(self):
Ethan Furman520ad572013-07-19 19:47:21 -0700580 return hash(self._name_)
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700581
Ethan Furmanca1b7942014-02-08 11:36:27 -0800582 def __reduce_ex__(self, proto):
Ethan Furmandc870522014-02-18 12:37:12 -0800583 return self.__class__, (self._value_, )
Ethan Furmanca1b7942014-02-08 11:36:27 -0800584
Ethan Furman33918c12013-09-27 23:02:02 -0700585 # DynamicClassAttribute is used to provide access to the `name` and
586 # `value` properties of enum members while keeping some measure of
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700587 # protection from modification, while still allowing for an enumeration
588 # to have members named `name` and `value`. This works because enumeration
589 # members are not set directly on the enum class -- __getattr__ is
590 # used to look them up.
591
Ethan Furmane03ea372013-09-25 07:14:41 -0700592 @DynamicClassAttribute
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700593 def name(self):
Ethan Furmanc850f342013-09-15 16:59:35 -0700594 """The name of the Enum member."""
Ethan Furman520ad572013-07-19 19:47:21 -0700595 return self._name_
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700596
Ethan Furmane03ea372013-09-25 07:14:41 -0700597 @DynamicClassAttribute
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700598 def value(self):
Ethan Furmanc850f342013-09-15 16:59:35 -0700599 """The value of the Enum member."""
Ethan Furman520ad572013-07-19 19:47:21 -0700600 return self._value_
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700601
Ethan Furman24e837f2015-03-18 17:27:57 -0700602 @classmethod
603 def _convert(cls, name, module, filter, source=None):
604 """
605 Create a new Enum subclass that replaces a collection of global constants
606 """
607 # convert all constants from source (or module) that pass filter() to
608 # a new Enum called name, and export the enum and its members back to
609 # module;
610 # also, replace the __reduce_ex__ method so unpickling works in
611 # previous Python versions
612 module_globals = vars(sys.modules[module])
613 if source:
614 source = vars(source)
615 else:
616 source = module_globals
Gregory P. Smith ext:(%20%5BGoogle%20Inc.%5D)6f20bd62016-06-03 19:14:52 +0000617 # We use an OrderedDict of sorted source keys so that the
618 # _value2member_map is populated in the same order every time
619 # for a consistent reverse mapping of number to name when there
620 # are multiple names for the same number rather than varying
621 # between runs due to hash randomization of the module dictionary.
Ethan Furman06339e72016-09-11 13:25:26 -0700622 members = [
623 (name, source[name])
624 for name in source.keys()
625 if filter(name)]
626 try:
627 # sort by value
628 members.sort(key=lambda t: (t[1], t[0]))
629 except TypeError:
630 # unless some values aren't comparable, in which case sort by name
631 members.sort(key=lambda t: t[0])
Ethan Furman24e837f2015-03-18 17:27:57 -0700632 cls = cls(name, members, module=module)
633 cls.__reduce_ex__ = _reduce_ex_by_name
634 module_globals.update(cls.__members__)
635 module_globals[name] = cls
636 return cls
637
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700638
639class IntEnum(int, Enum):
640 """Enum where members are also (and must be) ints"""
Ethan Furmanf24bb352013-07-18 17:05:39 -0700641
642
Ethan Furman24e837f2015-03-18 17:27:57 -0700643def _reduce_ex_by_name(self, proto):
644 return self.name
645
Ethan Furman65a5a472016-09-01 23:55:19 -0700646class Flag(Enum):
Ethan Furmanee47e5c2016-08-31 00:12:15 -0700647 """Support for flags"""
Ethan Furmanc16595e2016-09-10 23:36:59 -0700648
649 def _generate_next_value_(name, start, count, last_values):
Ethan Furmanee47e5c2016-08-31 00:12:15 -0700650 """
651 Generate the next value when not given.
652
653 name: the name of the member
654 start: the initital start value or None
655 count: the number of existing members
656 last_value: the last value assigned or None
657 """
658 if not count:
659 return start if start is not None else 1
Ethan Furmanc16595e2016-09-10 23:36:59 -0700660 for last_value in reversed(last_values):
661 try:
662 high_bit = _high_bit(last_value)
663 break
Ethan Furman3515dcc2016-09-18 13:15:41 -0700664 except Exception:
Ethan Furmanc16595e2016-09-10 23:36:59 -0700665 raise TypeError('Invalid Flag value: %r' % last_value) from None
Ethan Furmanee47e5c2016-08-31 00:12:15 -0700666 return 2 ** (high_bit+1)
667
668 @classmethod
669 def _missing_(cls, value):
670 original_value = value
671 if value < 0:
672 value = ~value
673 possible_member = cls._create_pseudo_member_(value)
Ethan Furmanee47e5c2016-08-31 00:12:15 -0700674 if original_value < 0:
675 possible_member = ~possible_member
676 return possible_member
677
678 @classmethod
679 def _create_pseudo_member_(cls, value):
Ethan Furman3515dcc2016-09-18 13:15:41 -0700680 """
681 Create a composite member iff value contains only members.
682 """
Ethan Furmanee47e5c2016-08-31 00:12:15 -0700683 pseudo_member = cls._value2member_map_.get(value, None)
684 if pseudo_member is None:
Ethan Furman3515dcc2016-09-18 13:15:41 -0700685 # verify all bits are accounted for
686 _, extra_flags = _decompose(cls, value)
687 if extra_flags:
688 raise ValueError("%r is not a valid %s" % (value, cls.__name__))
689 # construct a singleton enum pseudo-member
Ethan Furmanee47e5c2016-08-31 00:12:15 -0700690 pseudo_member = object.__new__(cls)
691 pseudo_member._name_ = None
692 pseudo_member._value_ = value
Ethan Furman28cf6632017-01-24 12:12:06 -0800693 # use setdefault in case another thread already created a composite
694 # with this value
695 pseudo_member = cls._value2member_map_.setdefault(value, pseudo_member)
Ethan Furmanee47e5c2016-08-31 00:12:15 -0700696 return pseudo_member
697
Ethan Furmanee47e5c2016-08-31 00:12:15 -0700698 def __contains__(self, other):
699 if not isinstance(other, self.__class__):
700 return NotImplemented
701 return other._value_ & self._value_ == other._value_
702
Ethan Furmanee47e5c2016-08-31 00:12:15 -0700703 def __repr__(self):
704 cls = self.__class__
705 if self._name_ is not None:
706 return '<%s.%s: %r>' % (cls.__name__, self._name_, self._value_)
Ethan Furman3515dcc2016-09-18 13:15:41 -0700707 members, uncovered = _decompose(cls, self._value_)
Ethan Furman27682d22016-09-04 11:39:01 -0700708 return '<%s.%s: %r>' % (
709 cls.__name__,
710 '|'.join([str(m._name_ or m._value_) for m in members]),
711 self._value_,
712 )
Ethan Furmanee47e5c2016-08-31 00:12:15 -0700713
714 def __str__(self):
715 cls = self.__class__
716 if self._name_ is not None:
717 return '%s.%s' % (cls.__name__, self._name_)
Ethan Furman3515dcc2016-09-18 13:15:41 -0700718 members, uncovered = _decompose(cls, self._value_)
Ethan Furmanee47e5c2016-08-31 00:12:15 -0700719 if len(members) == 1 and members[0]._name_ is None:
720 return '%s.%r' % (cls.__name__, members[0]._value_)
721 else:
722 return '%s.%s' % (
723 cls.__name__,
724 '|'.join([str(m._name_ or m._value_) for m in members]),
725 )
726
Ethan Furman25d94bb2016-09-02 16:32:32 -0700727 def __bool__(self):
728 return bool(self._value_)
729
Ethan Furmanee47e5c2016-08-31 00:12:15 -0700730 def __or__(self, other):
731 if not isinstance(other, self.__class__):
732 return NotImplemented
733 return self.__class__(self._value_ | other._value_)
734
735 def __and__(self, other):
736 if not isinstance(other, self.__class__):
737 return NotImplemented
738 return self.__class__(self._value_ & other._value_)
739
740 def __xor__(self, other):
741 if not isinstance(other, self.__class__):
742 return NotImplemented
743 return self.__class__(self._value_ ^ other._value_)
744
745 def __invert__(self):
Ethan Furman3515dcc2016-09-18 13:15:41 -0700746 members, uncovered = _decompose(self.__class__, self._value_)
747 inverted_members = [
748 m for m in self.__class__
749 if m not in members and not m._value_ & self._value_
750 ]
Ethan Furmanee47e5c2016-08-31 00:12:15 -0700751 inverted = reduce(_or_, inverted_members, self.__class__(0))
752 return self.__class__(inverted)
753
754
Ethan Furman65a5a472016-09-01 23:55:19 -0700755class IntFlag(int, Flag):
Ethan Furmanee47e5c2016-08-31 00:12:15 -0700756 """Support for integer-based Flags"""
757
758 @classmethod
Ethan Furman3515dcc2016-09-18 13:15:41 -0700759 def _missing_(cls, value):
760 if not isinstance(value, int):
761 raise ValueError("%r is not a valid %s" % (value, cls.__name__))
762 new_member = cls._create_pseudo_member_(value)
763 return new_member
764
765 @classmethod
Ethan Furmanee47e5c2016-08-31 00:12:15 -0700766 def _create_pseudo_member_(cls, value):
767 pseudo_member = cls._value2member_map_.get(value, None)
768 if pseudo_member is None:
Ethan Furman3515dcc2016-09-18 13:15:41 -0700769 need_to_create = [value]
770 # get unaccounted for bits
771 _, extra_flags = _decompose(cls, value)
772 # timer = 10
773 while extra_flags:
774 # timer -= 1
775 bit = _high_bit(extra_flags)
776 flag_value = 2 ** bit
777 if (flag_value not in cls._value2member_map_ and
778 flag_value not in need_to_create
779 ):
780 need_to_create.append(flag_value)
781 if extra_flags == -flag_value:
782 extra_flags = 0
783 else:
784 extra_flags ^= flag_value
785 for value in reversed(need_to_create):
786 # construct singleton pseudo-members
787 pseudo_member = int.__new__(cls, value)
788 pseudo_member._name_ = None
789 pseudo_member._value_ = value
Ethan Furman28cf6632017-01-24 12:12:06 -0800790 # use setdefault in case another thread already created a composite
791 # with this value
792 pseudo_member = cls._value2member_map_.setdefault(value, pseudo_member)
Ethan Furmanee47e5c2016-08-31 00:12:15 -0700793 return pseudo_member
794
Ethan Furmanee47e5c2016-08-31 00:12:15 -0700795 def __or__(self, other):
796 if not isinstance(other, (self.__class__, int)):
797 return NotImplemented
Ethan Furman3515dcc2016-09-18 13:15:41 -0700798 result = self.__class__(self._value_ | self.__class__(other)._value_)
799 return result
Ethan Furmanee47e5c2016-08-31 00:12:15 -0700800
801 def __and__(self, other):
802 if not isinstance(other, (self.__class__, int)):
803 return NotImplemented
804 return self.__class__(self._value_ & self.__class__(other)._value_)
805
806 def __xor__(self, other):
807 if not isinstance(other, (self.__class__, int)):
808 return NotImplemented
809 return self.__class__(self._value_ ^ self.__class__(other)._value_)
810
811 __ror__ = __or__
812 __rand__ = __and__
813 __rxor__ = __xor__
814
815 def __invert__(self):
Ethan Furman3515dcc2016-09-18 13:15:41 -0700816 result = self.__class__(~self._value_)
817 return result
Ethan Furmanee47e5c2016-08-31 00:12:15 -0700818
819
820def _high_bit(value):
Ethan Furman04439532016-09-02 15:50:21 -0700821 """returns index of highest bit, or -1 if value is zero or negative"""
Ethan Furman3515dcc2016-09-18 13:15:41 -0700822 return value.bit_length() - 1
Ethan Furmanee47e5c2016-08-31 00:12:15 -0700823
Ethan Furmanf24bb352013-07-18 17:05:39 -0700824def unique(enumeration):
825 """Class decorator for enumerations ensuring unique member values."""
826 duplicates = []
827 for name, member in enumeration.__members__.items():
828 if name != member.name:
829 duplicates.append((name, member.name))
830 if duplicates:
831 alias_details = ', '.join(
832 ["%s -> %s" % (alias, name) for (alias, name) in duplicates])
833 raise ValueError('duplicate values found in %r: %s' %
834 (enumeration, alias_details))
835 return enumeration
Ethan Furman3515dcc2016-09-18 13:15:41 -0700836
837def _decompose(flag, value):
838 """Extract all members from the value."""
839 # _decompose is only called if the value is not named
840 not_covered = value
841 negative = value < 0
Ethan Furman28cf6632017-01-24 12:12:06 -0800842 # issue29167: wrap accesses to _value2member_map_ in a list to avoid race
843 # conditions between iterating over it and having more psuedo-
844 # members added to it
Ethan Furman3515dcc2016-09-18 13:15:41 -0700845 if negative:
846 # only check for named flags
847 flags_to_check = [
848 (m, v)
Ethan Furman28cf6632017-01-24 12:12:06 -0800849 for v, m in list(flag._value2member_map_.items())
Ethan Furman3515dcc2016-09-18 13:15:41 -0700850 if m.name is not None
851 ]
852 else:
853 # check for named flags and powers-of-two flags
854 flags_to_check = [
855 (m, v)
Ethan Furman28cf6632017-01-24 12:12:06 -0800856 for v, m in list(flag._value2member_map_.items())
Ethan Furman3515dcc2016-09-18 13:15:41 -0700857 if m.name is not None or _power_of_two(v)
858 ]
859 members = []
860 for member, member_value in flags_to_check:
861 if member_value and member_value & value == member_value:
862 members.append(member)
863 not_covered &= ~member_value
864 if not members and value in flag._value2member_map_:
865 members.append(flag._value2member_map_[value])
866 members.sort(key=lambda m: m._value_, reverse=True)
867 if len(members) > 1 and members[0].value == value:
868 # we have the breakdown, don't need the value member itself
869 members.pop(0)
870 return members, not_covered
871
872def _power_of_two(value):
873 if value < 1:
874 return False
875 return value == 2 ** _high_bit(value)