blob: f5657a6eba29c19e139bff5089e01d5f1eb362d5 [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__ = [
Ethan Furman9bf7c2d2021-07-03 21:08:42 -07006 'EnumMeta',
7 'Enum', 'IntEnum', 'Flag', 'IntFlag',
8 'auto', 'unique',
Ethan Furmanc16595e2016-09-10 23:36:59 -07009 ]
Ethan Furman6b3d64a2013-06-14 16:55:46 -070010
11
Ethan Furman101e0742013-09-15 12:34:36 -070012def _is_descriptor(obj):
Ethan Furman6d3dfee2020-12-08 12:26:56 -080013 """
14 Returns True if obj is a descriptor, False otherwise.
15 """
Ethan Furman101e0742013-09-15 12:34:36 -070016 return (
17 hasattr(obj, '__get__') or
18 hasattr(obj, '__set__') or
Ethan Furman6d3dfee2020-12-08 12:26:56 -080019 hasattr(obj, '__delete__')
20 )
Ethan Furman101e0742013-09-15 12:34:36 -070021
Ethan Furman6b3d64a2013-06-14 16:55:46 -070022def _is_dunder(name):
Ethan Furman6d3dfee2020-12-08 12:26:56 -080023 """
24 Returns True if a __dunder__ name, False otherwise.
25 """
26 return (
27 len(name) > 4 and
Brennan D Baraban8b914d22019-03-03 14:09:11 -080028 name[:2] == name[-2:] == '__' and
29 name[2] != '_' and
Ethan Furman6d3dfee2020-12-08 12:26:56 -080030 name[-3] != '_'
31 )
Ethan Furman6b3d64a2013-06-14 16:55:46 -070032
33def _is_sunder(name):
Ethan Furman6d3dfee2020-12-08 12:26:56 -080034 """
35 Returns True if a _sunder_ name, False otherwise.
36 """
37 return (
38 len(name) > 2 and
Brennan D Baraban8b914d22019-03-03 14:09:11 -080039 name[0] == name[-1] == '_' and
Ethan Furman6b3d64a2013-06-14 16:55:46 -070040 name[1:2] != '_' and
Ethan Furman6d3dfee2020-12-08 12:26:56 -080041 name[-2:-1] != '_'
42 )
Ethan Furman6b3d64a2013-06-14 16:55:46 -070043
Ethan Furman7cf0aad2020-12-09 17:12:11 -080044def _is_private(cls_name, name):
45 # do not use `re` as `re` imports `enum`
46 pattern = '_%s__' % (cls_name, )
Ethan Furman2a9ab752021-10-20 19:48:37 -070047 pat_len = len(pattern)
Ethan Furman7cf0aad2020-12-09 17:12:11 -080048 if (
Ethan Furman2a9ab752021-10-20 19:48:37 -070049 len(name) > pat_len
Ethan Furman7cf0aad2020-12-09 17:12:11 -080050 and name.startswith(pattern)
Ethan Furman2a9ab752021-10-20 19:48:37 -070051 and name[pat_len:pat_len+1] != ['_']
Ethan Furman7cf0aad2020-12-09 17:12:11 -080052 and (name[-1] != '_' or name[-2] != '_')
53 ):
54 return True
55 else:
56 return False
57
Ethan Furman9bf7c2d2021-07-03 21:08:42 -070058def _make_class_unpicklable(cls):
Ethan Furman7aaeb2a2021-01-25 14:26:19 -080059 """
Ethan Furman9bf7c2d2021-07-03 21:08:42 -070060 Make the given class un-picklable.
Ethan Furman6d3dfee2020-12-08 12:26:56 -080061 """
Ethan Furmanca1b7942014-02-08 11:36:27 -080062 def _break_on_call_reduce(self, proto):
Ethan Furman6b3d64a2013-06-14 16:55:46 -070063 raise TypeError('%r cannot be pickled' % self)
Ethan Furman9bf7c2d2021-07-03 21:08:42 -070064 cls.__reduce_ex__ = _break_on_call_reduce
65 cls.__module__ = '<unknown>'
Ethan Furman7aaeb2a2021-01-25 14:26:19 -080066
Ethan Furman3515dcc2016-09-18 13:15:41 -070067_auto_null = object()
Ethan Furmanc16595e2016-09-10 23:36:59 -070068class auto:
69 """
70 Instances are replaced with an appropriate value in Enum class suites.
71 """
Ethan Furman3515dcc2016-09-18 13:15:41 -070072 value = _auto_null
Ethan Furmanc16595e2016-09-10 23:36:59 -070073
Ethan Furman101e0742013-09-15 12:34:36 -070074
Ethan Furman6b3d64a2013-06-14 16:55:46 -070075class _EnumDict(dict):
Ethan Furman6d3dfee2020-12-08 12:26:56 -080076 """
77 Track enum member order and ensure member names are not reused.
Ethan Furman6b3d64a2013-06-14 16:55:46 -070078
Ethan Furman9bf7c2d2021-07-03 21:08:42 -070079 EnumMeta will use the names found in self._member_names as the
Ethan Furman6b3d64a2013-06-14 16:55:46 -070080 enumeration member names.
Ethan Furman6b3d64a2013-06-14 16:55:46 -070081 """
82 def __init__(self):
83 super().__init__()
84 self._member_names = []
Ethan Furmanc16595e2016-09-10 23:36:59 -070085 self._last_values = []
Ethan Furmana4b1bb42018-01-22 07:56:37 -080086 self._ignore = []
Ethan Onstottd9a43e22020-04-28 13:20:55 -040087 self._auto_called = False
Ethan Furman6b3d64a2013-06-14 16:55:46 -070088
89 def __setitem__(self, key, value):
Ethan Furman6d3dfee2020-12-08 12:26:56 -080090 """
91 Changes anything not dundered or not a descriptor.
Ethan Furman6b3d64a2013-06-14 16:55:46 -070092
93 If an enum member name is used twice, an error is raised; duplicate
94 values are not checked for.
95
96 Single underscore (sunder) names are reserved.
Ethan Furman6b3d64a2013-06-14 16:55:46 -070097 """
Ethan Furman7cf0aad2020-12-09 17:12:11 -080098 if _is_private(self._cls_name, key):
Ethan Furman9bf7c2d2021-07-03 21:08:42 -070099 import warnings
100 warnings.warn(
101 "private variables, such as %r, will be normal attributes in 3.11"
102 % (key, ),
103 DeprecationWarning,
104 stacklevel=2,
105 )
106 if _is_sunder(key):
Ethan Furmanee47e5c2016-08-31 00:12:15 -0700107 if key not in (
Ethan Furman9bf7c2d2021-07-03 21:08:42 -0700108 '_order_', '_create_pseudo_member_',
Ethan Furmana4b1bb42018-01-22 07:56:37 -0800109 '_generate_next_value_', '_missing_', '_ignore_',
Ethan Furmanee47e5c2016-08-31 00:12:15 -0700110 ):
Ethan Furman9bf7c2d2021-07-03 21:08:42 -0700111 raise ValueError('_names_ are reserved for future Enum use')
Ethan Furmanc16595e2016-09-10 23:36:59 -0700112 if key == '_generate_next_value_':
Ethan Onstottd9a43e22020-04-28 13:20:55 -0400113 # check if members already defined as auto()
114 if self._auto_called:
115 raise TypeError("_generate_next_value_ must be defined before members")
Ethan Furman9bf7c2d2021-07-03 21:08:42 -0700116 setattr(self, '_generate_next_value', value)
Ethan Furmana4b1bb42018-01-22 07:56:37 -0800117 elif key == '_ignore_':
118 if isinstance(value, str):
119 value = value.replace(',',' ').split()
120 else:
121 value = list(value)
122 self._ignore = value
123 already = set(value) & set(self._member_names)
124 if already:
Ethan Furman6d3dfee2020-12-08 12:26:56 -0800125 raise ValueError(
126 '_ignore_ cannot specify already set names: %r'
127 % (already, )
128 )
Ethan Furman101e0742013-09-15 12:34:36 -0700129 elif _is_dunder(key):
Ethan Furmane8e61272016-08-20 07:19:31 -0700130 if key == '__order__':
131 key = '_order_'
Ethan Furman101e0742013-09-15 12:34:36 -0700132 elif key in self._member_names:
133 # descriptor overwriting an enum?
Ethan Furman9bf7c2d2021-07-03 21:08:42 -0700134 raise TypeError('Attempted to reuse key: %r' % key)
Ethan Furmana4b1bb42018-01-22 07:56:37 -0800135 elif key in self._ignore:
136 pass
Ethan Furman101e0742013-09-15 12:34:36 -0700137 elif not _is_descriptor(value):
138 if key in self:
139 # enum overwriting a descriptor?
Ethan Furmanee47e5c2016-08-31 00:12:15 -0700140 raise TypeError('%r already defined as: %r' % (key, self[key]))
Ethan Furmanc16595e2016-09-10 23:36:59 -0700141 if isinstance(value, auto):
Ethan Furman3515dcc2016-09-18 13:15:41 -0700142 if value.value == _auto_null:
Ethan Furman6d3dfee2020-12-08 12:26:56 -0800143 value.value = self._generate_next_value(
Ethan Furman9bf7c2d2021-07-03 21:08:42 -0700144 key,
145 1,
146 len(self._member_names),
147 self._last_values[:],
Ethan Furman6d3dfee2020-12-08 12:26:56 -0800148 )
Ethan Furmanfc23a942020-09-16 12:37:54 -0700149 self._auto_called = True
Ethan Furman3515dcc2016-09-18 13:15:41 -0700150 value = value.value
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700151 self._member_names.append(key)
Ethan Furmanc16595e2016-09-10 23:36:59 -0700152 self._last_values.append(value)
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700153 super().__setitem__(key, value)
154
Ethan Furmana6582872020-12-10 13:07:00 -0800155
Ethan Furman9bf7c2d2021-07-03 21:08:42 -0700156# Dummy value for Enum as EnumMeta explicitly checks for it, but of course
157# until EnumMeta finishes running the first time the Enum class doesn't exist.
158# This is also why there are checks in EnumMeta like `if Enum is not None`
159Enum = None
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700160
Ethan Furman9bf7c2d2021-07-03 21:08:42 -0700161class EnumMeta(type):
Ethan Furman6d3dfee2020-12-08 12:26:56 -0800162 """
163 Metaclass for Enum
164 """
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700165 @classmethod
Ethan Furman6ec0ade2020-12-24 10:05:02 -0800166 def __prepare__(metacls, cls, bases, **kwds):
Ethan Furman3064dbf2020-09-16 07:11:57 -0700167 # check that previous enum members do not exist
168 metacls._check_for_existing_members(cls, bases)
Ethan Furmanee47e5c2016-08-31 00:12:15 -0700169 # create the namespace dict
170 enum_dict = _EnumDict()
Ethan Furman7cf0aad2020-12-09 17:12:11 -0800171 enum_dict._cls_name = cls
Ethan Furmanee47e5c2016-08-31 00:12:15 -0700172 # inherit previous flags and _generate_next_value_ function
Ethan Furman3064dbf2020-09-16 07:11:57 -0700173 member_type, first_enum = metacls._get_mixins_(cls, bases)
Ethan Furmanee47e5c2016-08-31 00:12:15 -0700174 if first_enum is not None:
Ethan Furman6d3dfee2020-12-08 12:26:56 -0800175 enum_dict['_generate_next_value_'] = getattr(
176 first_enum, '_generate_next_value_', None,
177 )
Ethan Furmanee47e5c2016-08-31 00:12:15 -0700178 return enum_dict
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700179
Ethan Furman9bf7c2d2021-07-03 21:08:42 -0700180 def __new__(metacls, cls, bases, classdict, **kwds):
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700181 # an Enum class is final once enumeration items have been defined; it
182 # cannot be mixed with other types (int, float, etc.) if it has an
183 # inherited __new__ unless a new __new__ is defined (or the resulting
184 # class will fail).
Ethan Furmana4b1bb42018-01-22 07:56:37 -0800185 #
186 # remove any keys listed in _ignore_
187 classdict.setdefault('_ignore_', []).append('_ignore_')
188 ignore = classdict['_ignore_']
189 for key in ignore:
190 classdict.pop(key, None)
Ethan Furman3064dbf2020-09-16 07:11:57 -0700191 member_type, first_enum = metacls._get_mixins_(cls, bases)
Ethan Furmanc2667362020-12-07 00:17:31 -0800192 __new__, save_new, use_args = metacls._find_new_(
193 classdict, member_type, first_enum,
194 )
Ethan Furman9bf7c2d2021-07-03 21:08:42 -0700195
196 # save enum items into separate mapping so they don't get baked into
197 # the new class
198 enum_members = {k: classdict[k] for k in classdict._member_names}
199 for name in classdict._member_names:
200 del classdict[name]
201
202 # adjust the sunders
203 _order_ = classdict.pop('_order_', None)
204
205 # check for illegal enum names (any others?)
206 invalid_names = set(enum_members) & {'mro', ''}
207 if invalid_names:
208 raise ValueError('Invalid enum member name: {0}'.format(
209 ','.join(invalid_names)))
210
Ethan Furmanc314e602021-01-12 23:47:57 -0800211 # create a default docstring if one has not been provided
212 if '__doc__' not in classdict:
213 classdict['__doc__'] = 'An enumeration.'
Ethan Furman9bf7c2d2021-07-03 21:08:42 -0700214
215 enum_class = super().__new__(metacls, cls, bases, classdict, **kwds)
216 enum_class._member_names_ = [] # names in definition order
217 enum_class._member_map_ = {} # name->value map
218 enum_class._member_type_ = member_type
219
220 # save DynamicClassAttribute attributes from super classes so we know
221 # if we can take the shortcut of storing members in the class dict
222 dynamic_attributes = {
223 k for c in enum_class.mro()
224 for k, v in c.__dict__.items()
225 if isinstance(v, DynamicClassAttribute)
226 }
227
228 # Reverse value->name map for hashable values.
229 enum_class._value2member_map_ = {}
230
231 # If a custom type is mixed into the Enum, and it does not know how
232 # to pickle itself, pickle.dumps will succeed but pickle.loads will
233 # fail. Rather than have the error show up later and possibly far
234 # from the source, sabotage the pickle protocol for this class so
235 # that pickle.dumps also fails.
Ethan Furmanc314e602021-01-12 23:47:57 -0800236 #
Ethan Furman9bf7c2d2021-07-03 21:08:42 -0700237 # However, if the new class implements its own __reduce_ex__, do not
238 # sabotage -- it's on them to make sure it works correctly. We use
239 # __reduce_ex__ instead of any of the others as it is preferred by
240 # pickle over __reduce__, and it handles all pickle protocols.
241 if '__reduce_ex__' not in classdict:
242 if member_type is not object:
243 methods = ('__getnewargs_ex__', '__getnewargs__',
244 '__reduce_ex__', '__reduce__')
245 if not any(m in member_type.__dict__ for m in methods):
246 if '__new__' in classdict:
247 # too late, sabotage
248 _make_class_unpicklable(enum_class)
249 else:
250 # final attempt to verify that pickling would work:
251 # travel mro until __new__ is found, checking for
252 # __reduce__ and friends along the way -- if any of them
253 # are found before/when __new__ is found, pickling should
254 # work
255 sabotage = None
256 for chain in bases:
257 for base in chain.__mro__:
258 if base is object:
259 continue
260 elif any(m in base.__dict__ for m in methods):
261 # found one, we're good
262 sabotage = False
263 break
264 elif '__new__' in base.__dict__:
265 # not good
266 sabotage = True
267 break
268 if sabotage is not None:
269 break
270 if sabotage:
271 _make_class_unpicklable(enum_class)
272 # instantiate them, checking for duplicates as we go
273 # we instantiate first instead of checking for duplicates first in case
274 # a custom __new__ is doing something funky with the values -- such as
275 # auto-numbering ;)
276 for member_name in classdict._member_names:
277 value = enum_members[member_name]
278 if not isinstance(value, tuple):
279 args = (value, )
280 else:
281 args = value
282 if member_type is tuple: # special case for tuple enums
283 args = (args, ) # wrap it one more time
284 if not use_args:
285 enum_member = __new__(enum_class)
286 if not hasattr(enum_member, '_value_'):
287 enum_member._value_ = value
288 else:
289 enum_member = __new__(enum_class, *args)
290 if not hasattr(enum_member, '_value_'):
291 if member_type is object:
292 enum_member._value_ = value
293 else:
294 enum_member._value_ = member_type(*args)
295 value = enum_member._value_
296 enum_member._name_ = member_name
297 enum_member.__objclass__ = enum_class
298 enum_member.__init__(*args)
299 # If another member with the same value was already defined, the
300 # new member becomes an alias to the existing one.
301 for name, canonical_member in enum_class._member_map_.items():
302 if canonical_member._value_ == enum_member._value_:
303 enum_member = canonical_member
304 break
305 else:
306 # Aliases don't appear in member names (only in __members__).
307 enum_class._member_names_.append(member_name)
308 # performance boost for any member that would not shadow
309 # a DynamicClassAttribute
310 if member_name not in dynamic_attributes:
311 setattr(enum_class, member_name, enum_member)
312 # now add to _member_map_
313 enum_class._member_map_[member_name] = enum_member
314 try:
315 # This may fail if value is not hashable. We can't add the value
316 # to the map, and by-value lookups for this value will be
317 # linear.
318 enum_class._value2member_map_[value] = enum_member
319 except TypeError:
320 pass
321
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700322 # double check that repr and friends are not the mixin's or various
323 # things break (such as pickle)
Ethan Furman22415ad2020-09-15 16:28:25 -0700324 # however, if the method is defined in the Enum itself, don't replace
325 # it
Ethan Furmandc870522014-02-18 12:37:12 -0800326 for name in ('__repr__', '__str__', '__format__', '__reduce_ex__'):
Ethan Furman22415ad2020-09-15 16:28:25 -0700327 if name in classdict:
328 continue
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700329 class_method = getattr(enum_class, name)
330 obj_method = getattr(member_type, name, None)
331 enum_method = getattr(first_enum, name, None)
332 if obj_method is not None and obj_method is class_method:
333 setattr(enum_class, name, enum_method)
Ethan Furman9bf7c2d2021-07-03 21:08:42 -0700334
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700335 # replace any other __new__ with our own (as long as Enum is not None,
336 # anyway) -- again, this is to support pickle
337 if Enum is not None:
338 # if the user defined their own __new__, save it before it gets
339 # clobbered in case they subclass later
340 if save_new:
341 enum_class.__new_member__ = __new__
342 enum_class.__new__ = Enum.__new__
Ethan Furman9bf7c2d2021-07-03 21:08:42 -0700343
Ethan Furmane8e61272016-08-20 07:19:31 -0700344 # py3 support for definition order (helps keep py2/py3 code in sync)
345 if _order_ is not None:
346 if isinstance(_order_, str):
347 _order_ = _order_.replace(',', ' ').split()
348 if _order_ != enum_class._member_names_:
Ethan Furman9bf7c2d2021-07-03 21:08:42 -0700349 raise TypeError('member order does not match _order_')
350
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700351 return enum_class
352
Ethan Furman5de67b12016-04-13 23:52:09 -0700353 def __bool__(self):
354 """
355 classes/types should always be True.
356 """
357 return True
358
Ethan Furman9bf7c2d2021-07-03 21:08:42 -0700359 def __call__(cls, value, names=None, *, module=None, qualname=None, type=None, start=1):
Ethan Furman6d3dfee2020-12-08 12:26:56 -0800360 """
361 Either returns an existing member, or creates a new enum class.
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700362
363 This method is used both when an enum class is given a value to match
364 to an enumeration member (i.e. Color(3)) and for the functional API
Ethan Furman23bb6f42016-11-21 09:22:05 -0800365 (i.e. Color = Enum('Color', names='RED GREEN BLUE')).
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700366
Ethan Furman2da95042014-03-03 12:42:52 -0800367 When used for the functional API:
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700368
Ethan Furman2da95042014-03-03 12:42:52 -0800369 `value` will be the name of the new class.
370
371 `names` should be either a string of white-space/comma delimited names
Ethan Furmand9925a12014-09-16 20:35:55 -0700372 (values will start at `start`), or an iterator/mapping of name, value pairs.
Ethan Furman2da95042014-03-03 12:42:52 -0800373
374 `module` should be set to the module this class is being created in;
375 if it is not set, an attempt to find that module will be made, but if
376 it fails the class will not be picklable.
377
378 `qualname` should be set to the actual location this class can be found
379 at in its module; by default it is set to the global scope. If this is
380 not correct, unpickling will fail in some circumstances.
381
382 `type`, if set, will be mixed in as the first base class.
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700383 """
384 if names is None: # simple value lookup
385 return cls.__new__(cls, value)
386 # otherwise, functional API: we're creating a new Enum type
Ethan Furman6d3dfee2020-12-08 12:26:56 -0800387 return cls._create_(
388 value,
389 names,
390 module=module,
391 qualname=qualname,
392 type=type,
393 start=start,
394 )
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700395
Ethan Furman2a9ab752021-10-20 19:48:37 -0700396 def __contains__(cls, obj):
397 if not isinstance(obj, Enum):
398 import warnings
399 warnings.warn(
400 "in 3.12 __contains__ will no longer raise TypeError, but will return True if\n"
401 "obj is a member or a member's value",
402 DeprecationWarning,
403 stacklevel=2,
404 )
Rahul Jha94306522018-09-10 23:51:04 +0530405 raise TypeError(
406 "unsupported operand type(s) for 'in': '%s' and '%s'" % (
Ethan Furman2a9ab752021-10-20 19:48:37 -0700407 type(obj).__qualname__, cls.__class__.__qualname__))
408 return isinstance(obj, cls) and obj._name_ in cls._member_map_
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700409
Ethan Furman64a99722013-09-22 16:18:19 -0700410 def __delattr__(cls, attr):
411 # nicer error message when someone tries to delete an attribute
412 # (see issue19025).
413 if attr in cls._member_map_:
Ethan Furman9bf7c2d2021-07-03 21:08:42 -0700414 raise AttributeError("%s: cannot delete Enum member." % cls.__name__)
Ethan Furman64a99722013-09-22 16:18:19 -0700415 super().__delattr__(attr)
416
Ethan Furman388a3922013-08-12 06:51:41 -0700417 def __dir__(self):
Ethan Furman6d3dfee2020-12-08 12:26:56 -0800418 return (
419 ['__class__', '__doc__', '__members__', '__module__']
420 + self._member_names_
421 )
Ethan Furman388a3922013-08-12 06:51:41 -0700422
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700423 def __getattr__(cls, name):
Ethan Furman6d3dfee2020-12-08 12:26:56 -0800424 """
425 Return the enum member matching `name`
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700426
427 We use __getattr__ instead of descriptors or inserting into the enum
428 class' __dict__ in order to support `name` and `value` being both
429 properties for enum members (which live in the class' __dict__) and
430 enum members themselves.
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700431 """
432 if _is_dunder(name):
433 raise AttributeError(name)
434 try:
Ethan Furman520ad572013-07-19 19:47:21 -0700435 return cls._member_map_[name]
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700436 except KeyError:
437 raise AttributeError(name) from None
438
439 def __getitem__(cls, name):
Ethan Furman520ad572013-07-19 19:47:21 -0700440 return cls._member_map_[name]
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700441
442 def __iter__(cls):
Ethan Furman6d3dfee2020-12-08 12:26:56 -0800443 """
444 Returns members in definition order.
445 """
Ethan Furman520ad572013-07-19 19:47:21 -0700446 return (cls._member_map_[name] for name in cls._member_names_)
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700447
448 def __len__(cls):
Ethan Furman520ad572013-07-19 19:47:21 -0700449 return len(cls._member_names_)
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700450
Ethan Furman9bf7c2d2021-07-03 21:08:42 -0700451 @property
Ethan Furman2131a4a2013-09-14 18:11:24 -0700452 def __members__(cls):
Ethan Furman6d3dfee2020-12-08 12:26:56 -0800453 """
454 Returns a mapping of member name->value.
Ethan Furman2131a4a2013-09-14 18:11:24 -0700455
456 This mapping lists all enum members, including aliases. Note that this
457 is a read-only view of the internal mapping.
Ethan Furman2131a4a2013-09-14 18:11:24 -0700458 """
459 return MappingProxyType(cls._member_map_)
460
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700461 def __repr__(cls):
Ethan Furman9bf7c2d2021-07-03 21:08:42 -0700462 return "<enum %r>" % cls.__name__
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700463
Ethan Furman2131a4a2013-09-14 18:11:24 -0700464 def __reversed__(cls):
Ethan Furman6d3dfee2020-12-08 12:26:56 -0800465 """
466 Returns members in reverse definition order.
467 """
Ethan Furman2131a4a2013-09-14 18:11:24 -0700468 return (cls._member_map_[name] for name in reversed(cls._member_names_))
469
Ethan Furmanf203f2d2013-09-06 07:16:48 -0700470 def __setattr__(cls, name, value):
Ethan Furman6d3dfee2020-12-08 12:26:56 -0800471 """
472 Block attempts to reassign Enum members.
Ethan Furmanf203f2d2013-09-06 07:16:48 -0700473
474 A simple assignment to the class namespace only changes one of the
475 several possible ways to get an Enum member from the Enum class,
476 resulting in an inconsistent Enumeration.
Ethan Furmanf203f2d2013-09-06 07:16:48 -0700477 """
478 member_map = cls.__dict__.get('_member_map_', {})
479 if name in member_map:
Ethan Furman9bf7c2d2021-07-03 21:08:42 -0700480 raise AttributeError('Cannot reassign members.')
Ethan Furmanf203f2d2013-09-06 07:16:48 -0700481 super().__setattr__(name, value)
482
Ethan Furman9bf7c2d2021-07-03 21:08:42 -0700483 def _create_(cls, class_name, names, *, module=None, qualname=None, type=None, start=1):
Ethan Furman6d3dfee2020-12-08 12:26:56 -0800484 """
485 Convenience method to create a new Enum class.
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700486
487 `names` can be:
488
489 * A string containing member names, separated either with spaces or
Ethan Furmand9925a12014-09-16 20:35:55 -0700490 commas. Values are incremented by 1 from `start`.
491 * An iterable of member names. Values are incremented by 1 from `start`.
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700492 * An iterable of (member name, value) pairs.
Ethan Furmand9925a12014-09-16 20:35:55 -0700493 * A mapping of member name -> value pairs.
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700494 """
495 metacls = cls.__class__
496 bases = (cls, ) if type is None else (type, cls)
Ethan Furman3064dbf2020-09-16 07:11:57 -0700497 _, first_enum = cls._get_mixins_(cls, bases)
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700498 classdict = metacls.__prepare__(class_name, bases)
499
500 # special processing needed for names?
501 if isinstance(names, str):
502 names = names.replace(',', ' ').split()
Dong-hee Nadcc8ce42017-06-22 01:52:32 +0900503 if isinstance(names, (tuple, list)) and names and isinstance(names[0], str):
Ethan Furmanee47e5c2016-08-31 00:12:15 -0700504 original_names, names = names, []
Ethan Furmanc16595e2016-09-10 23:36:59 -0700505 last_values = []
Ethan Furmanee47e5c2016-08-31 00:12:15 -0700506 for count, name in enumerate(original_names):
Ethan Furmanc16595e2016-09-10 23:36:59 -0700507 value = first_enum._generate_next_value_(name, start, count, last_values[:])
508 last_values.append(value)
509 names.append((name, value))
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700510
511 # Here, names is either an iterable of (name, value) or a mapping.
512 for item in names:
513 if isinstance(item, str):
514 member_name, member_value = item, names[item]
515 else:
516 member_name, member_value = item
517 classdict[member_name] = member_value
Ethan Furman9bf7c2d2021-07-03 21:08:42 -0700518 enum_class = metacls.__new__(metacls, class_name, bases, classdict)
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700519
520 # TODO: replace the frame hack if a blessed way to know the calling
521 # module is ever developed
522 if module is None:
523 try:
524 module = sys._getframe(2).f_globals['__name__']
Pablo Galindo293dd232019-11-19 21:34:03 +0000525 except (AttributeError, ValueError, KeyError):
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700526 pass
527 if module is None:
Ethan Furman9bf7c2d2021-07-03 21:08:42 -0700528 _make_class_unpicklable(enum_class)
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700529 else:
Ethan Furman9bf7c2d2021-07-03 21:08:42 -0700530 enum_class.__module__ = module
Ethan Furmanca1b7942014-02-08 11:36:27 -0800531 if qualname is not None:
Ethan Furman9bf7c2d2021-07-03 21:08:42 -0700532 enum_class.__qualname__ = qualname
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700533
Ethan Furman9bf7c2d2021-07-03 21:08:42 -0700534 return enum_class
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700535
Ethan Furman9bf7c2d2021-07-03 21:08:42 -0700536 def _convert_(cls, name, module, filter, source=None):
orlnub1230fb9fad2018-09-12 20:28:53 +0300537 """
538 Create a new Enum subclass that replaces a collection of global constants
539 """
540 # convert all constants from source (or module) that pass filter() to
541 # a new Enum called name, and export the enum and its members back to
542 # module;
543 # also, replace the __reduce_ex__ method so unpickling works in
544 # previous Python versions
Ethan Furman9bf7c2d2021-07-03 21:08:42 -0700545 module_globals = vars(sys.modules[module])
orlnub1230fb9fad2018-09-12 20:28:53 +0300546 if source:
Ethan Furman9bf7c2d2021-07-03 21:08:42 -0700547 source = vars(source)
orlnub1230fb9fad2018-09-12 20:28:53 +0300548 else:
549 source = module_globals
550 # _value2member_map_ is populated in the same order every time
551 # for a consistent reverse mapping of number to name when there
552 # are multiple names for the same number.
553 members = [
554 (name, value)
555 for name, value in source.items()
556 if filter(name)]
557 try:
558 # sort by value
559 members.sort(key=lambda t: (t[1], t[0]))
560 except TypeError:
561 # unless some values aren't comparable, in which case sort by name
562 members.sort(key=lambda t: t[0])
Ethan Furman9bf7c2d2021-07-03 21:08:42 -0700563 cls = cls(name, members, module=module)
564 cls.__reduce_ex__ = _reduce_ex_by_name
565 module_globals.update(cls.__members__)
orlnub1230fb9fad2018-09-12 20:28:53 +0300566 module_globals[name] = cls
567 return cls
568
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700569 @staticmethod
Ethan Furman3064dbf2020-09-16 07:11:57 -0700570 def _check_for_existing_members(class_name, bases):
571 for chain in bases:
572 for base in chain.__mro__:
573 if issubclass(base, Enum) and base._member_names_:
Ethan Furman6d3dfee2020-12-08 12:26:56 -0800574 raise TypeError(
575 "%s: cannot extend enumeration %r"
576 % (class_name, base.__name__)
577 )
Ethan Furman3064dbf2020-09-16 07:11:57 -0700578
579 @staticmethod
580 def _get_mixins_(class_name, bases):
Ethan Furman6d3dfee2020-12-08 12:26:56 -0800581 """
582 Returns the type for creating enum members, and the first inherited
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700583 enum class.
584
585 bases: the tuple of bases that was given to __new__
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700586 """
587 if not bases:
588 return object, Enum
589
Ethan Furman5bdab642018-09-21 19:03:09 -0700590 def _find_data_type(bases):
Ethan Furman2a9ab752021-10-20 19:48:37 -0700591 data_types = set()
Ethan Furman5bdab642018-09-21 19:03:09 -0700592 for chain in bases:
Ethan Furmanbff01f32020-09-15 15:56:26 -0700593 candidate = None
Ethan Furman5bdab642018-09-21 19:03:09 -0700594 for base in chain.__mro__:
595 if base is object:
596 continue
Ethan Furmanc2667362020-12-07 00:17:31 -0800597 elif issubclass(base, Enum):
598 if base._member_type_ is not object:
Ethan Furman2a9ab752021-10-20 19:48:37 -0700599 data_types.add(base._member_type_)
Ethan Furmanc2667362020-12-07 00:17:31 -0800600 break
Ethan Furman5bdab642018-09-21 19:03:09 -0700601 elif '__new__' in base.__dict__:
Ethan Furmancd453852018-10-05 23:29:36 -0700602 if issubclass(base, Enum):
Ethan Furman5bdab642018-09-21 19:03:09 -0700603 continue
Ethan Furman2a9ab752021-10-20 19:48:37 -0700604 data_types.add(candidate or base)
Ethan Furmanbff01f32020-09-15 15:56:26 -0700605 break
Ethan Furmanc2667362020-12-07 00:17:31 -0800606 else:
Miss Islington (bot)0f993242021-06-15 14:07:37 -0700607 candidate = candidate or base
Ethan Furmanbff01f32020-09-15 15:56:26 -0700608 if len(data_types) > 1:
Ethan Furman3064dbf2020-09-16 07:11:57 -0700609 raise TypeError('%r: too many data types: %r' % (class_name, data_types))
Ethan Furmanbff01f32020-09-15 15:56:26 -0700610 elif data_types:
Ethan Furman2a9ab752021-10-20 19:48:37 -0700611 return data_types.pop()
Ethan Furmanbff01f32020-09-15 15:56:26 -0700612 else:
613 return None
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700614
Ethan Furman5bdab642018-09-21 19:03:09 -0700615 # ensure final parent class is an Enum derivative, find any concrete
616 # data type, and check that Enum has no members
617 first_enum = bases[-1]
618 if not issubclass(first_enum, Enum):
619 raise TypeError("new enumerations should be created as "
620 "`EnumName([mixin_type, ...] [data_type,] enum_type)`")
621 member_type = _find_data_type(bases) or object
622 if first_enum._member_names_:
623 raise TypeError("Cannot extend enumerations")
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700624 return member_type, first_enum
625
626 @staticmethod
627 def _find_new_(classdict, member_type, first_enum):
Ethan Furman6d3dfee2020-12-08 12:26:56 -0800628 """
629 Returns the __new__ to be used for creating the enum members.
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700630
631 classdict: the class dictionary given to __new__
632 member_type: the data type whose __new__ will be used by default
633 first_enum: enumeration to check for an overriding __new__
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700634 """
635 # now find the correct __new__, checking to see of one was defined
636 # by the user; also check earlier enum classes in case a __new__ was
637 # saved as __new_member__
638 __new__ = classdict.get('__new__', None)
639
640 # should __new__ be saved as __new_member__ later?
Ethan Furman9bf7c2d2021-07-03 21:08:42 -0700641 save_new = __new__ is not None
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700642
643 if __new__ is None:
644 # check all possibles for __new_member__ before falling back to
645 # __new__
646 for method in ('__new_member__', '__new__'):
647 for possible in (member_type, first_enum):
648 target = getattr(possible, method, None)
649 if target not in {
650 None,
651 None.__new__,
652 object.__new__,
653 Enum.__new__,
654 }:
655 __new__ = target
656 break
657 if __new__ is not None:
658 break
659 else:
660 __new__ = object.__new__
661
662 # if a non-object.__new__ is used then whatever value/tuple was
663 # assigned to the enum member name will be passed to __new__ and to the
664 # new enum member's __init__
Ethan Furman9bf7c2d2021-07-03 21:08:42 -0700665 if __new__ is object.__new__:
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700666 use_args = False
667 else:
668 use_args = True
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700669 return __new__, save_new, use_args
670
671
Ethan Furman9bf7c2d2021-07-03 21:08:42 -0700672class Enum(metaclass=EnumMeta):
Ethan Furman6d3dfee2020-12-08 12:26:56 -0800673 """
674 Generic enumeration.
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700675
676 Derive from this class to define new enumerations.
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700677 """
678 def __new__(cls, value):
679 # all enum instances are actually created during class construction
680 # without calling this method; this method is called by the metaclass'
681 # __call__ (i.e. Color(3) ), and by pickle
682 if type(value) is cls:
Ethan Furman23bb6f42016-11-21 09:22:05 -0800683 # For lookups like Color(Color.RED)
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700684 return value
685 # by-value search for a matching enum member
686 # see if it's in the reverse mapping (for hashable values)
Ethan Furman2aa27322013-07-19 19:35:56 -0700687 try:
Andrew Svetlov34ae04f2018-12-26 20:45:33 +0200688 return cls._value2member_map_[value]
689 except KeyError:
690 # Not found, no need to do long O(n) search
691 pass
Ethan Furman2aa27322013-07-19 19:35:56 -0700692 except TypeError:
693 # not there, now do long search -- O(n) behavior
Ethan Furman520ad572013-07-19 19:47:21 -0700694 for member in cls._member_map_.values():
Ethan Furman0081f232014-09-16 17:31:23 -0700695 if member._value_ == value:
Ethan Furman2aa27322013-07-19 19:35:56 -0700696 return member
Ethan Furmanee47e5c2016-08-31 00:12:15 -0700697 # still not found -- try _missing_ hook
Ethan Furman019f0a02018-09-12 11:43:34 -0700698 try:
699 exc = None
700 result = cls._missing_(value)
701 except Exception as e:
702 exc = e
703 result = None
Ethan Furman2a9ab752021-10-20 19:48:37 -0700704 try:
705 if isinstance(result, cls):
706 return result
707 else:
708 ve_exc = ValueError("%r is not a valid %s" % (value, cls.__qualname__))
709 if result is None and exc is None:
710 raise ve_exc
711 elif exc is None:
712 exc = TypeError(
713 'error in %s._missing_: returned %r instead of None or a valid member'
714 % (cls.__name__, result)
715 )
716 if not isinstance(exc, ValueError):
717 exc.__context__ = ve_exc
718 raise exc
719 finally:
720 # ensure all variables that could hold an exception are destroyed
721 exc = None
722 ve_exc = None
Ethan Furmanee47e5c2016-08-31 00:12:15 -0700723
Ethan Furmanc16595e2016-09-10 23:36:59 -0700724 def _generate_next_value_(name, start, count, last_values):
Ethan Furman6d3dfee2020-12-08 12:26:56 -0800725 """
726 Generate the next value when not given.
727
728 name: the name of the member
729 start: the initial start value or None
730 count: the number of existing members
731 last_value: the last value assigned or None
732 """
Ethan Furmanc16595e2016-09-10 23:36:59 -0700733 for last_value in reversed(last_values):
734 try:
735 return last_value + 1
736 except TypeError:
737 pass
738 else:
Ethan Furmanee47e5c2016-08-31 00:12:15 -0700739 return start
Ethan Furmanc16595e2016-09-10 23:36:59 -0700740
Ethan Furmanee47e5c2016-08-31 00:12:15 -0700741 @classmethod
742 def _missing_(cls, value):
Ethan Furmanc95ad7a2020-09-16 10:26:50 -0700743 return None
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700744
745 def __repr__(self):
Ethan Furman9bf7c2d2021-07-03 21:08:42 -0700746 return "<%s.%s: %r>" % (
747 self.__class__.__name__, self._name_, self._value_)
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700748
749 def __str__(self):
Ethan Furman9bf7c2d2021-07-03 21:08:42 -0700750 return "%s.%s" % (self.__class__.__name__, self._name_)
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700751
Ethan Furman388a3922013-08-12 06:51:41 -0700752 def __dir__(self):
Ethan Furman6d3dfee2020-12-08 12:26:56 -0800753 """
754 Returns all members and all public methods
755 """
Ethan Furman0ae550b2014-10-14 08:58:32 -0700756 added_behavior = [
757 m
758 for cls in self.__class__.mro()
759 for m in cls.__dict__
Ethan Furman354ecf12015-03-11 08:43:12 -0700760 if m[0] != '_' and m not in self._member_map_
Angelin BOOZ68526fe2020-09-21 15:11:06 +0200761 ] + [m for m in self.__dict__ if m[0] != '_']
Ethan Furmanec5f8eb2014-10-21 13:40:35 -0700762 return (['__class__', '__doc__', '__module__'] + added_behavior)
Ethan Furman388a3922013-08-12 06:51:41 -0700763
Ethan Furmanec15a822013-08-31 19:17:41 -0700764 def __format__(self, format_spec):
Ethan Furman6d3dfee2020-12-08 12:26:56 -0800765 """
766 Returns format using actual value type unless __str__ has been overridden.
767 """
Ethan Furmanec15a822013-08-31 19:17:41 -0700768 # mixed-in Enums should use the mixed-in type's __format__, otherwise
769 # we can get strange results with the Enum name showing up instead of
770 # the value
Ethan Furman9bf7c2d2021-07-03 21:08:42 -0700771
thatneat2f19e822019-07-04 11:28:37 -0700772 # pure Enum branch, or branch with __str__ explicitly overridden
Ethan Furman9bf7c2d2021-07-03 21:08:42 -0700773 str_overridden = type(self).__str__ not in (Enum.__str__, Flag.__str__)
thatneat2f19e822019-07-04 11:28:37 -0700774 if self._member_type_ is object or str_overridden:
Ethan Furmanec15a822013-08-31 19:17:41 -0700775 cls = str
776 val = str(self)
777 # mix-in branch
778 else:
779 cls = self._member_type_
Ethan Furman0081f232014-09-16 17:31:23 -0700780 val = self._value_
Ethan Furmanec15a822013-08-31 19:17:41 -0700781 return cls.__format__(val, format_spec)
782
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700783 def __hash__(self):
Ethan Furman520ad572013-07-19 19:47:21 -0700784 return hash(self._name_)
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700785
Ethan Furmanca1b7942014-02-08 11:36:27 -0800786 def __reduce_ex__(self, proto):
Ethan Furman9bf7c2d2021-07-03 21:08:42 -0700787 return self.__class__, (self._value_, )
Ethan Furmanca1b7942014-02-08 11:36:27 -0800788
Ethan Furman9bf7c2d2021-07-03 21:08:42 -0700789 # DynamicClassAttribute is used to provide access to the `name` and
790 # `value` properties of enum members while keeping some measure of
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700791 # protection from modification, while still allowing for an enumeration
792 # to have members named `name` and `value`. This works because enumeration
Ethan Furman9bf7c2d2021-07-03 21:08:42 -0700793 # members are not set directly on the enum class -- __getattr__ is
794 # used to look them up.
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700795
Ethan Furman9bf7c2d2021-07-03 21:08:42 -0700796 @DynamicClassAttribute
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700797 def name(self):
Ethan Furmanc850f342013-09-15 16:59:35 -0700798 """The name of the Enum member."""
Ethan Furman520ad572013-07-19 19:47:21 -0700799 return self._name_
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700800
Ethan Furman9bf7c2d2021-07-03 21:08:42 -0700801 @DynamicClassAttribute
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700802 def value(self):
Ethan Furmanc850f342013-09-15 16:59:35 -0700803 """The value of the Enum member."""
Ethan Furman520ad572013-07-19 19:47:21 -0700804 return self._value_
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700805
806
807class IntEnum(int, Enum):
Ethan Furman9bf7c2d2021-07-03 21:08:42 -0700808 """Enum where members are also (and must be) ints"""
Ethan Furman1b4addf2021-06-18 14:25:42 -0700809
Ethan Furman0063ff42020-09-21 17:23:13 -0700810
Ethan Furman9bf7c2d2021-07-03 21:08:42 -0700811def _reduce_ex_by_name(self, proto):
Ethan Furman24e837f2015-03-18 17:27:57 -0700812 return self.name
813
Ethan Furman9bf7c2d2021-07-03 21:08:42 -0700814class Flag(Enum):
Ethan Furman6d3dfee2020-12-08 12:26:56 -0800815 """
816 Support for flags
817 """
Ethan Furmanc16595e2016-09-10 23:36:59 -0700818
819 def _generate_next_value_(name, start, count, last_values):
Ethan Furmanee47e5c2016-08-31 00:12:15 -0700820 """
821 Generate the next value when not given.
822
823 name: the name of the member
HongWeipengbb16fb22019-09-21 13:22:54 +0800824 start: the initial start value or None
Ethan Furmanee47e5c2016-08-31 00:12:15 -0700825 count: the number of existing members
826 last_value: the last value assigned or None
827 """
828 if not count:
829 return start if start is not None else 1
Ethan Furman9bf7c2d2021-07-03 21:08:42 -0700830 for last_value in reversed(last_values):
831 try:
832 high_bit = _high_bit(last_value)
833 break
834 except Exception:
835 raise TypeError('Invalid Flag value: %r' % last_value) from None
Ethan Furmanee47e5c2016-08-31 00:12:15 -0700836 return 2 ** (high_bit+1)
837
838 @classmethod
839 def _missing_(cls, value):
Ethan Furman6d3dfee2020-12-08 12:26:56 -0800840 """
Ethan Furman9bf7c2d2021-07-03 21:08:42 -0700841 Returns member (possibly creating it) if one can be found for value.
Ethan Furman3515dcc2016-09-18 13:15:41 -0700842 """
Ethan Furman9bf7c2d2021-07-03 21:08:42 -0700843 original_value = value
Ethan Furman7aaeb2a2021-01-25 14:26:19 -0800844 if value < 0:
Ethan Furman9bf7c2d2021-07-03 21:08:42 -0700845 value = ~value
846 possible_member = cls._create_pseudo_member_(value)
847 if original_value < 0:
848 possible_member = ~possible_member
849 return possible_member
850
851 @classmethod
852 def _create_pseudo_member_(cls, value):
853 """
854 Create a composite member iff value contains only members.
855 """
856 pseudo_member = cls._value2member_map_.get(value, None)
857 if pseudo_member is None:
858 # verify all bits are accounted for
859 _, extra_flags = _decompose(cls, value)
860 if extra_flags:
861 raise ValueError("%r is not a valid %s" % (value, cls.__qualname__))
Ethan Furman3515dcc2016-09-18 13:15:41 -0700862 # construct a singleton enum pseudo-member
Ethan Furmanee47e5c2016-08-31 00:12:15 -0700863 pseudo_member = object.__new__(cls)
Ethan Furman7aaeb2a2021-01-25 14:26:19 -0800864 pseudo_member._name_ = None
Ethan Furman9bf7c2d2021-07-03 21:08:42 -0700865 pseudo_member._value_ = value
866 # use setdefault in case another thread already created a composite
867 # with this value
Ethan Furman28cf6632017-01-24 12:12:06 -0800868 pseudo_member = cls._value2member_map_.setdefault(value, pseudo_member)
Ethan Furmanee47e5c2016-08-31 00:12:15 -0700869 return pseudo_member
870
Ethan Furmanee47e5c2016-08-31 00:12:15 -0700871 def __contains__(self, other):
Ethan Furman6d3dfee2020-12-08 12:26:56 -0800872 """
873 Returns True if self has at least the same flags set as other.
874 """
Ethan Furmanee47e5c2016-08-31 00:12:15 -0700875 if not isinstance(other, self.__class__):
Rahul Jha94306522018-09-10 23:51:04 +0530876 raise TypeError(
877 "unsupported operand type(s) for 'in': '%s' and '%s'" % (
878 type(other).__qualname__, self.__class__.__qualname__))
Ethan Furmanee47e5c2016-08-31 00:12:15 -0700879 return other._value_ & self._value_ == other._value_
880
Ethan Furmanee47e5c2016-08-31 00:12:15 -0700881 def __repr__(self):
Ethan Furman9bf7c2d2021-07-03 21:08:42 -0700882 cls = self.__class__
883 if self._name_ is not None:
884 return '<%s.%s: %r>' % (cls.__name__, self._name_, self._value_)
885 members, uncovered = _decompose(cls, self._value_)
886 return '<%s.%s: %r>' % (
887 cls.__name__,
888 '|'.join([str(m._name_ or m._value_) for m in members]),
889 self._value_,
890 )
Ethan Furmanee47e5c2016-08-31 00:12:15 -0700891
892 def __str__(self):
893 cls = self.__class__
Ethan Furman9bf7c2d2021-07-03 21:08:42 -0700894 if self._name_ is not None:
895 return '%s.%s' % (cls.__name__, self._name_)
896 members, uncovered = _decompose(cls, self._value_)
897 if len(members) == 1 and members[0]._name_ is None:
898 return '%s.%r' % (cls.__name__, members[0]._value_)
Ethan Furmanee47e5c2016-08-31 00:12:15 -0700899 else:
Ethan Furman9bf7c2d2021-07-03 21:08:42 -0700900 return '%s.%s' % (
901 cls.__name__,
902 '|'.join([str(m._name_ or m._value_) for m in members]),
903 )
Ethan Furmanee47e5c2016-08-31 00:12:15 -0700904
Ethan Furman25d94bb2016-09-02 16:32:32 -0700905 def __bool__(self):
906 return bool(self._value_)
907
Ethan Furmanee47e5c2016-08-31 00:12:15 -0700908 def __or__(self, other):
909 if not isinstance(other, self.__class__):
910 return NotImplemented
911 return self.__class__(self._value_ | other._value_)
912
913 def __and__(self, other):
914 if not isinstance(other, self.__class__):
915 return NotImplemented
916 return self.__class__(self._value_ & other._value_)
917
918 def __xor__(self, other):
919 if not isinstance(other, self.__class__):
920 return NotImplemented
921 return self.__class__(self._value_ ^ other._value_)
922
923 def __invert__(self):
Ethan Furman9bf7c2d2021-07-03 21:08:42 -0700924 members, uncovered = _decompose(self.__class__, self._value_)
925 inverted = self.__class__(0)
926 for m in self.__class__:
927 if m not in members and not (m._value_ & self._value_):
928 inverted = inverted | m
929 return self.__class__(inverted)
Ethan Furmanee47e5c2016-08-31 00:12:15 -0700930
931
Ethan Furman9bf7c2d2021-07-03 21:08:42 -0700932class IntFlag(int, Flag):
Ethan Furman6d3dfee2020-12-08 12:26:56 -0800933 """
934 Support for integer-based Flags
935 """
Ethan Furmanee47e5c2016-08-31 00:12:15 -0700936
Ethan Furman9bf7c2d2021-07-03 21:08:42 -0700937 @classmethod
938 def _missing_(cls, value):
Ethan Furman1b4addf2021-06-18 14:25:42 -0700939 """
Ethan Furman9bf7c2d2021-07-03 21:08:42 -0700940 Returns member (possibly creating it) if one can be found for value.
Ethan Furman1b4addf2021-06-18 14:25:42 -0700941 """
Ethan Furman9bf7c2d2021-07-03 21:08:42 -0700942 if not isinstance(value, int):
943 raise ValueError("%r is not a valid %s" % (value, cls.__qualname__))
944 new_member = cls._create_pseudo_member_(value)
945 return new_member
946
947 @classmethod
948 def _create_pseudo_member_(cls, value):
949 """
950 Create a composite member iff value contains only members.
951 """
952 pseudo_member = cls._value2member_map_.get(value, None)
953 if pseudo_member is None:
954 need_to_create = [value]
955 # get unaccounted for bits
956 _, extra_flags = _decompose(cls, value)
957 # timer = 10
958 while extra_flags:
959 # timer -= 1
960 bit = _high_bit(extra_flags)
961 flag_value = 2 ** bit
962 if (flag_value not in cls._value2member_map_ and
963 flag_value not in need_to_create
964 ):
965 need_to_create.append(flag_value)
966 if extra_flags == -flag_value:
967 extra_flags = 0
968 else:
969 extra_flags ^= flag_value
970 for value in reversed(need_to_create):
971 # construct singleton pseudo-members
972 pseudo_member = int.__new__(cls, value)
973 pseudo_member._name_ = None
974 pseudo_member._value_ = value
975 # use setdefault in case another thread already created a composite
976 # with this value
977 pseudo_member = cls._value2member_map_.setdefault(value, pseudo_member)
978 return pseudo_member
Ethan Furman1b4addf2021-06-18 14:25:42 -0700979
Ethan Furmanee47e5c2016-08-31 00:12:15 -0700980 def __or__(self, other):
Ethan Furman9bf7c2d2021-07-03 21:08:42 -0700981 if not isinstance(other, (self.__class__, int)):
Ethan Furmanee47e5c2016-08-31 00:12:15 -0700982 return NotImplemented
Ethan Furman9bf7c2d2021-07-03 21:08:42 -0700983 result = self.__class__(self._value_ | self.__class__(other)._value_)
984 return result
Ethan Furmanee47e5c2016-08-31 00:12:15 -0700985
986 def __and__(self, other):
Ethan Furman9bf7c2d2021-07-03 21:08:42 -0700987 if not isinstance(other, (self.__class__, int)):
Ethan Furmanee47e5c2016-08-31 00:12:15 -0700988 return NotImplemented
Ethan Furman9bf7c2d2021-07-03 21:08:42 -0700989 return self.__class__(self._value_ & self.__class__(other)._value_)
Ethan Furmanee47e5c2016-08-31 00:12:15 -0700990
991 def __xor__(self, other):
Ethan Furman9bf7c2d2021-07-03 21:08:42 -0700992 if not isinstance(other, (self.__class__, int)):
Ethan Furmanee47e5c2016-08-31 00:12:15 -0700993 return NotImplemented
Ethan Furman9bf7c2d2021-07-03 21:08:42 -0700994 return self.__class__(self._value_ ^ self.__class__(other)._value_)
Ethan Furmanee47e5c2016-08-31 00:12:15 -0700995
996 __ror__ = __or__
997 __rand__ = __and__
998 __rxor__ = __xor__
Ethan Furman9bf7c2d2021-07-03 21:08:42 -0700999
1000 def __invert__(self):
1001 result = self.__class__(~self._value_)
1002 return result
1003
Ethan Furmanee47e5c2016-08-31 00:12:15 -07001004
1005def _high_bit(value):
Ethan Furman6d3dfee2020-12-08 12:26:56 -08001006 """
1007 returns index of highest bit, or -1 if value is zero or negative
1008 """
Ethan Furman3515dcc2016-09-18 13:15:41 -07001009 return value.bit_length() - 1
Ethan Furmanee47e5c2016-08-31 00:12:15 -07001010
Ethan Furmanf24bb352013-07-18 17:05:39 -07001011def unique(enumeration):
Ethan Furman6d3dfee2020-12-08 12:26:56 -08001012 """
1013 Class decorator for enumerations ensuring unique member values.
1014 """
Ethan Furmanf24bb352013-07-18 17:05:39 -07001015 duplicates = []
1016 for name, member in enumeration.__members__.items():
1017 if name != member.name:
1018 duplicates.append((name, member.name))
1019 if duplicates:
1020 alias_details = ', '.join(
1021 ["%s -> %s" % (alias, name) for (alias, name) in duplicates])
1022 raise ValueError('duplicate values found in %r: %s' %
1023 (enumeration, alias_details))
1024 return enumeration
Ethan Furman3515dcc2016-09-18 13:15:41 -07001025
Ethan Furman9bf7c2d2021-07-03 21:08:42 -07001026def _decompose(flag, value):
Ethan Furmanb7751062021-03-30 21:17:26 -07001027 """
Ethan Furman9bf7c2d2021-07-03 21:08:42 -07001028 Extract all members from the value.
Ethan Furmanb7751062021-03-30 21:17:26 -07001029 """
Ethan Furman9bf7c2d2021-07-03 21:08:42 -07001030 # _decompose is only called if the value is not named
1031 not_covered = value
1032 negative = value < 0
1033 members = []
1034 for member in flag:
1035 member_value = member.value
1036 if member_value and member_value & value == member_value:
1037 members.append(member)
1038 not_covered &= ~member_value
1039 if not negative:
1040 tmp = not_covered
1041 while tmp:
1042 flag_value = 2 ** _high_bit(tmp)
1043 if flag_value in flag._value2member_map_:
1044 members.append(flag._value2member_map_[flag_value])
1045 not_covered &= ~flag_value
1046 tmp &= ~flag_value
1047 if not members and value in flag._value2member_map_:
1048 members.append(flag._value2member_map_[value])
1049 members.sort(key=lambda m: m._value_, reverse=True)
1050 if len(members) > 1 and members[0].value == value:
1051 # we have the breakdown, don't need the value member itself
1052 members.pop(0)
1053 return members, not_covered