blob: 04d8ec1fa872f881f373afcbff374d14aa23d04f [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 Furmane5754ab2015-09-17 22:03:52 -07004
Ethan Furmanc16595e2016-09-10 23:36:59 -07005__all__ = [
6 'EnumMeta',
7 'Enum', 'IntEnum', 'Flag', 'IntFlag',
8 'auto', 'unique',
9 ]
Ethan Furman6b3d64a2013-06-14 16:55:46 -070010
11
Ethan Furman101e0742013-09-15 12:34:36 -070012def _is_descriptor(obj):
13 """Returns True if obj is a descriptor, False otherwise."""
14 return (
15 hasattr(obj, '__get__') or
16 hasattr(obj, '__set__') or
17 hasattr(obj, '__delete__'))
18
19
Ethan Furman6b3d64a2013-06-14 16:55:46 -070020def _is_dunder(name):
21 """Returns True if a __dunder__ name, False otherwise."""
22 return (name[:2] == name[-2:] == '__' and
23 name[2:3] != '_' and
Ethan Furman648f8602013-10-06 17:19:54 -070024 name[-3:-2] != '_' and
25 len(name) > 4)
Ethan Furman6b3d64a2013-06-14 16:55:46 -070026
27
28def _is_sunder(name):
29 """Returns True if a _sunder_ name, False otherwise."""
30 return (name[0] == name[-1] == '_' and
31 name[1:2] != '_' and
Ethan Furman648f8602013-10-06 17:19:54 -070032 name[-2:-1] != '_' and
33 len(name) > 2)
Ethan Furman6b3d64a2013-06-14 16:55:46 -070034
Ethan Furman6b3d64a2013-06-14 16:55:46 -070035def _make_class_unpicklable(cls):
36 """Make the given class un-picklable."""
Ethan Furmanca1b7942014-02-08 11:36:27 -080037 def _break_on_call_reduce(self, proto):
Ethan Furman6b3d64a2013-06-14 16:55:46 -070038 raise TypeError('%r cannot be pickled' % self)
Ethan Furmanca1b7942014-02-08 11:36:27 -080039 cls.__reduce_ex__ = _break_on_call_reduce
Ethan Furman6b3d64a2013-06-14 16:55:46 -070040 cls.__module__ = '<unknown>'
41
Ethan Furman3515dcc2016-09-18 13:15:41 -070042_auto_null = object()
Ethan Furmanc16595e2016-09-10 23:36:59 -070043class auto:
44 """
45 Instances are replaced with an appropriate value in Enum class suites.
46 """
Ethan Furman3515dcc2016-09-18 13:15:41 -070047 value = _auto_null
Ethan Furmanc16595e2016-09-10 23:36:59 -070048
Ethan Furman101e0742013-09-15 12:34:36 -070049
Ethan Furman6b3d64a2013-06-14 16:55:46 -070050class _EnumDict(dict):
Ethan Furman101e0742013-09-15 12:34:36 -070051 """Track enum member order and ensure member names are not reused.
Ethan Furman6b3d64a2013-06-14 16:55:46 -070052
53 EnumMeta will use the names found in self._member_names as the
54 enumeration member names.
55
56 """
57 def __init__(self):
58 super().__init__()
59 self._member_names = []
Ethan Furmanc16595e2016-09-10 23:36:59 -070060 self._last_values = []
Ethan Furmana4b1bb42018-01-22 07:56:37 -080061 self._ignore = []
Ethan Furman6b3d64a2013-06-14 16:55:46 -070062
63 def __setitem__(self, key, value):
Ethan Furman101e0742013-09-15 12:34:36 -070064 """Changes anything not dundered or not a descriptor.
Ethan Furman6b3d64a2013-06-14 16:55:46 -070065
66 If an enum member name is used twice, an error is raised; duplicate
67 values are not checked for.
68
69 Single underscore (sunder) names are reserved.
70
71 """
72 if _is_sunder(key):
Ethan Furmanee47e5c2016-08-31 00:12:15 -070073 if key not in (
Ethan Furman3515dcc2016-09-18 13:15:41 -070074 '_order_', '_create_pseudo_member_',
Ethan Furmana4b1bb42018-01-22 07:56:37 -080075 '_generate_next_value_', '_missing_', '_ignore_',
Ethan Furmanee47e5c2016-08-31 00:12:15 -070076 ):
Ethan Furmane8e61272016-08-20 07:19:31 -070077 raise ValueError('_names_ are reserved for future Enum use')
Ethan Furmanc16595e2016-09-10 23:36:59 -070078 if key == '_generate_next_value_':
79 setattr(self, '_generate_next_value', value)
Ethan Furmana4b1bb42018-01-22 07:56:37 -080080 elif key == '_ignore_':
81 if isinstance(value, str):
82 value = value.replace(',',' ').split()
83 else:
84 value = list(value)
85 self._ignore = value
86 already = set(value) & set(self._member_names)
87 if already:
88 raise ValueError('_ignore_ cannot specify already set names: %r' % (already, ))
Ethan Furman101e0742013-09-15 12:34:36 -070089 elif _is_dunder(key):
Ethan Furmane8e61272016-08-20 07:19:31 -070090 if key == '__order__':
91 key = '_order_'
Ethan Furman101e0742013-09-15 12:34:36 -070092 elif key in self._member_names:
93 # descriptor overwriting an enum?
94 raise TypeError('Attempted to reuse key: %r' % key)
Ethan Furmana4b1bb42018-01-22 07:56:37 -080095 elif key in self._ignore:
96 pass
Ethan Furman101e0742013-09-15 12:34:36 -070097 elif not _is_descriptor(value):
98 if key in self:
99 # enum overwriting a descriptor?
Ethan Furmanee47e5c2016-08-31 00:12:15 -0700100 raise TypeError('%r already defined as: %r' % (key, self[key]))
Ethan Furmanc16595e2016-09-10 23:36:59 -0700101 if isinstance(value, auto):
Ethan Furman3515dcc2016-09-18 13:15:41 -0700102 if value.value == _auto_null:
103 value.value = self._generate_next_value(key, 1, len(self._member_names), self._last_values[:])
104 value = value.value
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700105 self._member_names.append(key)
Ethan Furmanc16595e2016-09-10 23:36:59 -0700106 self._last_values.append(value)
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700107 super().__setitem__(key, value)
108
109
Ezio Melotti9a3777e2013-08-17 15:53:55 +0300110# Dummy value for Enum as EnumMeta explicitly checks for it, but of course
111# until EnumMeta finishes running the first time the Enum class doesn't exist.
112# This is also why there are checks in EnumMeta like `if Enum is not None`
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700113Enum = None
114
Ethan Furman332dbc72016-08-20 00:00:52 -0700115
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700116class EnumMeta(type):
117 """Metaclass for Enum"""
118 @classmethod
Ethan Furman332dbc72016-08-20 00:00:52 -0700119 def __prepare__(metacls, cls, bases):
Ethan Furmanee47e5c2016-08-31 00:12:15 -0700120 # create the namespace dict
121 enum_dict = _EnumDict()
122 # inherit previous flags and _generate_next_value_ function
123 member_type, first_enum = metacls._get_mixins_(bases)
124 if first_enum is not None:
125 enum_dict['_generate_next_value_'] = getattr(first_enum, '_generate_next_value_', None)
126 return enum_dict
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700127
Ethan Furman65a5a472016-09-01 23:55:19 -0700128 def __new__(metacls, cls, bases, classdict):
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700129 # an Enum class is final once enumeration items have been defined; it
130 # cannot be mixed with other types (int, float, etc.) if it has an
131 # inherited __new__ unless a new __new__ is defined (or the resulting
132 # class will fail).
Ethan Furmana4b1bb42018-01-22 07:56:37 -0800133 #
134 # remove any keys listed in _ignore_
135 classdict.setdefault('_ignore_', []).append('_ignore_')
136 ignore = classdict['_ignore_']
137 for key in ignore:
138 classdict.pop(key, None)
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700139 member_type, first_enum = metacls._get_mixins_(bases)
140 __new__, save_new, use_args = metacls._find_new_(classdict, member_type,
141 first_enum)
142
143 # save enum items into separate mapping so they don't get baked into
144 # the new class
Ethan Furmanee47e5c2016-08-31 00:12:15 -0700145 enum_members = {k: classdict[k] for k in classdict._member_names}
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700146 for name in classdict._member_names:
147 del classdict[name]
148
Ethan Furmane8e61272016-08-20 07:19:31 -0700149 # adjust the sunders
150 _order_ = classdict.pop('_order_', None)
151
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700152 # check for illegal enum names (any others?)
Ethan Furmanee47e5c2016-08-31 00:12:15 -0700153 invalid_names = set(enum_members) & {'mro', }
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700154 if invalid_names:
155 raise ValueError('Invalid enum member name: {0}'.format(
156 ','.join(invalid_names)))
157
Ethan Furman48a724f2015-04-11 23:23:06 -0700158 # create a default docstring if one has not been provided
159 if '__doc__' not in classdict:
160 classdict['__doc__'] = 'An enumeration.'
161
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700162 # create our new Enum type
163 enum_class = super().__new__(metacls, cls, bases, classdict)
Ethan Furman520ad572013-07-19 19:47:21 -0700164 enum_class._member_names_ = [] # names in definition order
INADA Naokie57f91a2018-06-19 01:14:26 +0900165 enum_class._member_map_ = {} # name->value map
Ethan Furman5e5a8232013-08-04 08:42:23 -0700166 enum_class._member_type_ = member_type
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700167
Ethan Furman354ecf12015-03-11 08:43:12 -0700168 # save attributes from super classes so we know if we can take
169 # the shortcut of storing members in the class dict
Ethan Furman3803ad42016-05-01 10:03:53 -0700170 base_attributes = {a for b in enum_class.mro() for a in b.__dict__}
Ethan Furman354ecf12015-03-11 08:43:12 -0700171
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700172 # Reverse value->name map for hashable values.
Ethan Furman520ad572013-07-19 19:47:21 -0700173 enum_class._value2member_map_ = {}
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700174
Ethan Furman2da95042014-03-03 12:42:52 -0800175 # If a custom type is mixed into the Enum, and it does not know how
176 # to pickle itself, pickle.dumps will succeed but pickle.loads will
177 # fail. Rather than have the error show up later and possibly far
178 # from the source, sabotage the pickle protocol for this class so
179 # that pickle.dumps also fails.
180 #
181 # However, if the new class implements its own __reduce_ex__, do not
182 # sabotage -- it's on them to make sure it works correctly. We use
183 # __reduce_ex__ instead of any of the others as it is preferred by
184 # pickle over __reduce__, and it handles all pickle protocols.
185 if '__reduce_ex__' not in classdict:
Ethan Furmandc870522014-02-18 12:37:12 -0800186 if member_type is not object:
187 methods = ('__getnewargs_ex__', '__getnewargs__',
188 '__reduce_ex__', '__reduce__')
Ethan Furman2da95042014-03-03 12:42:52 -0800189 if not any(m in member_type.__dict__ for m in methods):
Ethan Furmandc870522014-02-18 12:37:12 -0800190 _make_class_unpicklable(enum_class)
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700191
192 # instantiate them, checking for duplicates as we go
193 # we instantiate first instead of checking for duplicates first in case
194 # a custom __new__ is doing something funky with the values -- such as
195 # auto-numbering ;)
196 for member_name in classdict._member_names:
Ethan Furmanee47e5c2016-08-31 00:12:15 -0700197 value = enum_members[member_name]
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700198 if not isinstance(value, tuple):
199 args = (value, )
200 else:
201 args = value
202 if member_type is tuple: # special case for tuple enums
203 args = (args, ) # wrap it one more time
204 if not use_args:
205 enum_member = __new__(enum_class)
Ethan Furmanb41803e2013-07-25 13:50:45 -0700206 if not hasattr(enum_member, '_value_'):
207 enum_member._value_ = value
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700208 else:
209 enum_member = __new__(enum_class, *args)
Ethan Furmanb41803e2013-07-25 13:50:45 -0700210 if not hasattr(enum_member, '_value_'):
Ethan Furmanee47e5c2016-08-31 00:12:15 -0700211 if member_type is object:
212 enum_member._value_ = value
213 else:
214 enum_member._value_ = member_type(*args)
Ethan Furman520ad572013-07-19 19:47:21 -0700215 value = enum_member._value_
Ethan Furman520ad572013-07-19 19:47:21 -0700216 enum_member._name_ = member_name
Ethan Furmanc850f342013-09-15 16:59:35 -0700217 enum_member.__objclass__ = enum_class
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700218 enum_member.__init__(*args)
219 # If another member with the same value was already defined, the
220 # new member becomes an alias to the existing one.
Ethan Furman520ad572013-07-19 19:47:21 -0700221 for name, canonical_member in enum_class._member_map_.items():
Ethan Furman0081f232014-09-16 17:31:23 -0700222 if canonical_member._value_ == enum_member._value_:
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700223 enum_member = canonical_member
224 break
225 else:
226 # Aliases don't appear in member names (only in __members__).
Ethan Furman520ad572013-07-19 19:47:21 -0700227 enum_class._member_names_.append(member_name)
Ethan Furman354ecf12015-03-11 08:43:12 -0700228 # performance boost for any member that would not shadow
229 # a DynamicClassAttribute
230 if member_name not in base_attributes:
231 setattr(enum_class, member_name, enum_member)
232 # now add to _member_map_
Ethan Furman520ad572013-07-19 19:47:21 -0700233 enum_class._member_map_[member_name] = enum_member
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700234 try:
235 # This may fail if value is not hashable. We can't add the value
236 # to the map, and by-value lookups for this value will be
237 # linear.
Ethan Furman520ad572013-07-19 19:47:21 -0700238 enum_class._value2member_map_[value] = enum_member
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700239 except TypeError:
240 pass
241
242 # double check that repr and friends are not the mixin's or various
243 # things break (such as pickle)
Ethan Furmandc870522014-02-18 12:37:12 -0800244 for name in ('__repr__', '__str__', '__format__', '__reduce_ex__'):
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700245 class_method = getattr(enum_class, name)
246 obj_method = getattr(member_type, name, None)
247 enum_method = getattr(first_enum, name, None)
248 if obj_method is not None and obj_method is class_method:
249 setattr(enum_class, name, enum_method)
250
251 # replace any other __new__ with our own (as long as Enum is not None,
252 # anyway) -- again, this is to support pickle
253 if Enum is not None:
254 # if the user defined their own __new__, save it before it gets
255 # clobbered in case they subclass later
256 if save_new:
257 enum_class.__new_member__ = __new__
258 enum_class.__new__ = Enum.__new__
Ethan Furmane8e61272016-08-20 07:19:31 -0700259
260 # py3 support for definition order (helps keep py2/py3 code in sync)
261 if _order_ is not None:
262 if isinstance(_order_, str):
263 _order_ = _order_.replace(',', ' ').split()
264 if _order_ != enum_class._member_names_:
265 raise TypeError('member order does not match _order_')
266
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700267 return enum_class
268
Ethan Furman5de67b12016-04-13 23:52:09 -0700269 def __bool__(self):
270 """
271 classes/types should always be True.
272 """
273 return True
274
Ethan Furmand9925a12014-09-16 20:35:55 -0700275 def __call__(cls, value, names=None, *, module=None, qualname=None, type=None, start=1):
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700276 """Either returns an existing member, or creates a new enum class.
277
278 This method is used both when an enum class is given a value to match
279 to an enumeration member (i.e. Color(3)) and for the functional API
Ethan Furman23bb6f42016-11-21 09:22:05 -0800280 (i.e. Color = Enum('Color', names='RED GREEN BLUE')).
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700281
Ethan Furman2da95042014-03-03 12:42:52 -0800282 When used for the functional API:
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700283
Ethan Furman2da95042014-03-03 12:42:52 -0800284 `value` will be the name of the new class.
285
286 `names` should be either a string of white-space/comma delimited names
Ethan Furmand9925a12014-09-16 20:35:55 -0700287 (values will start at `start`), or an iterator/mapping of name, value pairs.
Ethan Furman2da95042014-03-03 12:42:52 -0800288
289 `module` should be set to the module this class is being created in;
290 if it is not set, an attempt to find that module will be made, but if
291 it fails the class will not be picklable.
292
293 `qualname` should be set to the actual location this class can be found
294 at in its module; by default it is set to the global scope. If this is
295 not correct, unpickling will fail in some circumstances.
296
297 `type`, if set, will be mixed in as the first base class.
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700298
299 """
300 if names is None: # simple value lookup
301 return cls.__new__(cls, value)
302 # otherwise, functional API: we're creating a new Enum type
Ethan Furmand9925a12014-09-16 20:35:55 -0700303 return cls._create_(value, names, module=module, qualname=qualname, type=type, start=start)
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700304
305 def __contains__(cls, member):
Ethan Furman0081f232014-09-16 17:31:23 -0700306 return isinstance(member, cls) and member._name_ in cls._member_map_
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700307
Ethan Furman64a99722013-09-22 16:18:19 -0700308 def __delattr__(cls, attr):
309 # nicer error message when someone tries to delete an attribute
310 # (see issue19025).
311 if attr in cls._member_map_:
312 raise AttributeError(
313 "%s: cannot delete Enum member." % cls.__name__)
314 super().__delattr__(attr)
315
Ethan Furman388a3922013-08-12 06:51:41 -0700316 def __dir__(self):
Ethan Furman64a99722013-09-22 16:18:19 -0700317 return (['__class__', '__doc__', '__members__', '__module__'] +
318 self._member_names_)
Ethan Furman388a3922013-08-12 06:51:41 -0700319
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700320 def __getattr__(cls, name):
321 """Return the enum member matching `name`
322
323 We use __getattr__ instead of descriptors or inserting into the enum
324 class' __dict__ in order to support `name` and `value` being both
325 properties for enum members (which live in the class' __dict__) and
326 enum members themselves.
327
328 """
329 if _is_dunder(name):
330 raise AttributeError(name)
331 try:
Ethan Furman520ad572013-07-19 19:47:21 -0700332 return cls._member_map_[name]
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700333 except KeyError:
334 raise AttributeError(name) from None
335
336 def __getitem__(cls, name):
Ethan Furman520ad572013-07-19 19:47:21 -0700337 return cls._member_map_[name]
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700338
339 def __iter__(cls):
Ethan Furman520ad572013-07-19 19:47:21 -0700340 return (cls._member_map_[name] for name in cls._member_names_)
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700341
342 def __len__(cls):
Ethan Furman520ad572013-07-19 19:47:21 -0700343 return len(cls._member_names_)
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700344
Ethan Furman2131a4a2013-09-14 18:11:24 -0700345 @property
346 def __members__(cls):
347 """Returns a mapping of member name->value.
348
349 This mapping lists all enum members, including aliases. Note that this
350 is a read-only view of the internal mapping.
351
352 """
353 return MappingProxyType(cls._member_map_)
354
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700355 def __repr__(cls):
356 return "<enum %r>" % cls.__name__
357
Ethan Furman2131a4a2013-09-14 18:11:24 -0700358 def __reversed__(cls):
359 return (cls._member_map_[name] for name in reversed(cls._member_names_))
360
Ethan Furmanf203f2d2013-09-06 07:16:48 -0700361 def __setattr__(cls, name, value):
362 """Block attempts to reassign Enum members.
363
364 A simple assignment to the class namespace only changes one of the
365 several possible ways to get an Enum member from the Enum class,
366 resulting in an inconsistent Enumeration.
367
368 """
369 member_map = cls.__dict__.get('_member_map_', {})
370 if name in member_map:
371 raise AttributeError('Cannot reassign members.')
372 super().__setattr__(name, value)
373
anentropicb8e21f12018-04-16 04:40:35 +0100374 def _create_(cls, class_name, names, *, module=None, qualname=None, type=None, start=1):
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700375 """Convenience method to create a new Enum class.
376
377 `names` can be:
378
379 * A string containing member names, separated either with spaces or
Ethan Furmand9925a12014-09-16 20:35:55 -0700380 commas. Values are incremented by 1 from `start`.
381 * An iterable of member names. Values are incremented by 1 from `start`.
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700382 * An iterable of (member name, value) pairs.
Ethan Furmand9925a12014-09-16 20:35:55 -0700383 * A mapping of member name -> value pairs.
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700384
385 """
386 metacls = cls.__class__
387 bases = (cls, ) if type is None else (type, cls)
Ethan Furmanee47e5c2016-08-31 00:12:15 -0700388 _, first_enum = cls._get_mixins_(bases)
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700389 classdict = metacls.__prepare__(class_name, bases)
390
391 # special processing needed for names?
392 if isinstance(names, str):
393 names = names.replace(',', ' ').split()
Dong-hee Nadcc8ce42017-06-22 01:52:32 +0900394 if isinstance(names, (tuple, list)) and names and isinstance(names[0], str):
Ethan Furmanee47e5c2016-08-31 00:12:15 -0700395 original_names, names = names, []
Ethan Furmanc16595e2016-09-10 23:36:59 -0700396 last_values = []
Ethan Furmanee47e5c2016-08-31 00:12:15 -0700397 for count, name in enumerate(original_names):
Ethan Furmanc16595e2016-09-10 23:36:59 -0700398 value = first_enum._generate_next_value_(name, start, count, last_values[:])
399 last_values.append(value)
400 names.append((name, value))
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700401
402 # Here, names is either an iterable of (name, value) or a mapping.
403 for item in names:
404 if isinstance(item, str):
405 member_name, member_value = item, names[item]
406 else:
407 member_name, member_value = item
408 classdict[member_name] = member_value
409 enum_class = metacls.__new__(metacls, class_name, bases, classdict)
410
411 # TODO: replace the frame hack if a blessed way to know the calling
412 # module is ever developed
413 if module is None:
414 try:
415 module = sys._getframe(2).f_globals['__name__']
416 except (AttributeError, ValueError) as exc:
417 pass
418 if module is None:
419 _make_class_unpicklable(enum_class)
420 else:
421 enum_class.__module__ = module
Ethan Furmanca1b7942014-02-08 11:36:27 -0800422 if qualname is not None:
423 enum_class.__qualname__ = qualname
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700424
425 return enum_class
426
427 @staticmethod
428 def _get_mixins_(bases):
429 """Returns the type for creating enum members, and the first inherited
430 enum class.
431
432 bases: the tuple of bases that was given to __new__
433
434 """
435 if not bases:
436 return object, Enum
437
438 # double check that we are not subclassing a class with existing
439 # enumeration members; while we're at it, see if any other data
440 # type has been mixed in so we can use the correct __new__
441 member_type = first_enum = None
442 for base in bases:
443 if (base is not Enum and
444 issubclass(base, Enum) and
Ethan Furman520ad572013-07-19 19:47:21 -0700445 base._member_names_):
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700446 raise TypeError("Cannot extend enumerations")
447 # base is now the last base in bases
448 if not issubclass(base, Enum):
449 raise TypeError("new enumerations must be created as "
450 "`ClassName([mixin_type,] enum_type)`")
451
452 # get correct mix-in type (either mix-in type of Enum subclass, or
453 # first base if last base is Enum)
454 if not issubclass(bases[0], Enum):
455 member_type = bases[0] # first data type
456 first_enum = bases[-1] # enum type
457 else:
458 for base in bases[0].__mro__:
459 # most common: (IntEnum, int, Enum, object)
460 # possible: (<Enum 'AutoIntEnum'>, <Enum 'IntEnum'>,
461 # <class 'int'>, <Enum 'Enum'>,
462 # <class 'object'>)
463 if issubclass(base, Enum):
464 if first_enum is None:
465 first_enum = base
466 else:
467 if member_type is None:
468 member_type = base
469
470 return member_type, first_enum
471
472 @staticmethod
473 def _find_new_(classdict, member_type, first_enum):
474 """Returns the __new__ to be used for creating the enum members.
475
476 classdict: the class dictionary given to __new__
477 member_type: the data type whose __new__ will be used by default
478 first_enum: enumeration to check for an overriding __new__
479
480 """
481 # now find the correct __new__, checking to see of one was defined
482 # by the user; also check earlier enum classes in case a __new__ was
483 # saved as __new_member__
484 __new__ = classdict.get('__new__', None)
485
486 # should __new__ be saved as __new_member__ later?
487 save_new = __new__ is not None
488
489 if __new__ is None:
490 # check all possibles for __new_member__ before falling back to
491 # __new__
492 for method in ('__new_member__', '__new__'):
493 for possible in (member_type, first_enum):
494 target = getattr(possible, method, None)
495 if target not in {
496 None,
497 None.__new__,
498 object.__new__,
499 Enum.__new__,
500 }:
501 __new__ = target
502 break
503 if __new__ is not None:
504 break
505 else:
506 __new__ = object.__new__
507
508 # if a non-object.__new__ is used then whatever value/tuple was
509 # assigned to the enum member name will be passed to __new__ and to the
510 # new enum member's __init__
511 if __new__ is object.__new__:
512 use_args = False
513 else:
514 use_args = True
515
516 return __new__, save_new, use_args
517
518
519class Enum(metaclass=EnumMeta):
520 """Generic enumeration.
521
522 Derive from this class to define new enumerations.
523
524 """
525 def __new__(cls, value):
526 # all enum instances are actually created during class construction
527 # without calling this method; this method is called by the metaclass'
528 # __call__ (i.e. Color(3) ), and by pickle
529 if type(value) is cls:
Ethan Furman23bb6f42016-11-21 09:22:05 -0800530 # For lookups like Color(Color.RED)
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700531 return value
532 # by-value search for a matching enum member
533 # see if it's in the reverse mapping (for hashable values)
Ethan Furman2aa27322013-07-19 19:35:56 -0700534 try:
Ethan Furman520ad572013-07-19 19:47:21 -0700535 if value in cls._value2member_map_:
536 return cls._value2member_map_[value]
Ethan Furman2aa27322013-07-19 19:35:56 -0700537 except TypeError:
538 # not there, now do long search -- O(n) behavior
Ethan Furman520ad572013-07-19 19:47:21 -0700539 for member in cls._member_map_.values():
Ethan Furman0081f232014-09-16 17:31:23 -0700540 if member._value_ == value:
Ethan Furman2aa27322013-07-19 19:35:56 -0700541 return member
Ethan Furmanee47e5c2016-08-31 00:12:15 -0700542 # still not found -- try _missing_ hook
543 return cls._missing_(value)
544
Ethan Furmanc16595e2016-09-10 23:36:59 -0700545 def _generate_next_value_(name, start, count, last_values):
546 for last_value in reversed(last_values):
547 try:
548 return last_value + 1
549 except TypeError:
550 pass
551 else:
Ethan Furmanee47e5c2016-08-31 00:12:15 -0700552 return start
Ethan Furmanc16595e2016-09-10 23:36:59 -0700553
Ethan Furmanee47e5c2016-08-31 00:12:15 -0700554 @classmethod
555 def _missing_(cls, value):
Ethan Furman0081f232014-09-16 17:31:23 -0700556 raise ValueError("%r is not a valid %s" % (value, cls.__name__))
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700557
558 def __repr__(self):
559 return "<%s.%s: %r>" % (
Ethan Furman520ad572013-07-19 19:47:21 -0700560 self.__class__.__name__, self._name_, self._value_)
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700561
562 def __str__(self):
Ethan Furman520ad572013-07-19 19:47:21 -0700563 return "%s.%s" % (self.__class__.__name__, self._name_)
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700564
Ethan Furman388a3922013-08-12 06:51:41 -0700565 def __dir__(self):
Ethan Furman0ae550b2014-10-14 08:58:32 -0700566 added_behavior = [
567 m
568 for cls in self.__class__.mro()
569 for m in cls.__dict__
Ethan Furman354ecf12015-03-11 08:43:12 -0700570 if m[0] != '_' and m not in self._member_map_
Ethan Furman0ae550b2014-10-14 08:58:32 -0700571 ]
Ethan Furmanec5f8eb2014-10-21 13:40:35 -0700572 return (['__class__', '__doc__', '__module__'] + added_behavior)
Ethan Furman388a3922013-08-12 06:51:41 -0700573
Ethan Furmanec15a822013-08-31 19:17:41 -0700574 def __format__(self, format_spec):
575 # mixed-in Enums should use the mixed-in type's __format__, otherwise
576 # we can get strange results with the Enum name showing up instead of
577 # the value
578
579 # pure Enum branch
580 if self._member_type_ is object:
581 cls = str
582 val = str(self)
583 # mix-in branch
584 else:
585 cls = self._member_type_
Ethan Furman0081f232014-09-16 17:31:23 -0700586 val = self._value_
Ethan Furmanec15a822013-08-31 19:17:41 -0700587 return cls.__format__(val, format_spec)
588
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700589 def __hash__(self):
Ethan Furman520ad572013-07-19 19:47:21 -0700590 return hash(self._name_)
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700591
Ethan Furmanca1b7942014-02-08 11:36:27 -0800592 def __reduce_ex__(self, proto):
Ethan Furmandc870522014-02-18 12:37:12 -0800593 return self.__class__, (self._value_, )
Ethan Furmanca1b7942014-02-08 11:36:27 -0800594
Ethan Furman33918c12013-09-27 23:02:02 -0700595 # DynamicClassAttribute is used to provide access to the `name` and
596 # `value` properties of enum members while keeping some measure of
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700597 # protection from modification, while still allowing for an enumeration
598 # to have members named `name` and `value`. This works because enumeration
599 # members are not set directly on the enum class -- __getattr__ is
600 # used to look them up.
601
Ethan Furmane03ea372013-09-25 07:14:41 -0700602 @DynamicClassAttribute
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700603 def name(self):
Ethan Furmanc850f342013-09-15 16:59:35 -0700604 """The name of the Enum member."""
Ethan Furman520ad572013-07-19 19:47:21 -0700605 return self._name_
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700606
Ethan Furmane03ea372013-09-25 07:14:41 -0700607 @DynamicClassAttribute
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700608 def value(self):
Ethan Furmanc850f342013-09-15 16:59:35 -0700609 """The value of the Enum member."""
Ethan Furman520ad572013-07-19 19:47:21 -0700610 return self._value_
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700611
Ethan Furman24e837f2015-03-18 17:27:57 -0700612 @classmethod
613 def _convert(cls, name, module, filter, source=None):
614 """
615 Create a new Enum subclass that replaces a collection of global constants
616 """
617 # convert all constants from source (or module) that pass filter() to
618 # a new Enum called name, and export the enum and its members back to
619 # module;
620 # also, replace the __reduce_ex__ method so unpickling works in
621 # previous Python versions
622 module_globals = vars(sys.modules[module])
623 if source:
624 source = vars(source)
625 else:
626 source = module_globals
INADA Naokie57f91a2018-06-19 01:14:26 +0900627 # _value2member_map_ is populated in the same order every time
Gregory P. Smith ext:(%20%5BGoogle%20Inc.%5D)6f20bd62016-06-03 19:14:52 +0000628 # for a consistent reverse mapping of number to name when there
INADA Naokie57f91a2018-06-19 01:14:26 +0900629 # are multiple names for the same number.
Ethan Furman06339e72016-09-11 13:25:26 -0700630 members = [
INADA Naokie57f91a2018-06-19 01:14:26 +0900631 (name, value)
632 for name, value in source.items()
Ethan Furman06339e72016-09-11 13:25:26 -0700633 if filter(name)]
634 try:
635 # sort by value
636 members.sort(key=lambda t: (t[1], t[0]))
637 except TypeError:
638 # unless some values aren't comparable, in which case sort by name
639 members.sort(key=lambda t: t[0])
Ethan Furman24e837f2015-03-18 17:27:57 -0700640 cls = cls(name, members, module=module)
641 cls.__reduce_ex__ = _reduce_ex_by_name
642 module_globals.update(cls.__members__)
643 module_globals[name] = cls
644 return cls
645
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700646
647class IntEnum(int, Enum):
648 """Enum where members are also (and must be) ints"""
Ethan Furmanf24bb352013-07-18 17:05:39 -0700649
650
Ethan Furman24e837f2015-03-18 17:27:57 -0700651def _reduce_ex_by_name(self, proto):
652 return self.name
653
Ethan Furman65a5a472016-09-01 23:55:19 -0700654class Flag(Enum):
Ethan Furmanee47e5c2016-08-31 00:12:15 -0700655 """Support for flags"""
Ethan Furmanc16595e2016-09-10 23:36:59 -0700656
657 def _generate_next_value_(name, start, count, last_values):
Ethan Furmanee47e5c2016-08-31 00:12:15 -0700658 """
659 Generate the next value when not given.
660
661 name: the name of the member
662 start: the initital start value or None
663 count: the number of existing members
664 last_value: the last value assigned or None
665 """
666 if not count:
667 return start if start is not None else 1
Ethan Furmanc16595e2016-09-10 23:36:59 -0700668 for last_value in reversed(last_values):
669 try:
670 high_bit = _high_bit(last_value)
671 break
Ethan Furman3515dcc2016-09-18 13:15:41 -0700672 except Exception:
Ethan Furmanc16595e2016-09-10 23:36:59 -0700673 raise TypeError('Invalid Flag value: %r' % last_value) from None
Ethan Furmanee47e5c2016-08-31 00:12:15 -0700674 return 2 ** (high_bit+1)
675
676 @classmethod
677 def _missing_(cls, value):
678 original_value = value
679 if value < 0:
680 value = ~value
681 possible_member = cls._create_pseudo_member_(value)
Ethan Furmanee47e5c2016-08-31 00:12:15 -0700682 if original_value < 0:
683 possible_member = ~possible_member
684 return possible_member
685
686 @classmethod
687 def _create_pseudo_member_(cls, value):
Ethan Furman3515dcc2016-09-18 13:15:41 -0700688 """
689 Create a composite member iff value contains only members.
690 """
Ethan Furmanee47e5c2016-08-31 00:12:15 -0700691 pseudo_member = cls._value2member_map_.get(value, None)
692 if pseudo_member is None:
Ethan Furman3515dcc2016-09-18 13:15:41 -0700693 # verify all bits are accounted for
694 _, extra_flags = _decompose(cls, value)
695 if extra_flags:
696 raise ValueError("%r is not a valid %s" % (value, cls.__name__))
697 # construct a singleton enum pseudo-member
Ethan Furmanee47e5c2016-08-31 00:12:15 -0700698 pseudo_member = object.__new__(cls)
699 pseudo_member._name_ = None
700 pseudo_member._value_ = value
Ethan Furman28cf6632017-01-24 12:12:06 -0800701 # use setdefault in case another thread already created a composite
702 # with this value
703 pseudo_member = cls._value2member_map_.setdefault(value, pseudo_member)
Ethan Furmanee47e5c2016-08-31 00:12:15 -0700704 return pseudo_member
705
Ethan Furmanee47e5c2016-08-31 00:12:15 -0700706 def __contains__(self, other):
707 if not isinstance(other, self.__class__):
708 return NotImplemented
709 return other._value_ & self._value_ == other._value_
710
Ethan Furmanee47e5c2016-08-31 00:12:15 -0700711 def __repr__(self):
712 cls = self.__class__
713 if self._name_ is not None:
714 return '<%s.%s: %r>' % (cls.__name__, self._name_, self._value_)
Ethan Furman3515dcc2016-09-18 13:15:41 -0700715 members, uncovered = _decompose(cls, self._value_)
Ethan Furman27682d22016-09-04 11:39:01 -0700716 return '<%s.%s: %r>' % (
717 cls.__name__,
718 '|'.join([str(m._name_ or m._value_) for m in members]),
719 self._value_,
720 )
Ethan Furmanee47e5c2016-08-31 00:12:15 -0700721
722 def __str__(self):
723 cls = self.__class__
724 if self._name_ is not None:
725 return '%s.%s' % (cls.__name__, self._name_)
Ethan Furman3515dcc2016-09-18 13:15:41 -0700726 members, uncovered = _decompose(cls, self._value_)
Ethan Furmanee47e5c2016-08-31 00:12:15 -0700727 if len(members) == 1 and members[0]._name_ is None:
728 return '%s.%r' % (cls.__name__, members[0]._value_)
729 else:
730 return '%s.%s' % (
731 cls.__name__,
732 '|'.join([str(m._name_ or m._value_) for m in members]),
733 )
734
Ethan Furman25d94bb2016-09-02 16:32:32 -0700735 def __bool__(self):
736 return bool(self._value_)
737
Ethan Furmanee47e5c2016-08-31 00:12:15 -0700738 def __or__(self, other):
739 if not isinstance(other, self.__class__):
740 return NotImplemented
741 return self.__class__(self._value_ | other._value_)
742
743 def __and__(self, other):
744 if not isinstance(other, self.__class__):
745 return NotImplemented
746 return self.__class__(self._value_ & other._value_)
747
748 def __xor__(self, other):
749 if not isinstance(other, self.__class__):
750 return NotImplemented
751 return self.__class__(self._value_ ^ other._value_)
752
753 def __invert__(self):
Ethan Furman3515dcc2016-09-18 13:15:41 -0700754 members, uncovered = _decompose(self.__class__, self._value_)
Serhiy Storchaka81108372017-09-26 00:55:55 +0300755 inverted = self.__class__(0)
756 for m in self.__class__:
757 if m not in members and not (m._value_ & self._value_):
758 inverted = inverted | m
Ethan Furmanee47e5c2016-08-31 00:12:15 -0700759 return self.__class__(inverted)
760
761
Ethan Furman65a5a472016-09-01 23:55:19 -0700762class IntFlag(int, Flag):
Ethan Furmanee47e5c2016-08-31 00:12:15 -0700763 """Support for integer-based Flags"""
764
765 @classmethod
Ethan Furman3515dcc2016-09-18 13:15:41 -0700766 def _missing_(cls, value):
767 if not isinstance(value, int):
768 raise ValueError("%r is not a valid %s" % (value, cls.__name__))
769 new_member = cls._create_pseudo_member_(value)
770 return new_member
771
772 @classmethod
Ethan Furmanee47e5c2016-08-31 00:12:15 -0700773 def _create_pseudo_member_(cls, value):
774 pseudo_member = cls._value2member_map_.get(value, None)
775 if pseudo_member is None:
Ethan Furman3515dcc2016-09-18 13:15:41 -0700776 need_to_create = [value]
777 # get unaccounted for bits
778 _, extra_flags = _decompose(cls, value)
779 # timer = 10
780 while extra_flags:
781 # timer -= 1
782 bit = _high_bit(extra_flags)
783 flag_value = 2 ** bit
784 if (flag_value not in cls._value2member_map_ and
785 flag_value not in need_to_create
786 ):
787 need_to_create.append(flag_value)
788 if extra_flags == -flag_value:
789 extra_flags = 0
790 else:
791 extra_flags ^= flag_value
792 for value in reversed(need_to_create):
793 # construct singleton pseudo-members
794 pseudo_member = int.__new__(cls, value)
795 pseudo_member._name_ = None
796 pseudo_member._value_ = value
Ethan Furman28cf6632017-01-24 12:12:06 -0800797 # use setdefault in case another thread already created a composite
798 # with this value
799 pseudo_member = cls._value2member_map_.setdefault(value, pseudo_member)
Ethan Furmanee47e5c2016-08-31 00:12:15 -0700800 return pseudo_member
801
Ethan Furmanee47e5c2016-08-31 00:12:15 -0700802 def __or__(self, other):
803 if not isinstance(other, (self.__class__, int)):
804 return NotImplemented
Ethan Furman3515dcc2016-09-18 13:15:41 -0700805 result = self.__class__(self._value_ | self.__class__(other)._value_)
806 return result
Ethan Furmanee47e5c2016-08-31 00:12:15 -0700807
808 def __and__(self, other):
809 if not isinstance(other, (self.__class__, int)):
810 return NotImplemented
811 return self.__class__(self._value_ & self.__class__(other)._value_)
812
813 def __xor__(self, other):
814 if not isinstance(other, (self.__class__, int)):
815 return NotImplemented
816 return self.__class__(self._value_ ^ self.__class__(other)._value_)
817
818 __ror__ = __or__
819 __rand__ = __and__
820 __rxor__ = __xor__
821
822 def __invert__(self):
Ethan Furman3515dcc2016-09-18 13:15:41 -0700823 result = self.__class__(~self._value_)
824 return result
Ethan Furmanee47e5c2016-08-31 00:12:15 -0700825
826
827def _high_bit(value):
Ethan Furman04439532016-09-02 15:50:21 -0700828 """returns index of highest bit, or -1 if value is zero or negative"""
Ethan Furman3515dcc2016-09-18 13:15:41 -0700829 return value.bit_length() - 1
Ethan Furmanee47e5c2016-08-31 00:12:15 -0700830
Ethan Furmanf24bb352013-07-18 17:05:39 -0700831def unique(enumeration):
832 """Class decorator for enumerations ensuring unique member values."""
833 duplicates = []
834 for name, member in enumeration.__members__.items():
835 if name != member.name:
836 duplicates.append((name, member.name))
837 if duplicates:
838 alias_details = ', '.join(
839 ["%s -> %s" % (alias, name) for (alias, name) in duplicates])
840 raise ValueError('duplicate values found in %r: %s' %
841 (enumeration, alias_details))
842 return enumeration
Ethan Furman3515dcc2016-09-18 13:15:41 -0700843
844def _decompose(flag, value):
845 """Extract all members from the value."""
846 # _decompose is only called if the value is not named
847 not_covered = value
848 negative = value < 0
Ethan Furman28cf6632017-01-24 12:12:06 -0800849 # issue29167: wrap accesses to _value2member_map_ in a list to avoid race
Ville Skyttä49b27342017-08-03 09:00:59 +0300850 # conditions between iterating over it and having more pseudo-
Ethan Furman28cf6632017-01-24 12:12:06 -0800851 # members added to it
Ethan Furman3515dcc2016-09-18 13:15:41 -0700852 if negative:
853 # only check for named 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
858 ]
859 else:
860 # check for named flags and powers-of-two flags
861 flags_to_check = [
862 (m, v)
Ethan Furman28cf6632017-01-24 12:12:06 -0800863 for v, m in list(flag._value2member_map_.items())
Ethan Furman3515dcc2016-09-18 13:15:41 -0700864 if m.name is not None or _power_of_two(v)
865 ]
866 members = []
867 for member, member_value in flags_to_check:
868 if member_value and member_value & value == member_value:
869 members.append(member)
870 not_covered &= ~member_value
871 if not members and value in flag._value2member_map_:
872 members.append(flag._value2member_map_[value])
873 members.sort(key=lambda m: m._value_, reverse=True)
874 if len(members) > 1 and members[0].value == value:
875 # we have the breakdown, don't need the value member itself
876 members.pop(0)
877 return members, not_covered
878
879def _power_of_two(value):
880 if value < 1:
881 return False
882 return value == 2 ** _high_bit(value)