blob: 0be1b60cd6d8aaa7aaebea10bab619d6bb908209 [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."""
Brennan D Baraban8b914d22019-03-03 14:09:11 -080022 return (len(name) > 4 and
23 name[:2] == name[-2:] == '__' and
24 name[2] != '_' and
25 name[-3] != '_')
Ethan Furman6b3d64a2013-06-14 16:55:46 -070026
27
28def _is_sunder(name):
29 """Returns True if a _sunder_ name, False otherwise."""
Brennan D Baraban8b914d22019-03-03 14:09:11 -080030 return (len(name) > 2 and
31 name[0] == name[-1] == '_' and
Ethan Furman6b3d64a2013-06-14 16:55:46 -070032 name[1:2] != '_' and
Brennan D Baraban8b914d22019-03-03 14:09:11 -080033 name[-2:-1] != '_')
34
Ethan Furman6b3d64a2013-06-14 16:55:46 -070035
Ethan Furman6b3d64a2013-06-14 16:55:46 -070036def _make_class_unpicklable(cls):
37 """Make the given class un-picklable."""
Ethan Furmanca1b7942014-02-08 11:36:27 -080038 def _break_on_call_reduce(self, proto):
Ethan Furman6b3d64a2013-06-14 16:55:46 -070039 raise TypeError('%r cannot be pickled' % self)
Ethan Furmanca1b7942014-02-08 11:36:27 -080040 cls.__reduce_ex__ = _break_on_call_reduce
Ethan Furman6b3d64a2013-06-14 16:55:46 -070041 cls.__module__ = '<unknown>'
42
Ethan Furman3515dcc2016-09-18 13:15:41 -070043_auto_null = object()
Ethan Furmanc16595e2016-09-10 23:36:59 -070044class auto:
45 """
46 Instances are replaced with an appropriate value in Enum class suites.
47 """
Ethan Furman3515dcc2016-09-18 13:15:41 -070048 value = _auto_null
Ethan Furmanc16595e2016-09-10 23:36:59 -070049
Ethan Furman101e0742013-09-15 12:34:36 -070050
Ethan Furman6b3d64a2013-06-14 16:55:46 -070051class _EnumDict(dict):
Ethan Furman101e0742013-09-15 12:34:36 -070052 """Track enum member order and ensure member names are not reused.
Ethan Furman6b3d64a2013-06-14 16:55:46 -070053
54 EnumMeta will use the names found in self._member_names as the
55 enumeration member names.
56
57 """
58 def __init__(self):
59 super().__init__()
60 self._member_names = []
Ethan Furmanc16595e2016-09-10 23:36:59 -070061 self._last_values = []
Ethan Furmana4b1bb42018-01-22 07:56:37 -080062 self._ignore = []
Ethan Furman6b3d64a2013-06-14 16:55:46 -070063
64 def __setitem__(self, key, value):
Ethan Furman101e0742013-09-15 12:34:36 -070065 """Changes anything not dundered or not a descriptor.
Ethan Furman6b3d64a2013-06-14 16:55:46 -070066
67 If an enum member name is used twice, an error is raised; duplicate
68 values are not checked for.
69
70 Single underscore (sunder) names are reserved.
71
72 """
73 if _is_sunder(key):
Ethan Furmanee47e5c2016-08-31 00:12:15 -070074 if key not in (
Ethan Furman3515dcc2016-09-18 13:15:41 -070075 '_order_', '_create_pseudo_member_',
Ethan Furmana4b1bb42018-01-22 07:56:37 -080076 '_generate_next_value_', '_missing_', '_ignore_',
Ethan Furmanee47e5c2016-08-31 00:12:15 -070077 ):
Ethan Furmane8e61272016-08-20 07:19:31 -070078 raise ValueError('_names_ are reserved for future Enum use')
Ethan Furmanc16595e2016-09-10 23:36:59 -070079 if key == '_generate_next_value_':
80 setattr(self, '_generate_next_value', value)
Ethan Furmana4b1bb42018-01-22 07:56:37 -080081 elif key == '_ignore_':
82 if isinstance(value, str):
83 value = value.replace(',',' ').split()
84 else:
85 value = list(value)
86 self._ignore = value
87 already = set(value) & set(self._member_names)
88 if already:
89 raise ValueError('_ignore_ cannot specify already set names: %r' % (already, ))
Ethan Furman101e0742013-09-15 12:34:36 -070090 elif _is_dunder(key):
Ethan Furmane8e61272016-08-20 07:19:31 -070091 if key == '__order__':
92 key = '_order_'
Ethan Furman101e0742013-09-15 12:34:36 -070093 elif key in self._member_names:
94 # descriptor overwriting an enum?
95 raise TypeError('Attempted to reuse key: %r' % key)
Ethan Furmana4b1bb42018-01-22 07:56:37 -080096 elif key in self._ignore:
97 pass
Ethan Furman101e0742013-09-15 12:34:36 -070098 elif not _is_descriptor(value):
99 if key in self:
100 # enum overwriting a descriptor?
Ethan Furmanee47e5c2016-08-31 00:12:15 -0700101 raise TypeError('%r already defined as: %r' % (key, self[key]))
Ethan Furmanc16595e2016-09-10 23:36:59 -0700102 if isinstance(value, auto):
Ethan Furman3515dcc2016-09-18 13:15:41 -0700103 if value.value == _auto_null:
104 value.value = self._generate_next_value(key, 1, len(self._member_names), self._last_values[:])
105 value = value.value
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700106 self._member_names.append(key)
Ethan Furmanc16595e2016-09-10 23:36:59 -0700107 self._last_values.append(value)
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700108 super().__setitem__(key, value)
109
110
Ezio Melotti9a3777e2013-08-17 15:53:55 +0300111# Dummy value for Enum as EnumMeta explicitly checks for it, but of course
112# until EnumMeta finishes running the first time the Enum class doesn't exist.
113# This is also why there are checks in EnumMeta like `if Enum is not None`
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700114Enum = None
115
Ethan Furman332dbc72016-08-20 00:00:52 -0700116
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700117class EnumMeta(type):
118 """Metaclass for Enum"""
119 @classmethod
Ethan Furman332dbc72016-08-20 00:00:52 -0700120 def __prepare__(metacls, cls, bases):
Ethan Furmanee47e5c2016-08-31 00:12:15 -0700121 # create the namespace dict
122 enum_dict = _EnumDict()
123 # inherit previous flags and _generate_next_value_ function
124 member_type, first_enum = metacls._get_mixins_(bases)
125 if first_enum is not None:
126 enum_dict['_generate_next_value_'] = getattr(first_enum, '_generate_next_value_', None)
127 return enum_dict
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700128
Ethan Furman65a5a472016-09-01 23:55:19 -0700129 def __new__(metacls, cls, bases, classdict):
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700130 # an Enum class is final once enumeration items have been defined; it
131 # cannot be mixed with other types (int, float, etc.) if it has an
132 # inherited __new__ unless a new __new__ is defined (or the resulting
133 # class will fail).
Ethan Furmana4b1bb42018-01-22 07:56:37 -0800134 #
135 # remove any keys listed in _ignore_
136 classdict.setdefault('_ignore_', []).append('_ignore_')
137 ignore = classdict['_ignore_']
138 for key in ignore:
139 classdict.pop(key, None)
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700140 member_type, first_enum = metacls._get_mixins_(bases)
141 __new__, save_new, use_args = metacls._find_new_(classdict, member_type,
142 first_enum)
143
144 # save enum items into separate mapping so they don't get baked into
145 # the new class
Ethan Furmanee47e5c2016-08-31 00:12:15 -0700146 enum_members = {k: classdict[k] for k in classdict._member_names}
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700147 for name in classdict._member_names:
148 del classdict[name]
149
Ethan Furmane8e61272016-08-20 07:19:31 -0700150 # adjust the sunders
151 _order_ = classdict.pop('_order_', None)
152
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700153 # check for illegal enum names (any others?)
Brennan D Baraban8b914d22019-03-03 14:09:11 -0800154 invalid_names = set(enum_members) & {'mro', ''}
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700155 if invalid_names:
156 raise ValueError('Invalid enum member name: {0}'.format(
157 ','.join(invalid_names)))
158
Ethan Furman48a724f2015-04-11 23:23:06 -0700159 # create a default docstring if one has not been provided
160 if '__doc__' not in classdict:
161 classdict['__doc__'] = 'An enumeration.'
162
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700163 # create our new Enum type
164 enum_class = super().__new__(metacls, cls, bases, classdict)
Ethan Furman520ad572013-07-19 19:47:21 -0700165 enum_class._member_names_ = [] # names in definition order
INADA Naokie57f91a2018-06-19 01:14:26 +0900166 enum_class._member_map_ = {} # name->value map
Ethan Furman5e5a8232013-08-04 08:42:23 -0700167 enum_class._member_type_ = member_type
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700168
orlnub1230fb9fad2018-09-12 20:28:53 +0300169 # save DynamicClassAttribute attributes from super classes so we know
170 # if we can take the shortcut of storing members in the class dict
171 dynamic_attributes = {k for c in enum_class.mro()
172 for k, v in c.__dict__.items()
173 if isinstance(v, DynamicClassAttribute)}
Ethan Furman354ecf12015-03-11 08:43:12 -0700174
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700175 # Reverse value->name map for hashable values.
Ethan Furman520ad572013-07-19 19:47:21 -0700176 enum_class._value2member_map_ = {}
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700177
Ethan Furman2da95042014-03-03 12:42:52 -0800178 # If a custom type is mixed into the Enum, and it does not know how
179 # to pickle itself, pickle.dumps will succeed but pickle.loads will
180 # fail. Rather than have the error show up later and possibly far
181 # from the source, sabotage the pickle protocol for this class so
182 # that pickle.dumps also fails.
183 #
184 # However, if the new class implements its own __reduce_ex__, do not
185 # sabotage -- it's on them to make sure it works correctly. We use
186 # __reduce_ex__ instead of any of the others as it is preferred by
187 # pickle over __reduce__, and it handles all pickle protocols.
188 if '__reduce_ex__' not in classdict:
Ethan Furmandc870522014-02-18 12:37:12 -0800189 if member_type is not object:
190 methods = ('__getnewargs_ex__', '__getnewargs__',
191 '__reduce_ex__', '__reduce__')
Ethan Furman2da95042014-03-03 12:42:52 -0800192 if not any(m in member_type.__dict__ for m in methods):
Ethan Furmandc870522014-02-18 12:37:12 -0800193 _make_class_unpicklable(enum_class)
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700194
195 # instantiate them, checking for duplicates as we go
196 # we instantiate first instead of checking for duplicates first in case
197 # a custom __new__ is doing something funky with the values -- such as
198 # auto-numbering ;)
199 for member_name in classdict._member_names:
Ethan Furmanee47e5c2016-08-31 00:12:15 -0700200 value = enum_members[member_name]
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700201 if not isinstance(value, tuple):
202 args = (value, )
203 else:
204 args = value
205 if member_type is tuple: # special case for tuple enums
206 args = (args, ) # wrap it one more time
207 if not use_args:
208 enum_member = __new__(enum_class)
Ethan Furmanb41803e2013-07-25 13:50:45 -0700209 if not hasattr(enum_member, '_value_'):
210 enum_member._value_ = value
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700211 else:
212 enum_member = __new__(enum_class, *args)
Ethan Furmanb41803e2013-07-25 13:50:45 -0700213 if not hasattr(enum_member, '_value_'):
Ethan Furmanee47e5c2016-08-31 00:12:15 -0700214 if member_type is object:
215 enum_member._value_ = value
216 else:
217 enum_member._value_ = member_type(*args)
Ethan Furman520ad572013-07-19 19:47:21 -0700218 value = enum_member._value_
Ethan Furman520ad572013-07-19 19:47:21 -0700219 enum_member._name_ = member_name
Ethan Furmanc850f342013-09-15 16:59:35 -0700220 enum_member.__objclass__ = enum_class
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700221 enum_member.__init__(*args)
222 # If another member with the same value was already defined, the
223 # new member becomes an alias to the existing one.
Ethan Furman520ad572013-07-19 19:47:21 -0700224 for name, canonical_member in enum_class._member_map_.items():
Ethan Furman0081f232014-09-16 17:31:23 -0700225 if canonical_member._value_ == enum_member._value_:
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700226 enum_member = canonical_member
227 break
228 else:
229 # Aliases don't appear in member names (only in __members__).
Ethan Furman520ad572013-07-19 19:47:21 -0700230 enum_class._member_names_.append(member_name)
Ethan Furman354ecf12015-03-11 08:43:12 -0700231 # performance boost for any member that would not shadow
232 # a DynamicClassAttribute
orlnub1230fb9fad2018-09-12 20:28:53 +0300233 if member_name not in dynamic_attributes:
Ethan Furman354ecf12015-03-11 08:43:12 -0700234 setattr(enum_class, member_name, enum_member)
235 # now add to _member_map_
Ethan Furman520ad572013-07-19 19:47:21 -0700236 enum_class._member_map_[member_name] = enum_member
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700237 try:
238 # This may fail if value is not hashable. We can't add the value
239 # to the map, and by-value lookups for this value will be
240 # linear.
Ethan Furman520ad572013-07-19 19:47:21 -0700241 enum_class._value2member_map_[value] = enum_member
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700242 except TypeError:
243 pass
244
245 # double check that repr and friends are not the mixin's or various
246 # things break (such as pickle)
Ethan Furmandc870522014-02-18 12:37:12 -0800247 for name in ('__repr__', '__str__', '__format__', '__reduce_ex__'):
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700248 class_method = getattr(enum_class, name)
249 obj_method = getattr(member_type, name, None)
250 enum_method = getattr(first_enum, name, None)
251 if obj_method is not None and obj_method is class_method:
252 setattr(enum_class, name, enum_method)
253
254 # replace any other __new__ with our own (as long as Enum is not None,
255 # anyway) -- again, this is to support pickle
256 if Enum is not None:
257 # if the user defined their own __new__, save it before it gets
258 # clobbered in case they subclass later
259 if save_new:
260 enum_class.__new_member__ = __new__
261 enum_class.__new__ = Enum.__new__
Ethan Furmane8e61272016-08-20 07:19:31 -0700262
263 # py3 support for definition order (helps keep py2/py3 code in sync)
264 if _order_ is not None:
265 if isinstance(_order_, str):
266 _order_ = _order_.replace(',', ' ').split()
267 if _order_ != enum_class._member_names_:
268 raise TypeError('member order does not match _order_')
269
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700270 return enum_class
271
Ethan Furman5de67b12016-04-13 23:52:09 -0700272 def __bool__(self):
273 """
274 classes/types should always be True.
275 """
276 return True
277
Ethan Furmand9925a12014-09-16 20:35:55 -0700278 def __call__(cls, value, names=None, *, module=None, qualname=None, type=None, start=1):
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700279 """Either returns an existing member, or creates a new enum class.
280
281 This method is used both when an enum class is given a value to match
282 to an enumeration member (i.e. Color(3)) and for the functional API
Ethan Furman23bb6f42016-11-21 09:22:05 -0800283 (i.e. Color = Enum('Color', names='RED GREEN BLUE')).
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700284
Ethan Furman2da95042014-03-03 12:42:52 -0800285 When used for the functional API:
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700286
Ethan Furman2da95042014-03-03 12:42:52 -0800287 `value` will be the name of the new class.
288
289 `names` should be either a string of white-space/comma delimited names
Ethan Furmand9925a12014-09-16 20:35:55 -0700290 (values will start at `start`), or an iterator/mapping of name, value pairs.
Ethan Furman2da95042014-03-03 12:42:52 -0800291
292 `module` should be set to the module this class is being created in;
293 if it is not set, an attempt to find that module will be made, but if
294 it fails the class will not be picklable.
295
296 `qualname` should be set to the actual location this class can be found
297 at in its module; by default it is set to the global scope. If this is
298 not correct, unpickling will fail in some circumstances.
299
300 `type`, if set, will be mixed in as the first base class.
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700301
302 """
303 if names is None: # simple value lookup
304 return cls.__new__(cls, value)
305 # otherwise, functional API: we're creating a new Enum type
Ethan Furmand9925a12014-09-16 20:35:55 -0700306 return cls._create_(value, names, module=module, qualname=qualname, type=type, start=start)
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700307
308 def __contains__(cls, member):
Rahul Jha94306522018-09-10 23:51:04 +0530309 if not isinstance(member, Enum):
310 raise TypeError(
311 "unsupported operand type(s) for 'in': '%s' and '%s'" % (
312 type(member).__qualname__, cls.__class__.__qualname__))
Ethan Furman0081f232014-09-16 17:31:23 -0700313 return isinstance(member, cls) and member._name_ in cls._member_map_
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700314
Ethan Furman64a99722013-09-22 16:18:19 -0700315 def __delattr__(cls, attr):
316 # nicer error message when someone tries to delete an attribute
317 # (see issue19025).
318 if attr in cls._member_map_:
319 raise AttributeError(
320 "%s: cannot delete Enum member." % cls.__name__)
321 super().__delattr__(attr)
322
Ethan Furman388a3922013-08-12 06:51:41 -0700323 def __dir__(self):
Ethan Furman64a99722013-09-22 16:18:19 -0700324 return (['__class__', '__doc__', '__members__', '__module__'] +
325 self._member_names_)
Ethan Furman388a3922013-08-12 06:51:41 -0700326
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700327 def __getattr__(cls, name):
328 """Return the enum member matching `name`
329
330 We use __getattr__ instead of descriptors or inserting into the enum
331 class' __dict__ in order to support `name` and `value` being both
332 properties for enum members (which live in the class' __dict__) and
333 enum members themselves.
334
335 """
336 if _is_dunder(name):
337 raise AttributeError(name)
338 try:
Ethan Furman520ad572013-07-19 19:47:21 -0700339 return cls._member_map_[name]
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700340 except KeyError:
341 raise AttributeError(name) from None
342
343 def __getitem__(cls, name):
Ethan Furman520ad572013-07-19 19:47:21 -0700344 return cls._member_map_[name]
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700345
346 def __iter__(cls):
Ethan Furman520ad572013-07-19 19:47:21 -0700347 return (cls._member_map_[name] for name in cls._member_names_)
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700348
349 def __len__(cls):
Ethan Furman520ad572013-07-19 19:47:21 -0700350 return len(cls._member_names_)
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700351
Ethan Furman2131a4a2013-09-14 18:11:24 -0700352 @property
353 def __members__(cls):
354 """Returns a mapping of member name->value.
355
356 This mapping lists all enum members, including aliases. Note that this
357 is a read-only view of the internal mapping.
358
359 """
360 return MappingProxyType(cls._member_map_)
361
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700362 def __repr__(cls):
363 return "<enum %r>" % cls.__name__
364
Ethan Furman2131a4a2013-09-14 18:11:24 -0700365 def __reversed__(cls):
366 return (cls._member_map_[name] for name in reversed(cls._member_names_))
367
Ethan Furmanf203f2d2013-09-06 07:16:48 -0700368 def __setattr__(cls, name, value):
369 """Block attempts to reassign Enum members.
370
371 A simple assignment to the class namespace only changes one of the
372 several possible ways to get an Enum member from the Enum class,
373 resulting in an inconsistent Enumeration.
374
375 """
376 member_map = cls.__dict__.get('_member_map_', {})
377 if name in member_map:
378 raise AttributeError('Cannot reassign members.')
379 super().__setattr__(name, value)
380
anentropicb8e21f12018-04-16 04:40:35 +0100381 def _create_(cls, class_name, names, *, module=None, qualname=None, type=None, start=1):
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700382 """Convenience method to create a new Enum class.
383
384 `names` can be:
385
386 * A string containing member names, separated either with spaces or
Ethan Furmand9925a12014-09-16 20:35:55 -0700387 commas. Values are incremented by 1 from `start`.
388 * An iterable of member names. Values are incremented by 1 from `start`.
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700389 * An iterable of (member name, value) pairs.
Ethan Furmand9925a12014-09-16 20:35:55 -0700390 * A mapping of member name -> value pairs.
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700391
392 """
393 metacls = cls.__class__
394 bases = (cls, ) if type is None else (type, cls)
Ethan Furmanee47e5c2016-08-31 00:12:15 -0700395 _, first_enum = cls._get_mixins_(bases)
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700396 classdict = metacls.__prepare__(class_name, bases)
397
398 # special processing needed for names?
399 if isinstance(names, str):
400 names = names.replace(',', ' ').split()
Dong-hee Nadcc8ce42017-06-22 01:52:32 +0900401 if isinstance(names, (tuple, list)) and names and isinstance(names[0], str):
Ethan Furmanee47e5c2016-08-31 00:12:15 -0700402 original_names, names = names, []
Ethan Furmanc16595e2016-09-10 23:36:59 -0700403 last_values = []
Ethan Furmanee47e5c2016-08-31 00:12:15 -0700404 for count, name in enumerate(original_names):
Ethan Furmanc16595e2016-09-10 23:36:59 -0700405 value = first_enum._generate_next_value_(name, start, count, last_values[:])
406 last_values.append(value)
407 names.append((name, value))
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700408
409 # Here, names is either an iterable of (name, value) or a mapping.
410 for item in names:
411 if isinstance(item, str):
412 member_name, member_value = item, names[item]
413 else:
414 member_name, member_value = item
415 classdict[member_name] = member_value
416 enum_class = metacls.__new__(metacls, class_name, bases, classdict)
417
418 # TODO: replace the frame hack if a blessed way to know the calling
419 # module is ever developed
420 if module is None:
421 try:
422 module = sys._getframe(2).f_globals['__name__']
Pablo Galindo293dd232019-11-19 21:34:03 +0000423 except (AttributeError, ValueError, KeyError):
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700424 pass
425 if module is None:
426 _make_class_unpicklable(enum_class)
427 else:
428 enum_class.__module__ = module
Ethan Furmanca1b7942014-02-08 11:36:27 -0800429 if qualname is not None:
430 enum_class.__qualname__ = qualname
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700431
432 return enum_class
433
orlnub1230fb9fad2018-09-12 20:28:53 +0300434 def _convert_(cls, name, module, filter, source=None):
435 """
436 Create a new Enum subclass that replaces a collection of global constants
437 """
438 # convert all constants from source (or module) that pass filter() to
439 # a new Enum called name, and export the enum and its members back to
440 # module;
441 # also, replace the __reduce_ex__ method so unpickling works in
442 # previous Python versions
443 module_globals = vars(sys.modules[module])
444 if source:
445 source = vars(source)
446 else:
447 source = module_globals
448 # _value2member_map_ is populated in the same order every time
449 # for a consistent reverse mapping of number to name when there
450 # are multiple names for the same number.
451 members = [
452 (name, value)
453 for name, value in source.items()
454 if filter(name)]
455 try:
456 # sort by value
457 members.sort(key=lambda t: (t[1], t[0]))
458 except TypeError:
459 # unless some values aren't comparable, in which case sort by name
460 members.sort(key=lambda t: t[0])
461 cls = cls(name, members, module=module)
462 cls.__reduce_ex__ = _reduce_ex_by_name
463 module_globals.update(cls.__members__)
464 module_globals[name] = cls
465 return cls
466
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700467 @staticmethod
468 def _get_mixins_(bases):
469 """Returns the type for creating enum members, and the first inherited
470 enum class.
471
472 bases: the tuple of bases that was given to __new__
473
474 """
475 if not bases:
476 return object, Enum
477
Ethan Furman5bdab642018-09-21 19:03:09 -0700478 def _find_data_type(bases):
479 for chain in bases:
480 for base in chain.__mro__:
481 if base is object:
482 continue
483 elif '__new__' in base.__dict__:
Ethan Furmancd453852018-10-05 23:29:36 -0700484 if issubclass(base, Enum):
Ethan Furman5bdab642018-09-21 19:03:09 -0700485 continue
486 return base
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700487
Ethan Furman5bdab642018-09-21 19:03:09 -0700488 # ensure final parent class is an Enum derivative, find any concrete
489 # data type, and check that Enum has no members
490 first_enum = bases[-1]
491 if not issubclass(first_enum, Enum):
492 raise TypeError("new enumerations should be created as "
493 "`EnumName([mixin_type, ...] [data_type,] enum_type)`")
494 member_type = _find_data_type(bases) or object
495 if first_enum._member_names_:
496 raise TypeError("Cannot extend enumerations")
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700497 return member_type, first_enum
498
499 @staticmethod
500 def _find_new_(classdict, member_type, first_enum):
501 """Returns the __new__ to be used for creating the enum members.
502
503 classdict: the class dictionary given to __new__
504 member_type: the data type whose __new__ will be used by default
505 first_enum: enumeration to check for an overriding __new__
506
507 """
508 # now find the correct __new__, checking to see of one was defined
509 # by the user; also check earlier enum classes in case a __new__ was
510 # saved as __new_member__
511 __new__ = classdict.get('__new__', None)
512
513 # should __new__ be saved as __new_member__ later?
514 save_new = __new__ is not None
515
516 if __new__ is None:
517 # check all possibles for __new_member__ before falling back to
518 # __new__
519 for method in ('__new_member__', '__new__'):
520 for possible in (member_type, first_enum):
521 target = getattr(possible, method, None)
522 if target not in {
523 None,
524 None.__new__,
525 object.__new__,
526 Enum.__new__,
527 }:
528 __new__ = target
529 break
530 if __new__ is not None:
531 break
532 else:
533 __new__ = object.__new__
534
535 # if a non-object.__new__ is used then whatever value/tuple was
536 # assigned to the enum member name will be passed to __new__ and to the
537 # new enum member's __init__
538 if __new__ is object.__new__:
539 use_args = False
540 else:
541 use_args = True
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700542 return __new__, save_new, use_args
543
544
545class Enum(metaclass=EnumMeta):
546 """Generic enumeration.
547
548 Derive from this class to define new enumerations.
549
550 """
551 def __new__(cls, value):
552 # all enum instances are actually created during class construction
553 # without calling this method; this method is called by the metaclass'
554 # __call__ (i.e. Color(3) ), and by pickle
555 if type(value) is cls:
Ethan Furman23bb6f42016-11-21 09:22:05 -0800556 # For lookups like Color(Color.RED)
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700557 return value
558 # by-value search for a matching enum member
559 # see if it's in the reverse mapping (for hashable values)
Ethan Furman2aa27322013-07-19 19:35:56 -0700560 try:
Andrew Svetlov34ae04f2018-12-26 20:45:33 +0200561 return cls._value2member_map_[value]
562 except KeyError:
563 # Not found, no need to do long O(n) search
564 pass
Ethan Furman2aa27322013-07-19 19:35:56 -0700565 except TypeError:
566 # not there, now do long search -- O(n) behavior
Ethan Furman520ad572013-07-19 19:47:21 -0700567 for member in cls._member_map_.values():
Ethan Furman0081f232014-09-16 17:31:23 -0700568 if member._value_ == value:
Ethan Furman2aa27322013-07-19 19:35:56 -0700569 return member
Ethan Furmanee47e5c2016-08-31 00:12:15 -0700570 # still not found -- try _missing_ hook
Ethan Furman019f0a02018-09-12 11:43:34 -0700571 try:
572 exc = None
573 result = cls._missing_(value)
574 except Exception as e:
575 exc = e
576 result = None
577 if isinstance(result, cls):
578 return result
579 else:
Walter Dörwald323842c2019-07-18 20:37:13 +0200580 ve_exc = ValueError("%r is not a valid %s" % (value, cls.__qualname__))
Ethan Furman019f0a02018-09-12 11:43:34 -0700581 if result is None and exc is None:
582 raise ve_exc
583 elif exc is None:
584 exc = TypeError(
585 'error in %s._missing_: returned %r instead of None or a valid member'
586 % (cls.__name__, result)
587 )
588 exc.__context__ = ve_exc
589 raise exc
Ethan Furmanee47e5c2016-08-31 00:12:15 -0700590
Ethan Furmanc16595e2016-09-10 23:36:59 -0700591 def _generate_next_value_(name, start, count, last_values):
592 for last_value in reversed(last_values):
593 try:
594 return last_value + 1
595 except TypeError:
596 pass
597 else:
Ethan Furmanee47e5c2016-08-31 00:12:15 -0700598 return start
Ethan Furmanc16595e2016-09-10 23:36:59 -0700599
Ethan Furmanee47e5c2016-08-31 00:12:15 -0700600 @classmethod
601 def _missing_(cls, value):
Walter Dörwald323842c2019-07-18 20:37:13 +0200602 raise ValueError("%r is not a valid %s" % (value, cls.__qualname__))
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700603
604 def __repr__(self):
605 return "<%s.%s: %r>" % (
Ethan Furman520ad572013-07-19 19:47:21 -0700606 self.__class__.__name__, self._name_, self._value_)
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700607
608 def __str__(self):
Ethan Furman520ad572013-07-19 19:47:21 -0700609 return "%s.%s" % (self.__class__.__name__, self._name_)
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700610
Ethan Furman388a3922013-08-12 06:51:41 -0700611 def __dir__(self):
Ethan Furman0ae550b2014-10-14 08:58:32 -0700612 added_behavior = [
613 m
614 for cls in self.__class__.mro()
615 for m in cls.__dict__
Ethan Furman354ecf12015-03-11 08:43:12 -0700616 if m[0] != '_' and m not in self._member_map_
Ethan Furman0ae550b2014-10-14 08:58:32 -0700617 ]
Ethan Furmanec5f8eb2014-10-21 13:40:35 -0700618 return (['__class__', '__doc__', '__module__'] + added_behavior)
Ethan Furman388a3922013-08-12 06:51:41 -0700619
Ethan Furmanec15a822013-08-31 19:17:41 -0700620 def __format__(self, format_spec):
621 # mixed-in Enums should use the mixed-in type's __format__, otherwise
622 # we can get strange results with the Enum name showing up instead of
623 # the value
624
thatneat2f19e822019-07-04 11:28:37 -0700625 # pure Enum branch, or branch with __str__ explicitly overridden
626 str_overridden = type(self).__str__ != Enum.__str__
627 if self._member_type_ is object or str_overridden:
Ethan Furmanec15a822013-08-31 19:17:41 -0700628 cls = str
629 val = str(self)
630 # mix-in branch
631 else:
632 cls = self._member_type_
Ethan Furman0081f232014-09-16 17:31:23 -0700633 val = self._value_
Ethan Furmanec15a822013-08-31 19:17:41 -0700634 return cls.__format__(val, format_spec)
635
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700636 def __hash__(self):
Ethan Furman520ad572013-07-19 19:47:21 -0700637 return hash(self._name_)
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700638
Ethan Furmanca1b7942014-02-08 11:36:27 -0800639 def __reduce_ex__(self, proto):
Ethan Furmandc870522014-02-18 12:37:12 -0800640 return self.__class__, (self._value_, )
Ethan Furmanca1b7942014-02-08 11:36:27 -0800641
Ethan Furman33918c12013-09-27 23:02:02 -0700642 # DynamicClassAttribute is used to provide access to the `name` and
643 # `value` properties of enum members while keeping some measure of
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700644 # protection from modification, while still allowing for an enumeration
645 # to have members named `name` and `value`. This works because enumeration
646 # members are not set directly on the enum class -- __getattr__ is
647 # used to look them up.
648
Ethan Furmane03ea372013-09-25 07:14:41 -0700649 @DynamicClassAttribute
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700650 def name(self):
Ethan Furmanc850f342013-09-15 16:59:35 -0700651 """The name of the Enum member."""
Ethan Furman520ad572013-07-19 19:47:21 -0700652 return self._name_
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700653
Ethan Furmane03ea372013-09-25 07:14:41 -0700654 @DynamicClassAttribute
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700655 def value(self):
Ethan Furmanc850f342013-09-15 16:59:35 -0700656 """The value of the Enum member."""
Ethan Furman520ad572013-07-19 19:47:21 -0700657 return self._value_
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700658
659
660class IntEnum(int, Enum):
661 """Enum where members are also (and must be) ints"""
Ethan Furmanf24bb352013-07-18 17:05:39 -0700662
663
Ethan Furman24e837f2015-03-18 17:27:57 -0700664def _reduce_ex_by_name(self, proto):
665 return self.name
666
Ethan Furman65a5a472016-09-01 23:55:19 -0700667class Flag(Enum):
Ethan Furmanee47e5c2016-08-31 00:12:15 -0700668 """Support for flags"""
Ethan Furmanc16595e2016-09-10 23:36:59 -0700669
670 def _generate_next_value_(name, start, count, last_values):
Ethan Furmanee47e5c2016-08-31 00:12:15 -0700671 """
672 Generate the next value when not given.
673
674 name: the name of the member
HongWeipengbb16fb22019-09-21 13:22:54 +0800675 start: the initial start value or None
Ethan Furmanee47e5c2016-08-31 00:12:15 -0700676 count: the number of existing members
677 last_value: the last value assigned or None
678 """
679 if not count:
680 return start if start is not None else 1
Ethan Furmanc16595e2016-09-10 23:36:59 -0700681 for last_value in reversed(last_values):
682 try:
683 high_bit = _high_bit(last_value)
684 break
Ethan Furman3515dcc2016-09-18 13:15:41 -0700685 except Exception:
Ethan Furmanc16595e2016-09-10 23:36:59 -0700686 raise TypeError('Invalid Flag value: %r' % last_value) from None
Ethan Furmanee47e5c2016-08-31 00:12:15 -0700687 return 2 ** (high_bit+1)
688
689 @classmethod
690 def _missing_(cls, value):
691 original_value = value
692 if value < 0:
693 value = ~value
694 possible_member = cls._create_pseudo_member_(value)
Ethan Furmanee47e5c2016-08-31 00:12:15 -0700695 if original_value < 0:
696 possible_member = ~possible_member
697 return possible_member
698
699 @classmethod
700 def _create_pseudo_member_(cls, value):
Ethan Furman3515dcc2016-09-18 13:15:41 -0700701 """
702 Create a composite member iff value contains only members.
703 """
Ethan Furmanee47e5c2016-08-31 00:12:15 -0700704 pseudo_member = cls._value2member_map_.get(value, None)
705 if pseudo_member is None:
Ethan Furman3515dcc2016-09-18 13:15:41 -0700706 # verify all bits are accounted for
707 _, extra_flags = _decompose(cls, value)
708 if extra_flags:
Walter Dörwald323842c2019-07-18 20:37:13 +0200709 raise ValueError("%r is not a valid %s" % (value, cls.__qualname__))
Ethan Furman3515dcc2016-09-18 13:15:41 -0700710 # construct a singleton enum pseudo-member
Ethan Furmanee47e5c2016-08-31 00:12:15 -0700711 pseudo_member = object.__new__(cls)
712 pseudo_member._name_ = None
713 pseudo_member._value_ = value
Ethan Furman28cf6632017-01-24 12:12:06 -0800714 # use setdefault in case another thread already created a composite
715 # with this value
716 pseudo_member = cls._value2member_map_.setdefault(value, pseudo_member)
Ethan Furmanee47e5c2016-08-31 00:12:15 -0700717 return pseudo_member
718
Ethan Furmanee47e5c2016-08-31 00:12:15 -0700719 def __contains__(self, other):
720 if not isinstance(other, self.__class__):
Rahul Jha94306522018-09-10 23:51:04 +0530721 raise TypeError(
722 "unsupported operand type(s) for 'in': '%s' and '%s'" % (
723 type(other).__qualname__, self.__class__.__qualname__))
Ethan Furmanee47e5c2016-08-31 00:12:15 -0700724 return other._value_ & self._value_ == other._value_
725
Ethan Furmanee47e5c2016-08-31 00:12:15 -0700726 def __repr__(self):
727 cls = self.__class__
728 if self._name_ is not None:
729 return '<%s.%s: %r>' % (cls.__name__, self._name_, self._value_)
Ethan Furman3515dcc2016-09-18 13:15:41 -0700730 members, uncovered = _decompose(cls, self._value_)
Ethan Furman27682d22016-09-04 11:39:01 -0700731 return '<%s.%s: %r>' % (
732 cls.__name__,
733 '|'.join([str(m._name_ or m._value_) for m in members]),
734 self._value_,
735 )
Ethan Furmanee47e5c2016-08-31 00:12:15 -0700736
737 def __str__(self):
738 cls = self.__class__
739 if self._name_ is not None:
740 return '%s.%s' % (cls.__name__, self._name_)
Ethan Furman3515dcc2016-09-18 13:15:41 -0700741 members, uncovered = _decompose(cls, self._value_)
Ethan Furmanee47e5c2016-08-31 00:12:15 -0700742 if len(members) == 1 and members[0]._name_ is None:
743 return '%s.%r' % (cls.__name__, members[0]._value_)
744 else:
745 return '%s.%s' % (
746 cls.__name__,
747 '|'.join([str(m._name_ or m._value_) for m in members]),
748 )
749
Ethan Furman25d94bb2016-09-02 16:32:32 -0700750 def __bool__(self):
751 return bool(self._value_)
752
Ethan Furmanee47e5c2016-08-31 00:12:15 -0700753 def __or__(self, other):
754 if not isinstance(other, self.__class__):
755 return NotImplemented
756 return self.__class__(self._value_ | other._value_)
757
758 def __and__(self, other):
759 if not isinstance(other, self.__class__):
760 return NotImplemented
761 return self.__class__(self._value_ & other._value_)
762
763 def __xor__(self, other):
764 if not isinstance(other, self.__class__):
765 return NotImplemented
766 return self.__class__(self._value_ ^ other._value_)
767
768 def __invert__(self):
Ethan Furman3515dcc2016-09-18 13:15:41 -0700769 members, uncovered = _decompose(self.__class__, self._value_)
Serhiy Storchaka81108372017-09-26 00:55:55 +0300770 inverted = self.__class__(0)
771 for m in self.__class__:
772 if m not in members and not (m._value_ & self._value_):
773 inverted = inverted | m
Ethan Furmanee47e5c2016-08-31 00:12:15 -0700774 return self.__class__(inverted)
775
776
Ethan Furman65a5a472016-09-01 23:55:19 -0700777class IntFlag(int, Flag):
Ethan Furmanee47e5c2016-08-31 00:12:15 -0700778 """Support for integer-based Flags"""
779
780 @classmethod
Ethan Furman3515dcc2016-09-18 13:15:41 -0700781 def _missing_(cls, value):
782 if not isinstance(value, int):
Walter Dörwald323842c2019-07-18 20:37:13 +0200783 raise ValueError("%r is not a valid %s" % (value, cls.__qualname__))
Ethan Furman3515dcc2016-09-18 13:15:41 -0700784 new_member = cls._create_pseudo_member_(value)
785 return new_member
786
787 @classmethod
Ethan Furmanee47e5c2016-08-31 00:12:15 -0700788 def _create_pseudo_member_(cls, value):
789 pseudo_member = cls._value2member_map_.get(value, None)
790 if pseudo_member is None:
Ethan Furman3515dcc2016-09-18 13:15:41 -0700791 need_to_create = [value]
792 # get unaccounted for bits
793 _, extra_flags = _decompose(cls, value)
794 # timer = 10
795 while extra_flags:
796 # timer -= 1
797 bit = _high_bit(extra_flags)
798 flag_value = 2 ** bit
799 if (flag_value not in cls._value2member_map_ and
800 flag_value not in need_to_create
801 ):
802 need_to_create.append(flag_value)
803 if extra_flags == -flag_value:
804 extra_flags = 0
805 else:
806 extra_flags ^= flag_value
807 for value in reversed(need_to_create):
808 # construct singleton pseudo-members
809 pseudo_member = int.__new__(cls, value)
810 pseudo_member._name_ = None
811 pseudo_member._value_ = value
Ethan Furman28cf6632017-01-24 12:12:06 -0800812 # use setdefault in case another thread already created a composite
813 # with this value
814 pseudo_member = cls._value2member_map_.setdefault(value, pseudo_member)
Ethan Furmanee47e5c2016-08-31 00:12:15 -0700815 return pseudo_member
816
Ethan Furmanee47e5c2016-08-31 00:12:15 -0700817 def __or__(self, other):
818 if not isinstance(other, (self.__class__, int)):
819 return NotImplemented
Ethan Furman3515dcc2016-09-18 13:15:41 -0700820 result = self.__class__(self._value_ | self.__class__(other)._value_)
821 return result
Ethan Furmanee47e5c2016-08-31 00:12:15 -0700822
823 def __and__(self, other):
824 if not isinstance(other, (self.__class__, int)):
825 return NotImplemented
826 return self.__class__(self._value_ & self.__class__(other)._value_)
827
828 def __xor__(self, other):
829 if not isinstance(other, (self.__class__, int)):
830 return NotImplemented
831 return self.__class__(self._value_ ^ self.__class__(other)._value_)
832
833 __ror__ = __or__
834 __rand__ = __and__
835 __rxor__ = __xor__
836
837 def __invert__(self):
Ethan Furman3515dcc2016-09-18 13:15:41 -0700838 result = self.__class__(~self._value_)
839 return result
Ethan Furmanee47e5c2016-08-31 00:12:15 -0700840
841
842def _high_bit(value):
Ethan Furman04439532016-09-02 15:50:21 -0700843 """returns index of highest bit, or -1 if value is zero or negative"""
Ethan Furman3515dcc2016-09-18 13:15:41 -0700844 return value.bit_length() - 1
Ethan Furmanee47e5c2016-08-31 00:12:15 -0700845
Ethan Furmanf24bb352013-07-18 17:05:39 -0700846def unique(enumeration):
847 """Class decorator for enumerations ensuring unique member values."""
848 duplicates = []
849 for name, member in enumeration.__members__.items():
850 if name != member.name:
851 duplicates.append((name, member.name))
852 if duplicates:
853 alias_details = ', '.join(
854 ["%s -> %s" % (alias, name) for (alias, name) in duplicates])
855 raise ValueError('duplicate values found in %r: %s' %
856 (enumeration, alias_details))
857 return enumeration
Ethan Furman3515dcc2016-09-18 13:15:41 -0700858
859def _decompose(flag, value):
860 """Extract all members from the value."""
861 # _decompose is only called if the value is not named
862 not_covered = value
863 negative = value < 0
Ethan Furman3515dcc2016-09-18 13:15:41 -0700864 members = []
HongWeipeng0b41a922019-11-27 06:36:02 +0800865 for member in flag:
866 member_value = member.value
Ethan Furman3515dcc2016-09-18 13:15:41 -0700867 if member_value and member_value & value == member_value:
868 members.append(member)
869 not_covered &= ~member_value
HongWeipeng0b41a922019-11-27 06:36:02 +0800870 if not negative:
871 tmp = not_covered
872 while tmp:
873 flag_value = 2 ** _high_bit(tmp)
874 if flag_value in flag._value2member_map_:
875 members.append(flag._value2member_map_[flag_value])
876 not_covered &= ~flag_value
877 tmp &= ~flag_value
Ethan Furman3515dcc2016-09-18 13:15:41 -0700878 if not members and value in flag._value2member_map_:
879 members.append(flag._value2member_map_[value])
880 members.sort(key=lambda m: m._value_, reverse=True)
881 if len(members) > 1 and members[0].value == value:
882 # we have the breakdown, don't need the value member itself
883 members.pop(0)
884 return members, not_covered