blob: d8303204ee8becf59f314f9c6852af39a9222ed3 [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
4from 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 Furmanc16595e2016-09-10 23:36:59 -070050class auto:
51 """
52 Instances are replaced with an appropriate value in Enum class suites.
53 """
54 pass
55
Ethan Furman101e0742013-09-15 12:34:36 -070056
Ethan Furman6b3d64a2013-06-14 16:55:46 -070057class _EnumDict(dict):
Ethan Furman101e0742013-09-15 12:34:36 -070058 """Track enum member order and ensure member names are not reused.
Ethan Furman6b3d64a2013-06-14 16:55:46 -070059
60 EnumMeta will use the names found in self._member_names as the
61 enumeration member names.
62
63 """
64 def __init__(self):
65 super().__init__()
66 self._member_names = []
Ethan Furmanc16595e2016-09-10 23:36:59 -070067 self._last_values = []
Ethan Furman6b3d64a2013-06-14 16:55:46 -070068
69 def __setitem__(self, key, value):
Ethan Furman101e0742013-09-15 12:34:36 -070070 """Changes anything not dundered or not a descriptor.
Ethan Furman6b3d64a2013-06-14 16:55:46 -070071
72 If an enum member name is used twice, an error is raised; duplicate
73 values are not checked for.
74
75 Single underscore (sunder) names are reserved.
76
77 """
78 if _is_sunder(key):
Ethan Furmanee47e5c2016-08-31 00:12:15 -070079 if key not in (
80 '_order_', '_create_pseudo_member_', '_decompose_',
81 '_generate_next_value_', '_missing_',
82 ):
Ethan Furmane8e61272016-08-20 07:19:31 -070083 raise ValueError('_names_ are reserved for future Enum use')
Ethan Furmanc16595e2016-09-10 23:36:59 -070084 if key == '_generate_next_value_':
85 setattr(self, '_generate_next_value', value)
Ethan Furman101e0742013-09-15 12:34:36 -070086 elif _is_dunder(key):
Ethan Furmane8e61272016-08-20 07:19:31 -070087 if key == '__order__':
88 key = '_order_'
Ethan Furman101e0742013-09-15 12:34:36 -070089 elif key in self._member_names:
90 # descriptor overwriting an enum?
91 raise TypeError('Attempted to reuse key: %r' % key)
92 elif not _is_descriptor(value):
93 if key in self:
94 # enum overwriting a descriptor?
Ethan Furmanee47e5c2016-08-31 00:12:15 -070095 raise TypeError('%r already defined as: %r' % (key, self[key]))
Ethan Furmanc16595e2016-09-10 23:36:59 -070096 if isinstance(value, auto):
97 value = self._generate_next_value(key, 1, len(self._member_names), self._last_values[:])
Ethan Furman6b3d64a2013-06-14 16:55:46 -070098 self._member_names.append(key)
Ethan Furmanc16595e2016-09-10 23:36:59 -070099 self._last_values.append(value)
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700100 super().__setitem__(key, value)
101
102
Ezio Melotti9a3777e2013-08-17 15:53:55 +0300103# Dummy value for Enum as EnumMeta explicitly checks for it, but of course
104# until EnumMeta finishes running the first time the Enum class doesn't exist.
105# This is also why there are checks in EnumMeta like `if Enum is not None`
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700106Enum = None
107
Ethan Furman332dbc72016-08-20 00:00:52 -0700108
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700109class EnumMeta(type):
110 """Metaclass for Enum"""
111 @classmethod
Ethan Furman332dbc72016-08-20 00:00:52 -0700112 def __prepare__(metacls, cls, bases):
Ethan Furmanee47e5c2016-08-31 00:12:15 -0700113 # create the namespace dict
114 enum_dict = _EnumDict()
115 # inherit previous flags and _generate_next_value_ function
116 member_type, first_enum = metacls._get_mixins_(bases)
117 if first_enum is not None:
118 enum_dict['_generate_next_value_'] = getattr(first_enum, '_generate_next_value_', None)
119 return enum_dict
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700120
Ethan Furman65a5a472016-09-01 23:55:19 -0700121 def __new__(metacls, cls, bases, classdict):
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700122 # an Enum class is final once enumeration items have been defined; it
123 # cannot be mixed with other types (int, float, etc.) if it has an
124 # inherited __new__ unless a new __new__ is defined (or the resulting
125 # class will fail).
126 member_type, first_enum = metacls._get_mixins_(bases)
127 __new__, save_new, use_args = metacls._find_new_(classdict, member_type,
128 first_enum)
129
130 # save enum items into separate mapping so they don't get baked into
131 # the new class
Ethan Furmanee47e5c2016-08-31 00:12:15 -0700132 enum_members = {k: classdict[k] for k in classdict._member_names}
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700133 for name in classdict._member_names:
134 del classdict[name]
135
Ethan Furmane8e61272016-08-20 07:19:31 -0700136 # adjust the sunders
137 _order_ = classdict.pop('_order_', None)
138
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700139 # check for illegal enum names (any others?)
Ethan Furmanee47e5c2016-08-31 00:12:15 -0700140 invalid_names = set(enum_members) & {'mro', }
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700141 if invalid_names:
142 raise ValueError('Invalid enum member name: {0}'.format(
143 ','.join(invalid_names)))
144
Ethan Furman48a724f2015-04-11 23:23:06 -0700145 # create a default docstring if one has not been provided
146 if '__doc__' not in classdict:
147 classdict['__doc__'] = 'An enumeration.'
148
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700149 # create our new Enum type
150 enum_class = super().__new__(metacls, cls, bases, classdict)
Ethan Furman520ad572013-07-19 19:47:21 -0700151 enum_class._member_names_ = [] # names in definition order
152 enum_class._member_map_ = OrderedDict() # name->value map
Ethan Furman5e5a8232013-08-04 08:42:23 -0700153 enum_class._member_type_ = member_type
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700154
Ethan Furman354ecf12015-03-11 08:43:12 -0700155 # save attributes from super classes so we know if we can take
156 # the shortcut of storing members in the class dict
Ethan Furman3803ad42016-05-01 10:03:53 -0700157 base_attributes = {a for b in enum_class.mro() for a in b.__dict__}
Ethan Furman354ecf12015-03-11 08:43:12 -0700158
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700159 # Reverse value->name map for hashable values.
Ethan Furman520ad572013-07-19 19:47:21 -0700160 enum_class._value2member_map_ = {}
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700161
Ethan Furman2da95042014-03-03 12:42:52 -0800162 # If a custom type is mixed into the Enum, and it does not know how
163 # to pickle itself, pickle.dumps will succeed but pickle.loads will
164 # fail. Rather than have the error show up later and possibly far
165 # from the source, sabotage the pickle protocol for this class so
166 # that pickle.dumps also fails.
167 #
168 # However, if the new class implements its own __reduce_ex__, do not
169 # sabotage -- it's on them to make sure it works correctly. We use
170 # __reduce_ex__ instead of any of the others as it is preferred by
171 # pickle over __reduce__, and it handles all pickle protocols.
172 if '__reduce_ex__' not in classdict:
Ethan Furmandc870522014-02-18 12:37:12 -0800173 if member_type is not object:
174 methods = ('__getnewargs_ex__', '__getnewargs__',
175 '__reduce_ex__', '__reduce__')
Ethan Furman2da95042014-03-03 12:42:52 -0800176 if not any(m in member_type.__dict__ for m in methods):
Ethan Furmandc870522014-02-18 12:37:12 -0800177 _make_class_unpicklable(enum_class)
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700178
179 # instantiate them, checking for duplicates as we go
180 # we instantiate first instead of checking for duplicates first in case
181 # a custom __new__ is doing something funky with the values -- such as
182 # auto-numbering ;)
183 for member_name in classdict._member_names:
Ethan Furmanee47e5c2016-08-31 00:12:15 -0700184 value = enum_members[member_name]
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700185 if not isinstance(value, tuple):
186 args = (value, )
187 else:
188 args = value
189 if member_type is tuple: # special case for tuple enums
190 args = (args, ) # wrap it one more time
191 if not use_args:
192 enum_member = __new__(enum_class)
Ethan Furmanb41803e2013-07-25 13:50:45 -0700193 if not hasattr(enum_member, '_value_'):
194 enum_member._value_ = value
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700195 else:
196 enum_member = __new__(enum_class, *args)
Ethan Furmanb41803e2013-07-25 13:50:45 -0700197 if not hasattr(enum_member, '_value_'):
Ethan Furmanee47e5c2016-08-31 00:12:15 -0700198 if member_type is object:
199 enum_member._value_ = value
200 else:
201 enum_member._value_ = member_type(*args)
Ethan Furman520ad572013-07-19 19:47:21 -0700202 value = enum_member._value_
Ethan Furman520ad572013-07-19 19:47:21 -0700203 enum_member._name_ = member_name
Ethan Furmanc850f342013-09-15 16:59:35 -0700204 enum_member.__objclass__ = enum_class
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700205 enum_member.__init__(*args)
206 # If another member with the same value was already defined, the
207 # new member becomes an alias to the existing one.
Ethan Furman520ad572013-07-19 19:47:21 -0700208 for name, canonical_member in enum_class._member_map_.items():
Ethan Furman0081f232014-09-16 17:31:23 -0700209 if canonical_member._value_ == enum_member._value_:
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700210 enum_member = canonical_member
211 break
212 else:
213 # Aliases don't appear in member names (only in __members__).
Ethan Furman520ad572013-07-19 19:47:21 -0700214 enum_class._member_names_.append(member_name)
Ethan Furman354ecf12015-03-11 08:43:12 -0700215 # performance boost for any member that would not shadow
216 # a DynamicClassAttribute
217 if member_name not in base_attributes:
218 setattr(enum_class, member_name, enum_member)
219 # now add to _member_map_
Ethan Furman520ad572013-07-19 19:47:21 -0700220 enum_class._member_map_[member_name] = enum_member
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700221 try:
222 # This may fail if value is not hashable. We can't add the value
223 # to the map, and by-value lookups for this value will be
224 # linear.
Ethan Furman520ad572013-07-19 19:47:21 -0700225 enum_class._value2member_map_[value] = enum_member
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700226 except TypeError:
227 pass
228
229 # double check that repr and friends are not the mixin's or various
230 # things break (such as pickle)
Ethan Furmandc870522014-02-18 12:37:12 -0800231 for name in ('__repr__', '__str__', '__format__', '__reduce_ex__'):
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700232 class_method = getattr(enum_class, name)
233 obj_method = getattr(member_type, name, None)
234 enum_method = getattr(first_enum, name, None)
235 if obj_method is not None and obj_method is class_method:
236 setattr(enum_class, name, enum_method)
237
238 # replace any other __new__ with our own (as long as Enum is not None,
239 # anyway) -- again, this is to support pickle
240 if Enum is not None:
241 # if the user defined their own __new__, save it before it gets
242 # clobbered in case they subclass later
243 if save_new:
244 enum_class.__new_member__ = __new__
245 enum_class.__new__ = Enum.__new__
Ethan Furmane8e61272016-08-20 07:19:31 -0700246
247 # py3 support for definition order (helps keep py2/py3 code in sync)
248 if _order_ is not None:
249 if isinstance(_order_, str):
250 _order_ = _order_.replace(',', ' ').split()
251 if _order_ != enum_class._member_names_:
252 raise TypeError('member order does not match _order_')
253
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700254 return enum_class
255
Ethan Furman5de67b12016-04-13 23:52:09 -0700256 def __bool__(self):
257 """
258 classes/types should always be True.
259 """
260 return True
261
Ethan Furmand9925a12014-09-16 20:35:55 -0700262 def __call__(cls, value, names=None, *, module=None, qualname=None, type=None, start=1):
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700263 """Either returns an existing member, or creates a new enum class.
264
265 This method is used both when an enum class is given a value to match
266 to an enumeration member (i.e. Color(3)) and for the functional API
267 (i.e. Color = Enum('Color', names='red green blue')).
268
Ethan Furman2da95042014-03-03 12:42:52 -0800269 When used for the functional API:
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700270
Ethan Furman2da95042014-03-03 12:42:52 -0800271 `value` will be the name of the new class.
272
273 `names` should be either a string of white-space/comma delimited names
Ethan Furmand9925a12014-09-16 20:35:55 -0700274 (values will start at `start`), or an iterator/mapping of name, value pairs.
Ethan Furman2da95042014-03-03 12:42:52 -0800275
276 `module` should be set to the module this class is being created in;
277 if it is not set, an attempt to find that module will be made, but if
278 it fails the class will not be picklable.
279
280 `qualname` should be set to the actual location this class can be found
281 at in its module; by default it is set to the global scope. If this is
282 not correct, unpickling will fail in some circumstances.
283
284 `type`, if set, will be mixed in as the first base class.
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700285
286 """
287 if names is None: # simple value lookup
288 return cls.__new__(cls, value)
289 # otherwise, functional API: we're creating a new Enum type
Ethan Furmand9925a12014-09-16 20:35:55 -0700290 return cls._create_(value, names, module=module, qualname=qualname, type=type, start=start)
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700291
292 def __contains__(cls, member):
Ethan Furman0081f232014-09-16 17:31:23 -0700293 return isinstance(member, cls) and member._name_ in cls._member_map_
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700294
Ethan Furman64a99722013-09-22 16:18:19 -0700295 def __delattr__(cls, attr):
296 # nicer error message when someone tries to delete an attribute
297 # (see issue19025).
298 if attr in cls._member_map_:
299 raise AttributeError(
300 "%s: cannot delete Enum member." % cls.__name__)
301 super().__delattr__(attr)
302
Ethan Furman388a3922013-08-12 06:51:41 -0700303 def __dir__(self):
Ethan Furman64a99722013-09-22 16:18:19 -0700304 return (['__class__', '__doc__', '__members__', '__module__'] +
305 self._member_names_)
Ethan Furman388a3922013-08-12 06:51:41 -0700306
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700307 def __getattr__(cls, name):
308 """Return the enum member matching `name`
309
310 We use __getattr__ instead of descriptors or inserting into the enum
311 class' __dict__ in order to support `name` and `value` being both
312 properties for enum members (which live in the class' __dict__) and
313 enum members themselves.
314
315 """
316 if _is_dunder(name):
317 raise AttributeError(name)
318 try:
Ethan Furman520ad572013-07-19 19:47:21 -0700319 return cls._member_map_[name]
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700320 except KeyError:
321 raise AttributeError(name) from None
322
323 def __getitem__(cls, name):
Ethan Furman520ad572013-07-19 19:47:21 -0700324 return cls._member_map_[name]
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700325
326 def __iter__(cls):
Ethan Furman520ad572013-07-19 19:47:21 -0700327 return (cls._member_map_[name] for name in cls._member_names_)
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700328
329 def __len__(cls):
Ethan Furman520ad572013-07-19 19:47:21 -0700330 return len(cls._member_names_)
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700331
Ethan Furman2131a4a2013-09-14 18:11:24 -0700332 @property
333 def __members__(cls):
334 """Returns a mapping of member name->value.
335
336 This mapping lists all enum members, including aliases. Note that this
337 is a read-only view of the internal mapping.
338
339 """
340 return MappingProxyType(cls._member_map_)
341
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700342 def __repr__(cls):
343 return "<enum %r>" % cls.__name__
344
Ethan Furman2131a4a2013-09-14 18:11:24 -0700345 def __reversed__(cls):
346 return (cls._member_map_[name] for name in reversed(cls._member_names_))
347
Ethan Furmanf203f2d2013-09-06 07:16:48 -0700348 def __setattr__(cls, name, value):
349 """Block attempts to reassign Enum members.
350
351 A simple assignment to the class namespace only changes one of the
352 several possible ways to get an Enum member from the Enum class,
353 resulting in an inconsistent Enumeration.
354
355 """
356 member_map = cls.__dict__.get('_member_map_', {})
357 if name in member_map:
358 raise AttributeError('Cannot reassign members.')
359 super().__setattr__(name, value)
360
Ethan Furmand9925a12014-09-16 20:35:55 -0700361 def _create_(cls, class_name, names=None, *, module=None, qualname=None, type=None, start=1):
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700362 """Convenience method to create a new Enum class.
363
364 `names` can be:
365
366 * A string containing member names, separated either with spaces or
Ethan Furmand9925a12014-09-16 20:35:55 -0700367 commas. Values are incremented by 1 from `start`.
368 * An iterable of member names. Values are incremented by 1 from `start`.
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700369 * An iterable of (member name, value) pairs.
Ethan Furmand9925a12014-09-16 20:35:55 -0700370 * A mapping of member name -> value pairs.
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700371
372 """
373 metacls = cls.__class__
374 bases = (cls, ) if type is None else (type, cls)
Ethan Furmanee47e5c2016-08-31 00:12:15 -0700375 _, first_enum = cls._get_mixins_(bases)
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700376 classdict = metacls.__prepare__(class_name, bases)
377
378 # special processing needed for names?
379 if isinstance(names, str):
380 names = names.replace(',', ' ').split()
381 if isinstance(names, (tuple, list)) and isinstance(names[0], str):
Ethan Furmanee47e5c2016-08-31 00:12:15 -0700382 original_names, names = names, []
Ethan Furmanc16595e2016-09-10 23:36:59 -0700383 last_values = []
Ethan Furmanee47e5c2016-08-31 00:12:15 -0700384 for count, name in enumerate(original_names):
Ethan Furmanc16595e2016-09-10 23:36:59 -0700385 value = first_enum._generate_next_value_(name, start, count, last_values[:])
386 last_values.append(value)
387 names.append((name, value))
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700388
389 # Here, names is either an iterable of (name, value) or a mapping.
390 for item in names:
391 if isinstance(item, str):
392 member_name, member_value = item, names[item]
393 else:
394 member_name, member_value = item
395 classdict[member_name] = member_value
396 enum_class = metacls.__new__(metacls, class_name, bases, classdict)
397
398 # TODO: replace the frame hack if a blessed way to know the calling
399 # module is ever developed
400 if module is None:
401 try:
402 module = sys._getframe(2).f_globals['__name__']
403 except (AttributeError, ValueError) as exc:
404 pass
405 if module is None:
406 _make_class_unpicklable(enum_class)
407 else:
408 enum_class.__module__ = module
Ethan Furmanca1b7942014-02-08 11:36:27 -0800409 if qualname is not None:
410 enum_class.__qualname__ = qualname
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700411
412 return enum_class
413
414 @staticmethod
415 def _get_mixins_(bases):
416 """Returns the type for creating enum members, and the first inherited
417 enum class.
418
419 bases: the tuple of bases that was given to __new__
420
421 """
422 if not bases:
423 return object, Enum
424
425 # double check that we are not subclassing a class with existing
426 # enumeration members; while we're at it, see if any other data
427 # type has been mixed in so we can use the correct __new__
428 member_type = first_enum = None
429 for base in bases:
430 if (base is not Enum and
431 issubclass(base, Enum) and
Ethan Furman520ad572013-07-19 19:47:21 -0700432 base._member_names_):
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700433 raise TypeError("Cannot extend enumerations")
434 # base is now the last base in bases
435 if not issubclass(base, Enum):
436 raise TypeError("new enumerations must be created as "
437 "`ClassName([mixin_type,] enum_type)`")
438
439 # get correct mix-in type (either mix-in type of Enum subclass, or
440 # first base if last base is Enum)
441 if not issubclass(bases[0], Enum):
442 member_type = bases[0] # first data type
443 first_enum = bases[-1] # enum type
444 else:
445 for base in bases[0].__mro__:
446 # most common: (IntEnum, int, Enum, object)
447 # possible: (<Enum 'AutoIntEnum'>, <Enum 'IntEnum'>,
448 # <class 'int'>, <Enum 'Enum'>,
449 # <class 'object'>)
450 if issubclass(base, Enum):
451 if first_enum is None:
452 first_enum = base
453 else:
454 if member_type is None:
455 member_type = base
456
457 return member_type, first_enum
458
459 @staticmethod
460 def _find_new_(classdict, member_type, first_enum):
461 """Returns the __new__ to be used for creating the enum members.
462
463 classdict: the class dictionary given to __new__
464 member_type: the data type whose __new__ will be used by default
465 first_enum: enumeration to check for an overriding __new__
466
467 """
468 # now find the correct __new__, checking to see of one was defined
469 # by the user; also check earlier enum classes in case a __new__ was
470 # saved as __new_member__
471 __new__ = classdict.get('__new__', None)
472
473 # should __new__ be saved as __new_member__ later?
474 save_new = __new__ is not None
475
476 if __new__ is None:
477 # check all possibles for __new_member__ before falling back to
478 # __new__
479 for method in ('__new_member__', '__new__'):
480 for possible in (member_type, first_enum):
481 target = getattr(possible, method, None)
482 if target not in {
483 None,
484 None.__new__,
485 object.__new__,
486 Enum.__new__,
487 }:
488 __new__ = target
489 break
490 if __new__ is not None:
491 break
492 else:
493 __new__ = object.__new__
494
495 # if a non-object.__new__ is used then whatever value/tuple was
496 # assigned to the enum member name will be passed to __new__ and to the
497 # new enum member's __init__
498 if __new__ is object.__new__:
499 use_args = False
500 else:
501 use_args = True
502
503 return __new__, save_new, use_args
504
505
506class Enum(metaclass=EnumMeta):
507 """Generic enumeration.
508
509 Derive from this class to define new enumerations.
510
511 """
512 def __new__(cls, value):
513 # all enum instances are actually created during class construction
514 # without calling this method; this method is called by the metaclass'
515 # __call__ (i.e. Color(3) ), and by pickle
516 if type(value) is cls:
517 # For lookups like Color(Color.red)
518 return value
519 # by-value search for a matching enum member
520 # see if it's in the reverse mapping (for hashable values)
Ethan Furman2aa27322013-07-19 19:35:56 -0700521 try:
Ethan Furman520ad572013-07-19 19:47:21 -0700522 if value in cls._value2member_map_:
523 return cls._value2member_map_[value]
Ethan Furman2aa27322013-07-19 19:35:56 -0700524 except TypeError:
525 # not there, now do long search -- O(n) behavior
Ethan Furman520ad572013-07-19 19:47:21 -0700526 for member in cls._member_map_.values():
Ethan Furman0081f232014-09-16 17:31:23 -0700527 if member._value_ == value:
Ethan Furman2aa27322013-07-19 19:35:56 -0700528 return member
Ethan Furmanee47e5c2016-08-31 00:12:15 -0700529 # still not found -- try _missing_ hook
530 return cls._missing_(value)
531
Ethan Furmanc16595e2016-09-10 23:36:59 -0700532 def _generate_next_value_(name, start, count, last_values):
533 for last_value in reversed(last_values):
534 try:
535 return last_value + 1
536 except TypeError:
537 pass
538 else:
Ethan Furmanee47e5c2016-08-31 00:12:15 -0700539 return start
Ethan Furmanc16595e2016-09-10 23:36:59 -0700540
Ethan Furmanee47e5c2016-08-31 00:12:15 -0700541 @classmethod
542 def _missing_(cls, value):
Ethan Furman0081f232014-09-16 17:31:23 -0700543 raise ValueError("%r is not a valid %s" % (value, cls.__name__))
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700544
545 def __repr__(self):
546 return "<%s.%s: %r>" % (
Ethan Furman520ad572013-07-19 19:47:21 -0700547 self.__class__.__name__, self._name_, self._value_)
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700548
549 def __str__(self):
Ethan Furman520ad572013-07-19 19:47:21 -0700550 return "%s.%s" % (self.__class__.__name__, self._name_)
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700551
Ethan Furman388a3922013-08-12 06:51:41 -0700552 def __dir__(self):
Ethan Furman0ae550b2014-10-14 08:58:32 -0700553 added_behavior = [
554 m
555 for cls in self.__class__.mro()
556 for m in cls.__dict__
Ethan Furman354ecf12015-03-11 08:43:12 -0700557 if m[0] != '_' and m not in self._member_map_
Ethan Furman0ae550b2014-10-14 08:58:32 -0700558 ]
Ethan Furmanec5f8eb2014-10-21 13:40:35 -0700559 return (['__class__', '__doc__', '__module__'] + added_behavior)
Ethan Furman388a3922013-08-12 06:51:41 -0700560
Ethan Furmanec15a822013-08-31 19:17:41 -0700561 def __format__(self, format_spec):
562 # mixed-in Enums should use the mixed-in type's __format__, otherwise
563 # we can get strange results with the Enum name showing up instead of
564 # the value
565
566 # pure Enum branch
567 if self._member_type_ is object:
568 cls = str
569 val = str(self)
570 # mix-in branch
571 else:
572 cls = self._member_type_
Ethan Furman0081f232014-09-16 17:31:23 -0700573 val = self._value_
Ethan Furmanec15a822013-08-31 19:17:41 -0700574 return cls.__format__(val, format_spec)
575
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700576 def __hash__(self):
Ethan Furman520ad572013-07-19 19:47:21 -0700577 return hash(self._name_)
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700578
Ethan Furmanca1b7942014-02-08 11:36:27 -0800579 def __reduce_ex__(self, proto):
Ethan Furmandc870522014-02-18 12:37:12 -0800580 return self.__class__, (self._value_, )
Ethan Furmanca1b7942014-02-08 11:36:27 -0800581
Ethan Furman33918c12013-09-27 23:02:02 -0700582 # DynamicClassAttribute is used to provide access to the `name` and
583 # `value` properties of enum members while keeping some measure of
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700584 # protection from modification, while still allowing for an enumeration
585 # to have members named `name` and `value`. This works because enumeration
586 # members are not set directly on the enum class -- __getattr__ is
587 # used to look them up.
588
Ethan Furmane03ea372013-09-25 07:14:41 -0700589 @DynamicClassAttribute
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700590 def name(self):
Ethan Furmanc850f342013-09-15 16:59:35 -0700591 """The name of the Enum member."""
Ethan Furman520ad572013-07-19 19:47:21 -0700592 return self._name_
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700593
Ethan Furmane03ea372013-09-25 07:14:41 -0700594 @DynamicClassAttribute
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700595 def value(self):
Ethan Furmanc850f342013-09-15 16:59:35 -0700596 """The value of the Enum member."""
Ethan Furman520ad572013-07-19 19:47:21 -0700597 return self._value_
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700598
Ethan Furman24e837f2015-03-18 17:27:57 -0700599 @classmethod
600 def _convert(cls, name, module, filter, source=None):
601 """
602 Create a new Enum subclass that replaces a collection of global constants
603 """
604 # convert all constants from source (or module) that pass filter() to
605 # a new Enum called name, and export the enum and its members back to
606 # module;
607 # also, replace the __reduce_ex__ method so unpickling works in
608 # previous Python versions
609 module_globals = vars(sys.modules[module])
610 if source:
611 source = vars(source)
612 else:
613 source = module_globals
Gregory P. Smith ext:(%20%5BGoogle%20Inc.%5D)6f20bd62016-06-03 19:14:52 +0000614 # We use an OrderedDict of sorted source keys so that the
615 # _value2member_map is populated in the same order every time
616 # for a consistent reverse mapping of number to name when there
617 # are multiple names for the same number rather than varying
618 # between runs due to hash randomization of the module dictionary.
Ethan Furman06339e72016-09-11 13:25:26 -0700619 members = [
620 (name, source[name])
621 for name in source.keys()
622 if filter(name)]
623 try:
624 # sort by value
625 members.sort(key=lambda t: (t[1], t[0]))
626 except TypeError:
627 # unless some values aren't comparable, in which case sort by name
628 members.sort(key=lambda t: t[0])
Ethan Furman24e837f2015-03-18 17:27:57 -0700629 cls = cls(name, members, module=module)
630 cls.__reduce_ex__ = _reduce_ex_by_name
631 module_globals.update(cls.__members__)
632 module_globals[name] = cls
633 return cls
634
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700635
636class IntEnum(int, Enum):
637 """Enum where members are also (and must be) ints"""
Ethan Furmanf24bb352013-07-18 17:05:39 -0700638
639
Ethan Furman24e837f2015-03-18 17:27:57 -0700640def _reduce_ex_by_name(self, proto):
641 return self.name
642
Ethan Furman65a5a472016-09-01 23:55:19 -0700643class Flag(Enum):
Ethan Furmanee47e5c2016-08-31 00:12:15 -0700644 """Support for flags"""
Ethan Furmanc16595e2016-09-10 23:36:59 -0700645
646 def _generate_next_value_(name, start, count, last_values):
Ethan Furmanee47e5c2016-08-31 00:12:15 -0700647 """
648 Generate the next value when not given.
649
650 name: the name of the member
651 start: the initital start value or None
652 count: the number of existing members
653 last_value: the last value assigned or None
654 """
655 if not count:
656 return start if start is not None else 1
Ethan Furmanc16595e2016-09-10 23:36:59 -0700657 for last_value in reversed(last_values):
658 try:
659 high_bit = _high_bit(last_value)
660 break
661 except TypeError:
662 raise TypeError('Invalid Flag value: %r' % last_value) from None
Ethan Furmanee47e5c2016-08-31 00:12:15 -0700663 return 2 ** (high_bit+1)
664
665 @classmethod
666 def _missing_(cls, value):
667 original_value = value
668 if value < 0:
669 value = ~value
670 possible_member = cls._create_pseudo_member_(value)
671 for member in possible_member._decompose_():
672 if member._name_ is None and member._value_ != 0:
673 raise ValueError('%r is not a valid %s' % (original_value, cls.__name__))
674 if original_value < 0:
675 possible_member = ~possible_member
676 return possible_member
677
678 @classmethod
679 def _create_pseudo_member_(cls, value):
680 pseudo_member = cls._value2member_map_.get(value, None)
681 if pseudo_member is None:
682 # construct a non-singleton enum pseudo-member
683 pseudo_member = object.__new__(cls)
684 pseudo_member._name_ = None
685 pseudo_member._value_ = value
686 cls._value2member_map_[value] = pseudo_member
687 return pseudo_member
688
689 def _decompose_(self):
690 """Extract all members from the value."""
691 value = self._value_
692 members = []
693 cls = self.__class__
694 for member in sorted(cls, key=lambda m: m._value_, reverse=True):
695 while _high_bit(value) > _high_bit(member._value_):
696 unknown = self._create_pseudo_member_(2 ** _high_bit(value))
697 members.append(unknown)
698 value &= ~unknown._value_
699 if (
700 (value & member._value_ == member._value_)
701 and (member._value_ or not members)
702 ):
703 value &= ~member._value_
704 members.append(member)
705 if not members or value:
706 members.append(self._create_pseudo_member_(value))
707 members = list(members)
708 return members
709
710 def __contains__(self, other):
711 if not isinstance(other, self.__class__):
712 return NotImplemented
713 return other._value_ & self._value_ == other._value_
714
715 def __iter__(self):
716 if self.value == 0:
717 return iter([])
718 else:
719 return iter(self._decompose_())
720
721 def __repr__(self):
722 cls = self.__class__
723 if self._name_ is not None:
724 return '<%s.%s: %r>' % (cls.__name__, self._name_, self._value_)
725 members = self._decompose_()
Ethan Furman27682d22016-09-04 11:39:01 -0700726 return '<%s.%s: %r>' % (
727 cls.__name__,
728 '|'.join([str(m._name_ or m._value_) for m in members]),
729 self._value_,
730 )
Ethan Furmanee47e5c2016-08-31 00:12:15 -0700731
732 def __str__(self):
733 cls = self.__class__
734 if self._name_ is not None:
735 return '%s.%s' % (cls.__name__, self._name_)
736 members = self._decompose_()
737 if len(members) == 1 and members[0]._name_ is None:
738 return '%s.%r' % (cls.__name__, members[0]._value_)
739 else:
740 return '%s.%s' % (
741 cls.__name__,
742 '|'.join([str(m._name_ or m._value_) for m in members]),
743 )
744
Ethan Furman25d94bb2016-09-02 16:32:32 -0700745 def __bool__(self):
746 return bool(self._value_)
747
Ethan Furmanee47e5c2016-08-31 00:12:15 -0700748 def __or__(self, other):
749 if not isinstance(other, self.__class__):
750 return NotImplemented
751 return self.__class__(self._value_ | other._value_)
752
753 def __and__(self, other):
754 if not isinstance(other, self.__class__):
755 return NotImplemented
756 return self.__class__(self._value_ & other._value_)
757
758 def __xor__(self, other):
759 if not isinstance(other, self.__class__):
760 return NotImplemented
761 return self.__class__(self._value_ ^ other._value_)
762
763 def __invert__(self):
764 members = self._decompose_()
765 inverted_members = [m for m in self.__class__ if m not in members and not m._value_ & self._value_]
766 inverted = reduce(_or_, inverted_members, self.__class__(0))
767 return self.__class__(inverted)
768
769
Ethan Furman65a5a472016-09-01 23:55:19 -0700770class IntFlag(int, Flag):
Ethan Furmanee47e5c2016-08-31 00:12:15 -0700771 """Support for integer-based Flags"""
772
773 @classmethod
774 def _create_pseudo_member_(cls, value):
775 pseudo_member = cls._value2member_map_.get(value, None)
776 if pseudo_member is None:
777 # construct a non-singleton enum pseudo-member
778 pseudo_member = int.__new__(cls, value)
779 pseudo_member._name_ = None
780 pseudo_member._value_ = value
781 cls._value2member_map_[value] = pseudo_member
782 return pseudo_member
783
784 @classmethod
785 def _missing_(cls, value):
786 possible_member = cls._create_pseudo_member_(value)
787 return possible_member
788
789 def __or__(self, other):
790 if not isinstance(other, (self.__class__, int)):
791 return NotImplemented
792 return self.__class__(self._value_ | self.__class__(other)._value_)
793
794 def __and__(self, other):
795 if not isinstance(other, (self.__class__, int)):
796 return NotImplemented
797 return self.__class__(self._value_ & self.__class__(other)._value_)
798
799 def __xor__(self, other):
800 if not isinstance(other, (self.__class__, int)):
801 return NotImplemented
802 return self.__class__(self._value_ ^ self.__class__(other)._value_)
803
804 __ror__ = __or__
805 __rand__ = __and__
806 __rxor__ = __xor__
807
808 def __invert__(self):
809 # members = self._decompose_()
810 # inverted_members = [m for m in self.__class__ if m not in members and not m._value_ & self._value_]
811 # inverted = reduce(_or_, inverted_members, self.__class__(0))
812 return self.__class__(~self._value_)
813
814
815
816
817def _high_bit(value):
Ethan Furman04439532016-09-02 15:50:21 -0700818 """returns index of highest bit, or -1 if value is zero or negative"""
819 return value.bit_length() - 1 if value > 0 else -1
Ethan Furmanee47e5c2016-08-31 00:12:15 -0700820
Ethan Furmanf24bb352013-07-18 17:05:39 -0700821def unique(enumeration):
822 """Class decorator for enumerations ensuring unique member values."""
823 duplicates = []
824 for name, member in enumeration.__members__.items():
825 if name != member.name:
826 duplicates.append((name, member.name))
827 if duplicates:
828 alias_details = ', '.join(
829 ["%s -> %s" % (alias, name) for (alias, name) in duplicates])
830 raise ValueError('duplicate values found in %r: %s' %
831 (enumeration, alias_details))
832 return enumeration