blob: fe7cb20fc06cdf6ad1127ab8851047a887a719cd [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 Furman6b3d64a2013-06-14 16:55:46 -07003
Ethan Furmanc7915072015-09-17 22:55:40 -07004# try _collections first to reduce startup cost
Ethan Furmane5754ab2015-09-17 22:03:52 -07005try:
6 from _collections import OrderedDict
7except ImportError:
8 from collections import OrderedDict
9
10
Ethan Furmanc16595e2016-09-10 23:36:59 -070011__all__ = [
12 'EnumMeta',
13 'Enum', 'IntEnum', 'Flag', 'IntFlag',
14 'auto', 'unique',
15 ]
Ethan Furman6b3d64a2013-06-14 16:55:46 -070016
17
Ethan Furman101e0742013-09-15 12:34:36 -070018def _is_descriptor(obj):
19 """Returns True if obj is a descriptor, False otherwise."""
20 return (
21 hasattr(obj, '__get__') or
22 hasattr(obj, '__set__') or
23 hasattr(obj, '__delete__'))
24
25
Ethan Furman6b3d64a2013-06-14 16:55:46 -070026def _is_dunder(name):
27 """Returns True if a __dunder__ name, False otherwise."""
28 return (name[:2] == name[-2:] == '__' and
29 name[2:3] != '_' and
Ethan Furman648f8602013-10-06 17:19:54 -070030 name[-3:-2] != '_' and
31 len(name) > 4)
Ethan Furman6b3d64a2013-06-14 16:55:46 -070032
33
34def _is_sunder(name):
35 """Returns True if a _sunder_ name, False otherwise."""
36 return (name[0] == name[-1] == '_' and
37 name[1:2] != '_' and
Ethan Furman648f8602013-10-06 17:19:54 -070038 name[-2:-1] != '_' and
39 len(name) > 2)
Ethan Furman6b3d64a2013-06-14 16:55:46 -070040
Ethan Furman6b3d64a2013-06-14 16:55:46 -070041def _make_class_unpicklable(cls):
42 """Make the given class un-picklable."""
Ethan Furmanca1b7942014-02-08 11:36:27 -080043 def _break_on_call_reduce(self, proto):
Ethan Furman6b3d64a2013-06-14 16:55:46 -070044 raise TypeError('%r cannot be pickled' % self)
Ethan Furmanca1b7942014-02-08 11:36:27 -080045 cls.__reduce_ex__ = _break_on_call_reduce
Ethan Furman6b3d64a2013-06-14 16:55:46 -070046 cls.__module__ = '<unknown>'
47
Ethan Furman3515dcc2016-09-18 13:15:41 -070048_auto_null = object()
Ethan Furmanc16595e2016-09-10 23:36:59 -070049class auto:
50 """
51 Instances are replaced with an appropriate value in Enum class suites.
52 """
Ethan Furman3515dcc2016-09-18 13:15:41 -070053 value = _auto_null
Ethan Furmanc16595e2016-09-10 23:36:59 -070054
Ethan Furman101e0742013-09-15 12:34:36 -070055
Ethan Furman6b3d64a2013-06-14 16:55:46 -070056class _EnumDict(dict):
Ethan Furman101e0742013-09-15 12:34:36 -070057 """Track enum member order and ensure member names are not reused.
Ethan Furman6b3d64a2013-06-14 16:55:46 -070058
59 EnumMeta will use the names found in self._member_names as the
60 enumeration member names.
61
62 """
63 def __init__(self):
64 super().__init__()
65 self._member_names = []
Ethan Furmanc16595e2016-09-10 23:36:59 -070066 self._last_values = []
Ethan Furman6b3d64a2013-06-14 16:55:46 -070067
68 def __setitem__(self, key, value):
Ethan Furman101e0742013-09-15 12:34:36 -070069 """Changes anything not dundered or not a descriptor.
Ethan Furman6b3d64a2013-06-14 16:55:46 -070070
71 If an enum member name is used twice, an error is raised; duplicate
72 values are not checked for.
73
74 Single underscore (sunder) names are reserved.
75
76 """
77 if _is_sunder(key):
Ethan Furmanee47e5c2016-08-31 00:12:15 -070078 if key not in (
Ethan Furman3515dcc2016-09-18 13:15:41 -070079 '_order_', '_create_pseudo_member_',
Ethan Furmanee47e5c2016-08-31 00:12:15 -070080 '_generate_next_value_', '_missing_',
81 ):
Ethan Furmane8e61272016-08-20 07:19:31 -070082 raise ValueError('_names_ are reserved for future Enum use')
Ethan Furmanc16595e2016-09-10 23:36:59 -070083 if key == '_generate_next_value_':
84 setattr(self, '_generate_next_value', value)
Ethan Furman101e0742013-09-15 12:34:36 -070085 elif _is_dunder(key):
Ethan Furmane8e61272016-08-20 07:19:31 -070086 if key == '__order__':
87 key = '_order_'
Ethan Furman101e0742013-09-15 12:34:36 -070088 elif key in self._member_names:
89 # descriptor overwriting an enum?
90 raise TypeError('Attempted to reuse key: %r' % key)
91 elif not _is_descriptor(value):
92 if key in self:
93 # enum overwriting a descriptor?
Ethan Furmanee47e5c2016-08-31 00:12:15 -070094 raise TypeError('%r already defined as: %r' % (key, self[key]))
Ethan Furmanc16595e2016-09-10 23:36:59 -070095 if isinstance(value, auto):
Ethan Furman3515dcc2016-09-18 13:15:41 -070096 if value.value == _auto_null:
97 value.value = self._generate_next_value(key, 1, len(self._member_names), self._last_values[:])
98 value = value.value
Ethan Furman6b3d64a2013-06-14 16:55:46 -070099 self._member_names.append(key)
Ethan Furmanc16595e2016-09-10 23:36:59 -0700100 self._last_values.append(value)
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700101 super().__setitem__(key, value)
102
103
Ezio Melotti9a3777e2013-08-17 15:53:55 +0300104# Dummy value for Enum as EnumMeta explicitly checks for it, but of course
105# until EnumMeta finishes running the first time the Enum class doesn't exist.
106# This is also why there are checks in EnumMeta like `if Enum is not None`
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700107Enum = None
108
Ethan Furman332dbc72016-08-20 00:00:52 -0700109
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700110class EnumMeta(type):
111 """Metaclass for Enum"""
112 @classmethod
Ethan Furman332dbc72016-08-20 00:00:52 -0700113 def __prepare__(metacls, cls, bases):
Ethan Furmanee47e5c2016-08-31 00:12:15 -0700114 # create the namespace dict
115 enum_dict = _EnumDict()
116 # inherit previous flags and _generate_next_value_ function
117 member_type, first_enum = metacls._get_mixins_(bases)
118 if first_enum is not None:
119 enum_dict['_generate_next_value_'] = getattr(first_enum, '_generate_next_value_', None)
120 return enum_dict
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700121
Ethan Furman65a5a472016-09-01 23:55:19 -0700122 def __new__(metacls, cls, bases, classdict):
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700123 # an Enum class is final once enumeration items have been defined; it
124 # cannot be mixed with other types (int, float, etc.) if it has an
125 # inherited __new__ unless a new __new__ is defined (or the resulting
126 # class will fail).
127 member_type, first_enum = metacls._get_mixins_(bases)
128 __new__, save_new, use_args = metacls._find_new_(classdict, member_type,
129 first_enum)
130
131 # save enum items into separate mapping so they don't get baked into
132 # the new class
Ethan Furmanee47e5c2016-08-31 00:12:15 -0700133 enum_members = {k: classdict[k] for k in classdict._member_names}
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700134 for name in classdict._member_names:
135 del classdict[name]
136
Ethan Furmane8e61272016-08-20 07:19:31 -0700137 # adjust the sunders
138 _order_ = classdict.pop('_order_', None)
139
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700140 # check for illegal enum names (any others?)
Ethan Furmanee47e5c2016-08-31 00:12:15 -0700141 invalid_names = set(enum_members) & {'mro', }
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700142 if invalid_names:
143 raise ValueError('Invalid enum member name: {0}'.format(
144 ','.join(invalid_names)))
145
Ethan Furman48a724f2015-04-11 23:23:06 -0700146 # create a default docstring if one has not been provided
147 if '__doc__' not in classdict:
148 classdict['__doc__'] = 'An enumeration.'
149
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700150 # create our new Enum type
151 enum_class = super().__new__(metacls, cls, bases, classdict)
Ethan Furman520ad572013-07-19 19:47:21 -0700152 enum_class._member_names_ = [] # names in definition order
153 enum_class._member_map_ = OrderedDict() # name->value map
Ethan Furman5e5a8232013-08-04 08:42:23 -0700154 enum_class._member_type_ = member_type
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700155
Ethan Furman354ecf12015-03-11 08:43:12 -0700156 # save attributes from super classes so we know if we can take
157 # the shortcut of storing members in the class dict
Ethan Furman3803ad42016-05-01 10:03:53 -0700158 base_attributes = {a for b in enum_class.mro() for a in b.__dict__}
Ethan Furman354ecf12015-03-11 08:43:12 -0700159
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700160 # Reverse value->name map for hashable values.
Ethan Furman520ad572013-07-19 19:47:21 -0700161 enum_class._value2member_map_ = {}
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700162
Ethan Furman2da95042014-03-03 12:42:52 -0800163 # If a custom type is mixed into the Enum, and it does not know how
164 # to pickle itself, pickle.dumps will succeed but pickle.loads will
165 # fail. Rather than have the error show up later and possibly far
166 # from the source, sabotage the pickle protocol for this class so
167 # that pickle.dumps also fails.
168 #
169 # However, if the new class implements its own __reduce_ex__, do not
170 # sabotage -- it's on them to make sure it works correctly. We use
171 # __reduce_ex__ instead of any of the others as it is preferred by
172 # pickle over __reduce__, and it handles all pickle protocols.
173 if '__reduce_ex__' not in classdict:
Ethan Furmandc870522014-02-18 12:37:12 -0800174 if member_type is not object:
175 methods = ('__getnewargs_ex__', '__getnewargs__',
176 '__reduce_ex__', '__reduce__')
Ethan Furman2da95042014-03-03 12:42:52 -0800177 if not any(m in member_type.__dict__ for m in methods):
Ethan Furmandc870522014-02-18 12:37:12 -0800178 _make_class_unpicklable(enum_class)
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700179
180 # instantiate them, checking for duplicates as we go
181 # we instantiate first instead of checking for duplicates first in case
182 # a custom __new__ is doing something funky with the values -- such as
183 # auto-numbering ;)
184 for member_name in classdict._member_names:
Ethan Furmanee47e5c2016-08-31 00:12:15 -0700185 value = enum_members[member_name]
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700186 if not isinstance(value, tuple):
187 args = (value, )
188 else:
189 args = value
190 if member_type is tuple: # special case for tuple enums
191 args = (args, ) # wrap it one more time
192 if not use_args:
193 enum_member = __new__(enum_class)
Ethan Furmanb41803e2013-07-25 13:50:45 -0700194 if not hasattr(enum_member, '_value_'):
195 enum_member._value_ = value
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700196 else:
197 enum_member = __new__(enum_class, *args)
Ethan Furmanb41803e2013-07-25 13:50:45 -0700198 if not hasattr(enum_member, '_value_'):
Ethan Furmanee47e5c2016-08-31 00:12:15 -0700199 if member_type is object:
200 enum_member._value_ = value
201 else:
202 enum_member._value_ = member_type(*args)
Ethan Furman520ad572013-07-19 19:47:21 -0700203 value = enum_member._value_
Ethan Furman520ad572013-07-19 19:47:21 -0700204 enum_member._name_ = member_name
Ethan Furmanc850f342013-09-15 16:59:35 -0700205 enum_member.__objclass__ = enum_class
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700206 enum_member.__init__(*args)
207 # If another member with the same value was already defined, the
208 # new member becomes an alias to the existing one.
Ethan Furman520ad572013-07-19 19:47:21 -0700209 for name, canonical_member in enum_class._member_map_.items():
Ethan Furman0081f232014-09-16 17:31:23 -0700210 if canonical_member._value_ == enum_member._value_:
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700211 enum_member = canonical_member
212 break
213 else:
214 # Aliases don't appear in member names (only in __members__).
Ethan Furman520ad572013-07-19 19:47:21 -0700215 enum_class._member_names_.append(member_name)
Ethan Furman354ecf12015-03-11 08:43:12 -0700216 # performance boost for any member that would not shadow
217 # a DynamicClassAttribute
218 if member_name not in base_attributes:
219 setattr(enum_class, member_name, enum_member)
220 # now add to _member_map_
Ethan Furman520ad572013-07-19 19:47:21 -0700221 enum_class._member_map_[member_name] = enum_member
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700222 try:
223 # This may fail if value is not hashable. We can't add the value
224 # to the map, and by-value lookups for this value will be
225 # linear.
Ethan Furman520ad572013-07-19 19:47:21 -0700226 enum_class._value2member_map_[value] = enum_member
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700227 except TypeError:
228 pass
229
230 # double check that repr and friends are not the mixin's or various
231 # things break (such as pickle)
Ethan Furmandc870522014-02-18 12:37:12 -0800232 for name in ('__repr__', '__str__', '__format__', '__reduce_ex__'):
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700233 class_method = getattr(enum_class, name)
234 obj_method = getattr(member_type, name, None)
235 enum_method = getattr(first_enum, name, None)
236 if obj_method is not None and obj_method is class_method:
237 setattr(enum_class, name, enum_method)
238
239 # replace any other __new__ with our own (as long as Enum is not None,
240 # anyway) -- again, this is to support pickle
241 if Enum is not None:
242 # if the user defined their own __new__, save it before it gets
243 # clobbered in case they subclass later
244 if save_new:
245 enum_class.__new_member__ = __new__
246 enum_class.__new__ = Enum.__new__
Ethan Furmane8e61272016-08-20 07:19:31 -0700247
248 # py3 support for definition order (helps keep py2/py3 code in sync)
249 if _order_ is not None:
250 if isinstance(_order_, str):
251 _order_ = _order_.replace(',', ' ').split()
252 if _order_ != enum_class._member_names_:
253 raise TypeError('member order does not match _order_')
254
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700255 return enum_class
256
Ethan Furman5de67b12016-04-13 23:52:09 -0700257 def __bool__(self):
258 """
259 classes/types should always be True.
260 """
261 return True
262
Ethan Furmand9925a12014-09-16 20:35:55 -0700263 def __call__(cls, value, names=None, *, module=None, qualname=None, type=None, start=1):
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700264 """Either returns an existing member, or creates a new enum class.
265
266 This method is used both when an enum class is given a value to match
267 to an enumeration member (i.e. Color(3)) and for the functional API
Ethan Furman23bb6f42016-11-21 09:22:05 -0800268 (i.e. Color = Enum('Color', names='RED GREEN BLUE')).
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700269
Ethan Furman2da95042014-03-03 12:42:52 -0800270 When used for the functional API:
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700271
Ethan Furman2da95042014-03-03 12:42:52 -0800272 `value` will be the name of the new class.
273
274 `names` should be either a string of white-space/comma delimited names
Ethan Furmand9925a12014-09-16 20:35:55 -0700275 (values will start at `start`), or an iterator/mapping of name, value pairs.
Ethan Furman2da95042014-03-03 12:42:52 -0800276
277 `module` should be set to the module this class is being created in;
278 if it is not set, an attempt to find that module will be made, but if
279 it fails the class will not be picklable.
280
281 `qualname` should be set to the actual location this class can be found
282 at in its module; by default it is set to the global scope. If this is
283 not correct, unpickling will fail in some circumstances.
284
285 `type`, if set, will be mixed in as the first base class.
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700286
287 """
288 if names is None: # simple value lookup
289 return cls.__new__(cls, value)
290 # otherwise, functional API: we're creating a new Enum type
Ethan Furmand9925a12014-09-16 20:35:55 -0700291 return cls._create_(value, names, module=module, qualname=qualname, type=type, start=start)
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700292
293 def __contains__(cls, member):
Ethan Furman0081f232014-09-16 17:31:23 -0700294 return isinstance(member, cls) and member._name_ in cls._member_map_
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700295
Ethan Furman64a99722013-09-22 16:18:19 -0700296 def __delattr__(cls, attr):
297 # nicer error message when someone tries to delete an attribute
298 # (see issue19025).
299 if attr in cls._member_map_:
300 raise AttributeError(
301 "%s: cannot delete Enum member." % cls.__name__)
302 super().__delattr__(attr)
303
Ethan Furman388a3922013-08-12 06:51:41 -0700304 def __dir__(self):
Ethan Furman64a99722013-09-22 16:18:19 -0700305 return (['__class__', '__doc__', '__members__', '__module__'] +
306 self._member_names_)
Ethan Furman388a3922013-08-12 06:51:41 -0700307
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700308 def __getattr__(cls, name):
309 """Return the enum member matching `name`
310
311 We use __getattr__ instead of descriptors or inserting into the enum
312 class' __dict__ in order to support `name` and `value` being both
313 properties for enum members (which live in the class' __dict__) and
314 enum members themselves.
315
316 """
317 if _is_dunder(name):
318 raise AttributeError(name)
319 try:
Ethan Furman520ad572013-07-19 19:47:21 -0700320 return cls._member_map_[name]
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700321 except KeyError:
322 raise AttributeError(name) from None
323
324 def __getitem__(cls, name):
Ethan Furman520ad572013-07-19 19:47:21 -0700325 return cls._member_map_[name]
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700326
327 def __iter__(cls):
Ethan Furman520ad572013-07-19 19:47:21 -0700328 return (cls._member_map_[name] for name in cls._member_names_)
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700329
330 def __len__(cls):
Ethan Furman520ad572013-07-19 19:47:21 -0700331 return len(cls._member_names_)
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700332
Ethan Furman2131a4a2013-09-14 18:11:24 -0700333 @property
334 def __members__(cls):
335 """Returns a mapping of member name->value.
336
337 This mapping lists all enum members, including aliases. Note that this
338 is a read-only view of the internal mapping.
339
340 """
341 return MappingProxyType(cls._member_map_)
342
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700343 def __repr__(cls):
344 return "<enum %r>" % cls.__name__
345
Ethan Furman2131a4a2013-09-14 18:11:24 -0700346 def __reversed__(cls):
347 return (cls._member_map_[name] for name in reversed(cls._member_names_))
348
Ethan Furmanf203f2d2013-09-06 07:16:48 -0700349 def __setattr__(cls, name, value):
350 """Block attempts to reassign Enum members.
351
352 A simple assignment to the class namespace only changes one of the
353 several possible ways to get an Enum member from the Enum class,
354 resulting in an inconsistent Enumeration.
355
356 """
357 member_map = cls.__dict__.get('_member_map_', {})
358 if name in member_map:
359 raise AttributeError('Cannot reassign members.')
360 super().__setattr__(name, value)
361
Ethan Furmand9925a12014-09-16 20:35:55 -0700362 def _create_(cls, class_name, names=None, *, module=None, qualname=None, type=None, start=1):
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700363 """Convenience method to create a new Enum class.
364
365 `names` can be:
366
367 * A string containing member names, separated either with spaces or
Ethan Furmand9925a12014-09-16 20:35:55 -0700368 commas. Values are incremented by 1 from `start`.
369 * An iterable of member names. Values are incremented by 1 from `start`.
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700370 * An iterable of (member name, value) pairs.
Ethan Furmand9925a12014-09-16 20:35:55 -0700371 * A mapping of member name -> value pairs.
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700372
373 """
374 metacls = cls.__class__
375 bases = (cls, ) if type is None else (type, cls)
Ethan Furmanee47e5c2016-08-31 00:12:15 -0700376 _, first_enum = cls._get_mixins_(bases)
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700377 classdict = metacls.__prepare__(class_name, bases)
378
379 # special processing needed for names?
380 if isinstance(names, str):
381 names = names.replace(',', ' ').split()
Dong-hee Nadcc8ce42017-06-22 01:52:32 +0900382 if isinstance(names, (tuple, list)) and names and isinstance(names[0], str):
Ethan Furmanee47e5c2016-08-31 00:12:15 -0700383 original_names, names = names, []
Ethan Furmanc16595e2016-09-10 23:36:59 -0700384 last_values = []
Ethan Furmanee47e5c2016-08-31 00:12:15 -0700385 for count, name in enumerate(original_names):
Ethan Furmanc16595e2016-09-10 23:36:59 -0700386 value = first_enum._generate_next_value_(name, start, count, last_values[:])
387 last_values.append(value)
388 names.append((name, value))
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700389
390 # Here, names is either an iterable of (name, value) or a mapping.
391 for item in names:
392 if isinstance(item, str):
393 member_name, member_value = item, names[item]
394 else:
395 member_name, member_value = item
396 classdict[member_name] = member_value
397 enum_class = metacls.__new__(metacls, class_name, bases, classdict)
398
399 # TODO: replace the frame hack if a blessed way to know the calling
400 # module is ever developed
401 if module is None:
402 try:
403 module = sys._getframe(2).f_globals['__name__']
404 except (AttributeError, ValueError) as exc:
405 pass
406 if module is None:
407 _make_class_unpicklable(enum_class)
408 else:
409 enum_class.__module__ = module
Ethan Furmanca1b7942014-02-08 11:36:27 -0800410 if qualname is not None:
411 enum_class.__qualname__ = qualname
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700412
413 return enum_class
414
415 @staticmethod
416 def _get_mixins_(bases):
417 """Returns the type for creating enum members, and the first inherited
418 enum class.
419
420 bases: the tuple of bases that was given to __new__
421
422 """
423 if not bases:
424 return object, Enum
425
426 # double check that we are not subclassing a class with existing
427 # enumeration members; while we're at it, see if any other data
428 # type has been mixed in so we can use the correct __new__
429 member_type = first_enum = None
430 for base in bases:
431 if (base is not Enum and
432 issubclass(base, Enum) and
Ethan Furman520ad572013-07-19 19:47:21 -0700433 base._member_names_):
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700434 raise TypeError("Cannot extend enumerations")
435 # base is now the last base in bases
436 if not issubclass(base, Enum):
437 raise TypeError("new enumerations must be created as "
438 "`ClassName([mixin_type,] enum_type)`")
439
440 # get correct mix-in type (either mix-in type of Enum subclass, or
441 # first base if last base is Enum)
442 if not issubclass(bases[0], Enum):
443 member_type = bases[0] # first data type
444 first_enum = bases[-1] # enum type
445 else:
446 for base in bases[0].__mro__:
447 # most common: (IntEnum, int, Enum, object)
448 # possible: (<Enum 'AutoIntEnum'>, <Enum 'IntEnum'>,
449 # <class 'int'>, <Enum 'Enum'>,
450 # <class 'object'>)
451 if issubclass(base, Enum):
452 if first_enum is None:
453 first_enum = base
454 else:
455 if member_type is None:
456 member_type = base
457
458 return member_type, first_enum
459
460 @staticmethod
461 def _find_new_(classdict, member_type, first_enum):
462 """Returns the __new__ to be used for creating the enum members.
463
464 classdict: the class dictionary given to __new__
465 member_type: the data type whose __new__ will be used by default
466 first_enum: enumeration to check for an overriding __new__
467
468 """
469 # now find the correct __new__, checking to see of one was defined
470 # by the user; also check earlier enum classes in case a __new__ was
471 # saved as __new_member__
472 __new__ = classdict.get('__new__', None)
473
474 # should __new__ be saved as __new_member__ later?
475 save_new = __new__ is not None
476
477 if __new__ is None:
478 # check all possibles for __new_member__ before falling back to
479 # __new__
480 for method in ('__new_member__', '__new__'):
481 for possible in (member_type, first_enum):
482 target = getattr(possible, method, None)
483 if target not in {
484 None,
485 None.__new__,
486 object.__new__,
487 Enum.__new__,
488 }:
489 __new__ = target
490 break
491 if __new__ is not None:
492 break
493 else:
494 __new__ = object.__new__
495
496 # if a non-object.__new__ is used then whatever value/tuple was
497 # assigned to the enum member name will be passed to __new__ and to the
498 # new enum member's __init__
499 if __new__ is object.__new__:
500 use_args = False
501 else:
502 use_args = True
503
504 return __new__, save_new, use_args
505
506
507class Enum(metaclass=EnumMeta):
508 """Generic enumeration.
509
510 Derive from this class to define new enumerations.
511
512 """
513 def __new__(cls, value):
514 # all enum instances are actually created during class construction
515 # without calling this method; this method is called by the metaclass'
516 # __call__ (i.e. Color(3) ), and by pickle
517 if type(value) is cls:
Ethan Furman23bb6f42016-11-21 09:22:05 -0800518 # For lookups like Color(Color.RED)
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700519 return value
520 # by-value search for a matching enum member
521 # see if it's in the reverse mapping (for hashable values)
Ethan Furman2aa27322013-07-19 19:35:56 -0700522 try:
Ethan Furman520ad572013-07-19 19:47:21 -0700523 if value in cls._value2member_map_:
524 return cls._value2member_map_[value]
Ethan Furman2aa27322013-07-19 19:35:56 -0700525 except TypeError:
526 # not there, now do long search -- O(n) behavior
Ethan Furman520ad572013-07-19 19:47:21 -0700527 for member in cls._member_map_.values():
Ethan Furman0081f232014-09-16 17:31:23 -0700528 if member._value_ == value:
Ethan Furman2aa27322013-07-19 19:35:56 -0700529 return member
Ethan Furmanee47e5c2016-08-31 00:12:15 -0700530 # still not found -- try _missing_ hook
531 return cls._missing_(value)
532
Ethan Furmanc16595e2016-09-10 23:36:59 -0700533 def _generate_next_value_(name, start, count, last_values):
534 for last_value in reversed(last_values):
535 try:
536 return last_value + 1
537 except TypeError:
538 pass
539 else:
Ethan Furmanee47e5c2016-08-31 00:12:15 -0700540 return start
Ethan Furmanc16595e2016-09-10 23:36:59 -0700541
Ethan Furmanee47e5c2016-08-31 00:12:15 -0700542 @classmethod
543 def _missing_(cls, value):
Ethan Furman0081f232014-09-16 17:31:23 -0700544 raise ValueError("%r is not a valid %s" % (value, cls.__name__))
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700545
546 def __repr__(self):
547 return "<%s.%s: %r>" % (
Ethan Furman520ad572013-07-19 19:47:21 -0700548 self.__class__.__name__, self._name_, self._value_)
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700549
550 def __str__(self):
Ethan Furman520ad572013-07-19 19:47:21 -0700551 return "%s.%s" % (self.__class__.__name__, self._name_)
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700552
Ethan Furman388a3922013-08-12 06:51:41 -0700553 def __dir__(self):
Ethan Furman0ae550b2014-10-14 08:58:32 -0700554 added_behavior = [
555 m
556 for cls in self.__class__.mro()
557 for m in cls.__dict__
Ethan Furman354ecf12015-03-11 08:43:12 -0700558 if m[0] != '_' and m not in self._member_map_
Ethan Furman0ae550b2014-10-14 08:58:32 -0700559 ]
Ethan Furmanec5f8eb2014-10-21 13:40:35 -0700560 return (['__class__', '__doc__', '__module__'] + added_behavior)
Ethan Furman388a3922013-08-12 06:51:41 -0700561
Ethan Furmanec15a822013-08-31 19:17:41 -0700562 def __format__(self, format_spec):
563 # mixed-in Enums should use the mixed-in type's __format__, otherwise
564 # we can get strange results with the Enum name showing up instead of
565 # the value
566
567 # pure Enum branch
568 if self._member_type_ is object:
569 cls = str
570 val = str(self)
571 # mix-in branch
572 else:
573 cls = self._member_type_
Ethan Furman0081f232014-09-16 17:31:23 -0700574 val = self._value_
Ethan Furmanec15a822013-08-31 19:17:41 -0700575 return cls.__format__(val, format_spec)
576
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700577 def __hash__(self):
Ethan Furman520ad572013-07-19 19:47:21 -0700578 return hash(self._name_)
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700579
Ethan Furmanca1b7942014-02-08 11:36:27 -0800580 def __reduce_ex__(self, proto):
Ethan Furmandc870522014-02-18 12:37:12 -0800581 return self.__class__, (self._value_, )
Ethan Furmanca1b7942014-02-08 11:36:27 -0800582
Ethan Furman33918c12013-09-27 23:02:02 -0700583 # DynamicClassAttribute is used to provide access to the `name` and
584 # `value` properties of enum members while keeping some measure of
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700585 # protection from modification, while still allowing for an enumeration
586 # to have members named `name` and `value`. This works because enumeration
587 # members are not set directly on the enum class -- __getattr__ is
588 # used to look them up.
589
Ethan Furmane03ea372013-09-25 07:14:41 -0700590 @DynamicClassAttribute
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700591 def name(self):
Ethan Furmanc850f342013-09-15 16:59:35 -0700592 """The name of the Enum member."""
Ethan Furman520ad572013-07-19 19:47:21 -0700593 return self._name_
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700594
Ethan Furmane03ea372013-09-25 07:14:41 -0700595 @DynamicClassAttribute
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700596 def value(self):
Ethan Furmanc850f342013-09-15 16:59:35 -0700597 """The value of the Enum member."""
Ethan Furman520ad572013-07-19 19:47:21 -0700598 return self._value_
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700599
Ethan Furman24e837f2015-03-18 17:27:57 -0700600 @classmethod
601 def _convert(cls, name, module, filter, source=None):
602 """
603 Create a new Enum subclass that replaces a collection of global constants
604 """
605 # convert all constants from source (or module) that pass filter() to
606 # a new Enum called name, and export the enum and its members back to
607 # module;
608 # also, replace the __reduce_ex__ method so unpickling works in
609 # previous Python versions
610 module_globals = vars(sys.modules[module])
611 if source:
612 source = vars(source)
613 else:
614 source = module_globals
Gregory P. Smith ext:(%20%5BGoogle%20Inc.%5D)6f20bd62016-06-03 19:14:52 +0000615 # We use an OrderedDict of sorted source keys so that the
616 # _value2member_map is populated in the same order every time
617 # for a consistent reverse mapping of number to name when there
618 # are multiple names for the same number rather than varying
619 # between runs due to hash randomization of the module dictionary.
Ethan Furman06339e72016-09-11 13:25:26 -0700620 members = [
621 (name, source[name])
622 for name in source.keys()
623 if filter(name)]
624 try:
625 # sort by value
626 members.sort(key=lambda t: (t[1], t[0]))
627 except TypeError:
628 # unless some values aren't comparable, in which case sort by name
629 members.sort(key=lambda t: t[0])
Ethan Furman24e837f2015-03-18 17:27:57 -0700630 cls = cls(name, members, module=module)
631 cls.__reduce_ex__ = _reduce_ex_by_name
632 module_globals.update(cls.__members__)
633 module_globals[name] = cls
634 return cls
635
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700636
637class IntEnum(int, Enum):
638 """Enum where members are also (and must be) ints"""
Ethan Furmanf24bb352013-07-18 17:05:39 -0700639
640
Ethan Furman24e837f2015-03-18 17:27:57 -0700641def _reduce_ex_by_name(self, proto):
642 return self.name
643
Ethan Furman65a5a472016-09-01 23:55:19 -0700644class Flag(Enum):
Ethan Furmanee47e5c2016-08-31 00:12:15 -0700645 """Support for flags"""
Ethan Furmanc16595e2016-09-10 23:36:59 -0700646
647 def _generate_next_value_(name, start, count, last_values):
Ethan Furmanee47e5c2016-08-31 00:12:15 -0700648 """
649 Generate the next value when not given.
650
651 name: the name of the member
652 start: the initital start value or None
653 count: the number of existing members
654 last_value: the last value assigned or None
655 """
656 if not count:
657 return start if start is not None else 1
Ethan Furmanc16595e2016-09-10 23:36:59 -0700658 for last_value in reversed(last_values):
659 try:
660 high_bit = _high_bit(last_value)
661 break
Ethan Furman3515dcc2016-09-18 13:15:41 -0700662 except Exception:
Ethan Furmanc16595e2016-09-10 23:36:59 -0700663 raise TypeError('Invalid Flag value: %r' % last_value) from None
Ethan Furmanee47e5c2016-08-31 00:12:15 -0700664 return 2 ** (high_bit+1)
665
666 @classmethod
667 def _missing_(cls, value):
668 original_value = value
669 if value < 0:
670 value = ~value
671 possible_member = cls._create_pseudo_member_(value)
Ethan Furmanee47e5c2016-08-31 00:12:15 -0700672 if original_value < 0:
673 possible_member = ~possible_member
674 return possible_member
675
676 @classmethod
677 def _create_pseudo_member_(cls, value):
Ethan Furman3515dcc2016-09-18 13:15:41 -0700678 """
679 Create a composite member iff value contains only members.
680 """
Ethan Furmanee47e5c2016-08-31 00:12:15 -0700681 pseudo_member = cls._value2member_map_.get(value, None)
682 if pseudo_member is None:
Ethan Furman3515dcc2016-09-18 13:15:41 -0700683 # verify all bits are accounted for
684 _, extra_flags = _decompose(cls, value)
685 if extra_flags:
686 raise ValueError("%r is not a valid %s" % (value, cls.__name__))
687 # construct a singleton enum pseudo-member
Ethan Furmanee47e5c2016-08-31 00:12:15 -0700688 pseudo_member = object.__new__(cls)
689 pseudo_member._name_ = None
690 pseudo_member._value_ = value
Ethan Furman28cf6632017-01-24 12:12:06 -0800691 # use setdefault in case another thread already created a composite
692 # with this value
693 pseudo_member = cls._value2member_map_.setdefault(value, pseudo_member)
Ethan Furmanee47e5c2016-08-31 00:12:15 -0700694 return pseudo_member
695
Ethan Furmanee47e5c2016-08-31 00:12:15 -0700696 def __contains__(self, other):
697 if not isinstance(other, self.__class__):
698 return NotImplemented
699 return other._value_ & self._value_ == other._value_
700
Ethan Furmanee47e5c2016-08-31 00:12:15 -0700701 def __repr__(self):
702 cls = self.__class__
703 if self._name_ is not None:
704 return '<%s.%s: %r>' % (cls.__name__, self._name_, self._value_)
Ethan Furman3515dcc2016-09-18 13:15:41 -0700705 members, uncovered = _decompose(cls, self._value_)
Ethan Furman27682d22016-09-04 11:39:01 -0700706 return '<%s.%s: %r>' % (
707 cls.__name__,
708 '|'.join([str(m._name_ or m._value_) for m in members]),
709 self._value_,
710 )
Ethan Furmanee47e5c2016-08-31 00:12:15 -0700711
712 def __str__(self):
713 cls = self.__class__
714 if self._name_ is not None:
715 return '%s.%s' % (cls.__name__, self._name_)
Ethan Furman3515dcc2016-09-18 13:15:41 -0700716 members, uncovered = _decompose(cls, self._value_)
Ethan Furmanee47e5c2016-08-31 00:12:15 -0700717 if len(members) == 1 and members[0]._name_ is None:
718 return '%s.%r' % (cls.__name__, members[0]._value_)
719 else:
720 return '%s.%s' % (
721 cls.__name__,
722 '|'.join([str(m._name_ or m._value_) for m in members]),
723 )
724
Ethan Furman25d94bb2016-09-02 16:32:32 -0700725 def __bool__(self):
726 return bool(self._value_)
727
Ethan Furmanee47e5c2016-08-31 00:12:15 -0700728 def __or__(self, other):
729 if not isinstance(other, self.__class__):
730 return NotImplemented
731 return self.__class__(self._value_ | other._value_)
732
733 def __and__(self, other):
734 if not isinstance(other, self.__class__):
735 return NotImplemented
736 return self.__class__(self._value_ & other._value_)
737
738 def __xor__(self, other):
739 if not isinstance(other, self.__class__):
740 return NotImplemented
741 return self.__class__(self._value_ ^ other._value_)
742
743 def __invert__(self):
Ethan Furman3515dcc2016-09-18 13:15:41 -0700744 members, uncovered = _decompose(self.__class__, self._value_)
Serhiy Storchaka81108372017-09-26 00:55:55 +0300745 inverted = self.__class__(0)
746 for m in self.__class__:
747 if m not in members and not (m._value_ & self._value_):
748 inverted = inverted | m
Ethan Furmanee47e5c2016-08-31 00:12:15 -0700749 return self.__class__(inverted)
750
751
Ethan Furman65a5a472016-09-01 23:55:19 -0700752class IntFlag(int, Flag):
Ethan Furmanee47e5c2016-08-31 00:12:15 -0700753 """Support for integer-based Flags"""
754
755 @classmethod
Ethan Furman3515dcc2016-09-18 13:15:41 -0700756 def _missing_(cls, value):
757 if not isinstance(value, int):
758 raise ValueError("%r is not a valid %s" % (value, cls.__name__))
759 new_member = cls._create_pseudo_member_(value)
760 return new_member
761
762 @classmethod
Ethan Furmanee47e5c2016-08-31 00:12:15 -0700763 def _create_pseudo_member_(cls, value):
764 pseudo_member = cls._value2member_map_.get(value, None)
765 if pseudo_member is None:
Ethan Furman3515dcc2016-09-18 13:15:41 -0700766 need_to_create = [value]
767 # get unaccounted for bits
768 _, extra_flags = _decompose(cls, value)
769 # timer = 10
770 while extra_flags:
771 # timer -= 1
772 bit = _high_bit(extra_flags)
773 flag_value = 2 ** bit
774 if (flag_value not in cls._value2member_map_ and
775 flag_value not in need_to_create
776 ):
777 need_to_create.append(flag_value)
778 if extra_flags == -flag_value:
779 extra_flags = 0
780 else:
781 extra_flags ^= flag_value
782 for value in reversed(need_to_create):
783 # construct singleton pseudo-members
784 pseudo_member = int.__new__(cls, value)
785 pseudo_member._name_ = None
786 pseudo_member._value_ = value
Ethan Furman28cf6632017-01-24 12:12:06 -0800787 # use setdefault in case another thread already created a composite
788 # with this value
789 pseudo_member = cls._value2member_map_.setdefault(value, pseudo_member)
Ethan Furmanee47e5c2016-08-31 00:12:15 -0700790 return pseudo_member
791
Ethan Furmanee47e5c2016-08-31 00:12:15 -0700792 def __or__(self, other):
793 if not isinstance(other, (self.__class__, int)):
794 return NotImplemented
Ethan Furman3515dcc2016-09-18 13:15:41 -0700795 result = self.__class__(self._value_ | self.__class__(other)._value_)
796 return result
Ethan Furmanee47e5c2016-08-31 00:12:15 -0700797
798 def __and__(self, other):
799 if not isinstance(other, (self.__class__, int)):
800 return NotImplemented
801 return self.__class__(self._value_ & self.__class__(other)._value_)
802
803 def __xor__(self, other):
804 if not isinstance(other, (self.__class__, int)):
805 return NotImplemented
806 return self.__class__(self._value_ ^ self.__class__(other)._value_)
807
808 __ror__ = __or__
809 __rand__ = __and__
810 __rxor__ = __xor__
811
812 def __invert__(self):
Ethan Furman3515dcc2016-09-18 13:15:41 -0700813 result = self.__class__(~self._value_)
814 return result
Ethan Furmanee47e5c2016-08-31 00:12:15 -0700815
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"""
Ethan Furman3515dcc2016-09-18 13:15:41 -0700819 return value.bit_length() - 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
Ethan Furman3515dcc2016-09-18 13:15:41 -0700833
834def _decompose(flag, value):
835 """Extract all members from the value."""
836 # _decompose is only called if the value is not named
837 not_covered = value
838 negative = value < 0
Ethan Furman28cf6632017-01-24 12:12:06 -0800839 # issue29167: wrap accesses to _value2member_map_ in a list to avoid race
Ville Skyttä49b27342017-08-03 09:00:59 +0300840 # conditions between iterating over it and having more pseudo-
Ethan Furman28cf6632017-01-24 12:12:06 -0800841 # members added to it
Ethan Furman3515dcc2016-09-18 13:15:41 -0700842 if negative:
843 # only check for named flags
844 flags_to_check = [
845 (m, v)
Ethan Furman28cf6632017-01-24 12:12:06 -0800846 for v, m in list(flag._value2member_map_.items())
Ethan Furman3515dcc2016-09-18 13:15:41 -0700847 if m.name is not None
848 ]
849 else:
850 # check for named flags and powers-of-two flags
851 flags_to_check = [
852 (m, v)
Ethan Furman28cf6632017-01-24 12:12:06 -0800853 for v, m in list(flag._value2member_map_.items())
Ethan Furman3515dcc2016-09-18 13:15:41 -0700854 if m.name is not None or _power_of_two(v)
855 ]
856 members = []
857 for member, member_value in flags_to_check:
858 if member_value and member_value & value == member_value:
859 members.append(member)
860 not_covered &= ~member_value
861 if not members and value in flag._value2member_map_:
862 members.append(flag._value2member_map_[value])
863 members.sort(key=lambda m: m._value_, reverse=True)
864 if len(members) > 1 and members[0].value == value:
865 # we have the breakdown, don't need the value member itself
866 members.pop(0)
867 return members, not_covered
868
869def _power_of_two(value):
870 if value < 1:
871 return False
872 return value == 2 ** _high_bit(value)