blob: 8ca385420da029bf3dc72647e4608e79eed8f094 [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 Furmanc314e602021-01-12 23:47:57 -08003from builtins import property as _bltin_property
Ethan Furman6b3d64a2013-06-14 16:55:46 -07004
Ethan Furmane5754ab2015-09-17 22:03:52 -07005
Ethan Furmanc16595e2016-09-10 23:36:59 -07006__all__ = [
7 'EnumMeta',
Ethan Furman0063ff42020-09-21 17:23:13 -07008 'Enum', 'IntEnum', 'StrEnum', 'Flag', 'IntFlag',
Ethan Furmanc16595e2016-09-10 23:36:59 -07009 'auto', 'unique',
Ethan Furmanc314e602021-01-12 23:47:57 -080010 'property',
Ethan Furmanc16595e2016-09-10 23:36:59 -070011 ]
Ethan Furman6b3d64a2013-06-14 16:55:46 -070012
13
Ethan Furman101e0742013-09-15 12:34:36 -070014def _is_descriptor(obj):
Ethan Furman6d3dfee2020-12-08 12:26:56 -080015 """
16 Returns True if obj is a descriptor, False otherwise.
17 """
Ethan Furman101e0742013-09-15 12:34:36 -070018 return (
19 hasattr(obj, '__get__') or
20 hasattr(obj, '__set__') or
Ethan Furman6d3dfee2020-12-08 12:26:56 -080021 hasattr(obj, '__delete__')
22 )
Ethan Furman101e0742013-09-15 12:34:36 -070023
Ethan Furman6b3d64a2013-06-14 16:55:46 -070024def _is_dunder(name):
Ethan Furman6d3dfee2020-12-08 12:26:56 -080025 """
26 Returns True if a __dunder__ name, False otherwise.
27 """
28 return (
29 len(name) > 4 and
Brennan D Baraban8b914d22019-03-03 14:09:11 -080030 name[:2] == name[-2:] == '__' and
31 name[2] != '_' and
Ethan Furman6d3dfee2020-12-08 12:26:56 -080032 name[-3] != '_'
33 )
Ethan Furman6b3d64a2013-06-14 16:55:46 -070034
35def _is_sunder(name):
Ethan Furman6d3dfee2020-12-08 12:26:56 -080036 """
37 Returns True if a _sunder_ name, False otherwise.
38 """
39 return (
40 len(name) > 2 and
Brennan D Baraban8b914d22019-03-03 14:09:11 -080041 name[0] == name[-1] == '_' and
Ethan Furman6b3d64a2013-06-14 16:55:46 -070042 name[1:2] != '_' and
Ethan Furman6d3dfee2020-12-08 12:26:56 -080043 name[-2:-1] != '_'
44 )
Ethan Furman6b3d64a2013-06-14 16:55:46 -070045
Ethan Furman7cf0aad2020-12-09 17:12:11 -080046def _is_private(cls_name, name):
47 # do not use `re` as `re` imports `enum`
48 pattern = '_%s__' % (cls_name, )
49 if (
50 len(name) >= 5
51 and name.startswith(pattern)
52 and name[len(pattern)] != '_'
53 and (name[-1] != '_' or name[-2] != '_')
54 ):
55 return True
56 else:
57 return False
58
Ethan Furmanc314e602021-01-12 23:47:57 -080059def _make_class_unpicklable(obj):
Ethan Furman6d3dfee2020-12-08 12:26:56 -080060 """
Ethan Furmanc314e602021-01-12 23:47:57 -080061 Make the given obj un-picklable.
62
63 obj should be either a dictionary, on an Enum
Ethan Furman6d3dfee2020-12-08 12:26:56 -080064 """
Ethan Furmanca1b7942014-02-08 11:36:27 -080065 def _break_on_call_reduce(self, proto):
Ethan Furman6b3d64a2013-06-14 16:55:46 -070066 raise TypeError('%r cannot be pickled' % self)
Ethan Furmanc314e602021-01-12 23:47:57 -080067 if isinstance(obj, dict):
68 obj['__reduce_ex__'] = _break_on_call_reduce
69 obj['__module__'] = '<unknown>'
70 else:
71 setattr(obj, '__reduce_ex__', _break_on_call_reduce)
72 setattr(obj, '__module__', '<unknown>')
Ethan Furman6b3d64a2013-06-14 16:55:46 -070073
Ethan Furman3515dcc2016-09-18 13:15:41 -070074_auto_null = object()
Ethan Furmanc16595e2016-09-10 23:36:59 -070075class auto:
76 """
77 Instances are replaced with an appropriate value in Enum class suites.
78 """
Ethan Furman3515dcc2016-09-18 13:15:41 -070079 value = _auto_null
Ethan Furmanc16595e2016-09-10 23:36:59 -070080
Ethan Furmanc314e602021-01-12 23:47:57 -080081class property(DynamicClassAttribute):
82 """
83 This is a descriptor, used to define attributes that act differently
84 when accessed through an enum member and through an enum class.
85 Instance access is the same as property(), but access to an attribute
86 through the enum class will instead look in the class' _member_map_ for
87 a corresponding enum member.
88 """
89
90 def __get__(self, instance, ownerclass=None):
91 if instance is None:
92 try:
93 return ownerclass._member_map_[self.name]
94 except KeyError:
95 raise AttributeError('%r not found in %r' % (self.name, ownerclass.__name__))
96 else:
97 if self.fget is None:
98 raise AttributeError('%s: cannot read attribute %r' % (ownerclass.__name__, self.name))
99 else:
100 return self.fget(instance)
101
102 def __set__(self, instance, value):
103 if self.fset is None:
104 raise AttributeError("%s: cannot set attribute %r" % (self.clsname, self.name))
105 else:
106 return self.fset(instance, value)
107
108 def __delete__(self, instance):
109 if self.fdel is None:
110 raise AttributeError("%s: cannot delete attribute %r" % (self.clsname, self.name))
111 else:
112 return self.fdel(instance)
113
114 def __set_name__(self, ownerclass, name):
115 self.name = name
116 self.clsname = ownerclass.__name__
117
118
119class _proto_member:
120 """
121 intermediate step for enum members between class execution and final creation
122 """
123
124 def __init__(self, value):
125 self.value = value
126
127 def __set_name__(self, enum_class, member_name):
128 """
129 convert each quasi-member into an instance of the new enum class
130 """
131 # first step: remove ourself from enum_class
132 delattr(enum_class, member_name)
133 # second step: create member based on enum_class
134 value = self.value
135 if not isinstance(value, tuple):
136 args = (value, )
137 else:
138 args = value
139 if enum_class._member_type_ is tuple: # special case for tuple enums
140 args = (args, ) # wrap it one more time
141 if not enum_class._use_args_:
142 enum_member = enum_class._new_member_(enum_class)
143 if not hasattr(enum_member, '_value_'):
144 enum_member._value_ = value
145 else:
146 enum_member = enum_class._new_member_(enum_class, *args)
147 if not hasattr(enum_member, '_value_'):
148 if enum_class._member_type_ is object:
149 enum_member._value_ = value
150 else:
151 enum_member._value_ = enum_class._member_type_(*args)
152 value = enum_member._value_
153 enum_member._name_ = member_name
154 enum_member.__objclass__ = enum_class
155 enum_member.__init__(*args)
156 # If another member with the same value was already defined, the
157 # new member becomes an alias to the existing one.
158 for name, canonical_member in enum_class._member_map_.items():
159 if canonical_member._value_ == enum_member._value_:
160 enum_member = canonical_member
161 break
162 else:
163 # no other instances found, record this member in _member_names_
164 enum_class._member_names_.append(member_name)
165 # get redirect in place before adding to _member_map_
166 # but check for other instances in parent classes first
167 need_override = False
168 descriptor = None
169 for base in enum_class.__mro__[1:]:
170 descriptor = base.__dict__.get(member_name)
171 if descriptor is not None:
172 if isinstance(descriptor, (property, DynamicClassAttribute)):
173 break
174 else:
175 need_override = True
176 # keep looking for an enum.property
177 if descriptor and not need_override:
178 # previous enum.property found, no further action needed
179 pass
180 else:
181 redirect = property()
182 redirect.__set_name__(enum_class, member_name)
183 if descriptor and need_override:
184 # previous enum.property found, but some other inherited attribute
185 # is in the way; copy fget, fset, fdel to this one
186 redirect.fget = descriptor.fget
187 redirect.fset = descriptor.fset
188 redirect.fdel = descriptor.fdel
189 setattr(enum_class, member_name, redirect)
190 # now add to _member_map_ (even aliases)
191 enum_class._member_map_[member_name] = enum_member
192 try:
193 # This may fail if value is not hashable. We can't add the value
194 # to the map, and by-value lookups for this value will be
195 # linear.
196 enum_class._value2member_map_[value] = enum_member
197 except TypeError:
198 pass
199
Ethan Furman101e0742013-09-15 12:34:36 -0700200
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700201class _EnumDict(dict):
Ethan Furman6d3dfee2020-12-08 12:26:56 -0800202 """
203 Track enum member order and ensure member names are not reused.
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700204
205 EnumMeta will use the names found in self._member_names as the
206 enumeration member names.
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700207 """
208 def __init__(self):
209 super().__init__()
210 self._member_names = []
Ethan Furmanc16595e2016-09-10 23:36:59 -0700211 self._last_values = []
Ethan Furmana4b1bb42018-01-22 07:56:37 -0800212 self._ignore = []
Ethan Onstottd9a43e22020-04-28 13:20:55 -0400213 self._auto_called = False
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700214
215 def __setitem__(self, key, value):
Ethan Furman6d3dfee2020-12-08 12:26:56 -0800216 """
217 Changes anything not dundered or not a descriptor.
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700218
219 If an enum member name is used twice, an error is raised; duplicate
220 values are not checked for.
221
222 Single underscore (sunder) names are reserved.
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700223 """
Ethan Furman7cf0aad2020-12-09 17:12:11 -0800224 if _is_private(self._cls_name, key):
225 # do nothing, name will be a normal attribute
226 pass
227 elif _is_sunder(key):
Ethan Furmanee47e5c2016-08-31 00:12:15 -0700228 if key not in (
Ethan Furman3515dcc2016-09-18 13:15:41 -0700229 '_order_', '_create_pseudo_member_',
Ethan Furmana4b1bb42018-01-22 07:56:37 -0800230 '_generate_next_value_', '_missing_', '_ignore_',
Ethan Furmanee47e5c2016-08-31 00:12:15 -0700231 ):
Ethan Furman6d3dfee2020-12-08 12:26:56 -0800232 raise ValueError(
233 '_sunder_ names, such as %r, are reserved for future Enum use'
234 % (key, )
235 )
Ethan Furmanc16595e2016-09-10 23:36:59 -0700236 if key == '_generate_next_value_':
Ethan Onstottd9a43e22020-04-28 13:20:55 -0400237 # check if members already defined as auto()
238 if self._auto_called:
239 raise TypeError("_generate_next_value_ must be defined before members")
Ethan Furmanc16595e2016-09-10 23:36:59 -0700240 setattr(self, '_generate_next_value', value)
Ethan Furmana4b1bb42018-01-22 07:56:37 -0800241 elif key == '_ignore_':
242 if isinstance(value, str):
243 value = value.replace(',',' ').split()
244 else:
245 value = list(value)
246 self._ignore = value
247 already = set(value) & set(self._member_names)
248 if already:
Ethan Furman6d3dfee2020-12-08 12:26:56 -0800249 raise ValueError(
250 '_ignore_ cannot specify already set names: %r'
251 % (already, )
252 )
Ethan Furman101e0742013-09-15 12:34:36 -0700253 elif _is_dunder(key):
Ethan Furmane8e61272016-08-20 07:19:31 -0700254 if key == '__order__':
255 key = '_order_'
Ethan Furman101e0742013-09-15 12:34:36 -0700256 elif key in self._member_names:
257 # descriptor overwriting an enum?
Ethan Furmana6582872020-12-10 13:07:00 -0800258 raise TypeError('%r already defined as: %r' % (key, self[key]))
Ethan Furmana4b1bb42018-01-22 07:56:37 -0800259 elif key in self._ignore:
260 pass
Ethan Furman101e0742013-09-15 12:34:36 -0700261 elif not _is_descriptor(value):
262 if key in self:
263 # enum overwriting a descriptor?
Ethan Furmanee47e5c2016-08-31 00:12:15 -0700264 raise TypeError('%r already defined as: %r' % (key, self[key]))
Ethan Furmanc16595e2016-09-10 23:36:59 -0700265 if isinstance(value, auto):
Ethan Furman3515dcc2016-09-18 13:15:41 -0700266 if value.value == _auto_null:
Ethan Furman6d3dfee2020-12-08 12:26:56 -0800267 value.value = self._generate_next_value(
268 key,
269 1,
270 len(self._member_names),
271 self._last_values[:],
272 )
Ethan Furmanfc23a942020-09-16 12:37:54 -0700273 self._auto_called = True
Ethan Furman3515dcc2016-09-18 13:15:41 -0700274 value = value.value
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700275 self._member_names.append(key)
Ethan Furmanc16595e2016-09-10 23:36:59 -0700276 self._last_values.append(value)
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700277 super().__setitem__(key, value)
278
Ethan Furmana6582872020-12-10 13:07:00 -0800279 def update(self, members, **more_members):
280 try:
281 for name in members.keys():
282 self[name] = members[name]
283 except AttributeError:
284 for name, value in members:
285 self[name] = value
286 for name, value in more_members.items():
287 self[name] = value
288
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700289
Ezio Melotti9a3777e2013-08-17 15:53:55 +0300290# Dummy value for Enum as EnumMeta explicitly checks for it, but of course
291# until EnumMeta finishes running the first time the Enum class doesn't exist.
292# This is also why there are checks in EnumMeta like `if Enum is not None`
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700293Enum = None
294
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700295class EnumMeta(type):
Ethan Furman6d3dfee2020-12-08 12:26:56 -0800296 """
297 Metaclass for Enum
298 """
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700299 @classmethod
Ethan Furman6ec0ade2020-12-24 10:05:02 -0800300 def __prepare__(metacls, cls, bases, **kwds):
Ethan Furman3064dbf2020-09-16 07:11:57 -0700301 # check that previous enum members do not exist
302 metacls._check_for_existing_members(cls, bases)
Ethan Furmanee47e5c2016-08-31 00:12:15 -0700303 # create the namespace dict
304 enum_dict = _EnumDict()
Ethan Furman7cf0aad2020-12-09 17:12:11 -0800305 enum_dict._cls_name = cls
Ethan Furmanee47e5c2016-08-31 00:12:15 -0700306 # inherit previous flags and _generate_next_value_ function
Ethan Furman3064dbf2020-09-16 07:11:57 -0700307 member_type, first_enum = metacls._get_mixins_(cls, bases)
Ethan Furmanee47e5c2016-08-31 00:12:15 -0700308 if first_enum is not None:
Ethan Furman6d3dfee2020-12-08 12:26:56 -0800309 enum_dict['_generate_next_value_'] = getattr(
310 first_enum, '_generate_next_value_', None,
311 )
Ethan Furmanee47e5c2016-08-31 00:12:15 -0700312 return enum_dict
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700313
Ethan Furman6bd94de2020-12-09 16:41:22 -0800314 def __new__(metacls, cls, bases, classdict, **kwds):
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700315 # an Enum class is final once enumeration items have been defined; it
316 # cannot be mixed with other types (int, float, etc.) if it has an
317 # inherited __new__ unless a new __new__ is defined (or the resulting
318 # class will fail).
Ethan Furmana4b1bb42018-01-22 07:56:37 -0800319 #
320 # remove any keys listed in _ignore_
321 classdict.setdefault('_ignore_', []).append('_ignore_')
322 ignore = classdict['_ignore_']
323 for key in ignore:
324 classdict.pop(key, None)
Ethan Furmanc314e602021-01-12 23:47:57 -0800325 #
326 # grab member names
327 member_names = classdict._member_names
328 #
329 # check for illegal enum names (any others?)
330 invalid_names = set(member_names) & {'mro', ''}
331 if invalid_names:
332 raise ValueError('Invalid enum member name: {0}'.format(
333 ','.join(invalid_names)))
334 #
335 # adjust the sunders
336 _order_ = classdict.pop('_order_', None)
337 # convert to normal dict
338 classdict = dict(classdict.items())
339 #
340 # data type of member and the controlling Enum class
Ethan Furman3064dbf2020-09-16 07:11:57 -0700341 member_type, first_enum = metacls._get_mixins_(cls, bases)
Ethan Furmanc2667362020-12-07 00:17:31 -0800342 __new__, save_new, use_args = metacls._find_new_(
343 classdict, member_type, first_enum,
344 )
Ethan Furmanc314e602021-01-12 23:47:57 -0800345 classdict['_new_member_'] = __new__
346 classdict['_use_args_'] = use_args
347 #
348 # convert future enum members into temporary _proto_members
349 for name in member_names:
350 classdict[name] = _proto_member(classdict[name])
351 #
352 # house keeping structures
353 classdict['_member_names_'] = []
354 classdict['_member_map_'] = {}
355 classdict['_value2member_map_'] = {}
356 classdict['_member_type_'] = member_type
357 #
Ethan Furman2da95042014-03-03 12:42:52 -0800358 # If a custom type is mixed into the Enum, and it does not know how
359 # to pickle itself, pickle.dumps will succeed but pickle.loads will
360 # fail. Rather than have the error show up later and possibly far
361 # from the source, sabotage the pickle protocol for this class so
362 # that pickle.dumps also fails.
363 #
364 # However, if the new class implements its own __reduce_ex__, do not
365 # sabotage -- it's on them to make sure it works correctly. We use
366 # __reduce_ex__ instead of any of the others as it is preferred by
367 # pickle over __reduce__, and it handles all pickle protocols.
368 if '__reduce_ex__' not in classdict:
Ethan Furmandc870522014-02-18 12:37:12 -0800369 if member_type is not object:
370 methods = ('__getnewargs_ex__', '__getnewargs__',
371 '__reduce_ex__', '__reduce__')
Ethan Furman2da95042014-03-03 12:42:52 -0800372 if not any(m in member_type.__dict__ for m in methods):
Ethan Furmanc314e602021-01-12 23:47:57 -0800373 _make_class_unpicklable(classdict)
374 #
375 # create a default docstring if one has not been provided
376 if '__doc__' not in classdict:
377 classdict['__doc__'] = 'An enumeration.'
378 try:
379 exc = None
380 enum_class = super().__new__(metacls, cls, bases, classdict, **kwds)
381 except RuntimeError as e:
382 # any exceptions raised by member.__new__ will get converted to a
383 # RuntimeError, so get that original exception back and raise it instead
384 exc = e.__cause__ or e
385 if exc is not None:
386 raise exc
387 #
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700388 # double check that repr and friends are not the mixin's or various
389 # things break (such as pickle)
Ethan Furman22415ad2020-09-15 16:28:25 -0700390 # however, if the method is defined in the Enum itself, don't replace
391 # it
Ethan Furmandc870522014-02-18 12:37:12 -0800392 for name in ('__repr__', '__str__', '__format__', '__reduce_ex__'):
Ethan Furman22415ad2020-09-15 16:28:25 -0700393 if name in classdict:
394 continue
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700395 class_method = getattr(enum_class, name)
396 obj_method = getattr(member_type, name, None)
397 enum_method = getattr(first_enum, name, None)
398 if obj_method is not None and obj_method is class_method:
399 setattr(enum_class, name, enum_method)
Ethan Furmanc314e602021-01-12 23:47:57 -0800400 #
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700401 # replace any other __new__ with our own (as long as Enum is not None,
402 # anyway) -- again, this is to support pickle
403 if Enum is not None:
404 # if the user defined their own __new__, save it before it gets
405 # clobbered in case they subclass later
406 if save_new:
407 enum_class.__new_member__ = __new__
408 enum_class.__new__ = Enum.__new__
Ethan Furmanc314e602021-01-12 23:47:57 -0800409 #
Ethan Furmane8e61272016-08-20 07:19:31 -0700410 # py3 support for definition order (helps keep py2/py3 code in sync)
411 if _order_ is not None:
412 if isinstance(_order_, str):
413 _order_ = _order_.replace(',', ' ').split()
414 if _order_ != enum_class._member_names_:
415 raise TypeError('member order does not match _order_')
Ethan Furmanc314e602021-01-12 23:47:57 -0800416 #
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700417 return enum_class
418
Ethan Furman5de67b12016-04-13 23:52:09 -0700419 def __bool__(self):
420 """
421 classes/types should always be True.
422 """
423 return True
424
Ethan Furmand9925a12014-09-16 20:35:55 -0700425 def __call__(cls, value, names=None, *, module=None, qualname=None, type=None, start=1):
Ethan Furman6d3dfee2020-12-08 12:26:56 -0800426 """
427 Either returns an existing member, or creates a new enum class.
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700428
429 This method is used both when an enum class is given a value to match
430 to an enumeration member (i.e. Color(3)) and for the functional API
Ethan Furman23bb6f42016-11-21 09:22:05 -0800431 (i.e. Color = Enum('Color', names='RED GREEN BLUE')).
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700432
Ethan Furman2da95042014-03-03 12:42:52 -0800433 When used for the functional API:
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700434
Ethan Furman2da95042014-03-03 12:42:52 -0800435 `value` will be the name of the new class.
436
437 `names` should be either a string of white-space/comma delimited names
Ethan Furmand9925a12014-09-16 20:35:55 -0700438 (values will start at `start`), or an iterator/mapping of name, value pairs.
Ethan Furman2da95042014-03-03 12:42:52 -0800439
440 `module` should be set to the module this class is being created in;
441 if it is not set, an attempt to find that module will be made, but if
442 it fails the class will not be picklable.
443
444 `qualname` should be set to the actual location this class can be found
445 at in its module; by default it is set to the global scope. If this is
446 not correct, unpickling will fail in some circumstances.
447
448 `type`, if set, will be mixed in as the first base class.
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700449 """
450 if names is None: # simple value lookup
451 return cls.__new__(cls, value)
452 # otherwise, functional API: we're creating a new Enum type
Ethan Furman6d3dfee2020-12-08 12:26:56 -0800453 return cls._create_(
454 value,
455 names,
456 module=module,
457 qualname=qualname,
458 type=type,
459 start=start,
460 )
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700461
462 def __contains__(cls, member):
Rahul Jha94306522018-09-10 23:51:04 +0530463 if not isinstance(member, Enum):
464 raise TypeError(
465 "unsupported operand type(s) for 'in': '%s' and '%s'" % (
466 type(member).__qualname__, cls.__class__.__qualname__))
Ethan Furman0081f232014-09-16 17:31:23 -0700467 return isinstance(member, cls) and member._name_ in cls._member_map_
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700468
Ethan Furman64a99722013-09-22 16:18:19 -0700469 def __delattr__(cls, attr):
470 # nicer error message when someone tries to delete an attribute
471 # (see issue19025).
472 if attr in cls._member_map_:
Ethan Furman6d3dfee2020-12-08 12:26:56 -0800473 raise AttributeError("%s: cannot delete Enum member %r." % (cls.__name__, attr))
Ethan Furman64a99722013-09-22 16:18:19 -0700474 super().__delattr__(attr)
475
Ethan Furman388a3922013-08-12 06:51:41 -0700476 def __dir__(self):
Ethan Furman6d3dfee2020-12-08 12:26:56 -0800477 return (
478 ['__class__', '__doc__', '__members__', '__module__']
479 + self._member_names_
480 )
Ethan Furman388a3922013-08-12 06:51:41 -0700481
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700482 def __getattr__(cls, name):
Ethan Furman6d3dfee2020-12-08 12:26:56 -0800483 """
484 Return the enum member matching `name`
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700485
486 We use __getattr__ instead of descriptors or inserting into the enum
487 class' __dict__ in order to support `name` and `value` being both
488 properties for enum members (which live in the class' __dict__) and
489 enum members themselves.
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700490 """
491 if _is_dunder(name):
492 raise AttributeError(name)
493 try:
Ethan Furman520ad572013-07-19 19:47:21 -0700494 return cls._member_map_[name]
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700495 except KeyError:
496 raise AttributeError(name) from None
497
498 def __getitem__(cls, name):
Ethan Furman520ad572013-07-19 19:47:21 -0700499 return cls._member_map_[name]
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700500
501 def __iter__(cls):
Ethan Furman6d3dfee2020-12-08 12:26:56 -0800502 """
503 Returns members in definition order.
504 """
Ethan Furman520ad572013-07-19 19:47:21 -0700505 return (cls._member_map_[name] for name in cls._member_names_)
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700506
507 def __len__(cls):
Ethan Furman520ad572013-07-19 19:47:21 -0700508 return len(cls._member_names_)
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700509
Ethan Furmanc314e602021-01-12 23:47:57 -0800510 @_bltin_property
Ethan Furman2131a4a2013-09-14 18:11:24 -0700511 def __members__(cls):
Ethan Furman6d3dfee2020-12-08 12:26:56 -0800512 """
513 Returns a mapping of member name->value.
Ethan Furman2131a4a2013-09-14 18:11:24 -0700514
515 This mapping lists all enum members, including aliases. Note that this
516 is a read-only view of the internal mapping.
Ethan Furman2131a4a2013-09-14 18:11:24 -0700517 """
518 return MappingProxyType(cls._member_map_)
519
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700520 def __repr__(cls):
521 return "<enum %r>" % cls.__name__
522
Ethan Furman2131a4a2013-09-14 18:11:24 -0700523 def __reversed__(cls):
Ethan Furman6d3dfee2020-12-08 12:26:56 -0800524 """
525 Returns members in reverse definition order.
526 """
Ethan Furman2131a4a2013-09-14 18:11:24 -0700527 return (cls._member_map_[name] for name in reversed(cls._member_names_))
528
Ethan Furmanf203f2d2013-09-06 07:16:48 -0700529 def __setattr__(cls, name, value):
Ethan Furman6d3dfee2020-12-08 12:26:56 -0800530 """
531 Block attempts to reassign Enum members.
Ethan Furmanf203f2d2013-09-06 07:16:48 -0700532
533 A simple assignment to the class namespace only changes one of the
534 several possible ways to get an Enum member from the Enum class,
535 resulting in an inconsistent Enumeration.
Ethan Furmanf203f2d2013-09-06 07:16:48 -0700536 """
537 member_map = cls.__dict__.get('_member_map_', {})
538 if name in member_map:
539 raise AttributeError('Cannot reassign members.')
540 super().__setattr__(name, value)
541
anentropicb8e21f12018-04-16 04:40:35 +0100542 def _create_(cls, class_name, names, *, module=None, qualname=None, type=None, start=1):
Ethan Furman6d3dfee2020-12-08 12:26:56 -0800543 """
544 Convenience method to create a new Enum class.
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700545
546 `names` can be:
547
548 * A string containing member names, separated either with spaces or
Ethan Furmand9925a12014-09-16 20:35:55 -0700549 commas. Values are incremented by 1 from `start`.
550 * An iterable of member names. Values are incremented by 1 from `start`.
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700551 * An iterable of (member name, value) pairs.
Ethan Furmand9925a12014-09-16 20:35:55 -0700552 * A mapping of member name -> value pairs.
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700553 """
554 metacls = cls.__class__
555 bases = (cls, ) if type is None else (type, cls)
Ethan Furman3064dbf2020-09-16 07:11:57 -0700556 _, first_enum = cls._get_mixins_(cls, bases)
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700557 classdict = metacls.__prepare__(class_name, bases)
558
559 # special processing needed for names?
560 if isinstance(names, str):
561 names = names.replace(',', ' ').split()
Dong-hee Nadcc8ce42017-06-22 01:52:32 +0900562 if isinstance(names, (tuple, list)) and names and isinstance(names[0], str):
Ethan Furmanee47e5c2016-08-31 00:12:15 -0700563 original_names, names = names, []
Ethan Furmanc16595e2016-09-10 23:36:59 -0700564 last_values = []
Ethan Furmanee47e5c2016-08-31 00:12:15 -0700565 for count, name in enumerate(original_names):
Ethan Furmanc16595e2016-09-10 23:36:59 -0700566 value = first_enum._generate_next_value_(name, start, count, last_values[:])
567 last_values.append(value)
568 names.append((name, value))
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700569
570 # Here, names is either an iterable of (name, value) or a mapping.
571 for item in names:
572 if isinstance(item, str):
573 member_name, member_value = item, names[item]
574 else:
575 member_name, member_value = item
576 classdict[member_name] = member_value
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700577
578 # TODO: replace the frame hack if a blessed way to know the calling
579 # module is ever developed
580 if module is None:
581 try:
582 module = sys._getframe(2).f_globals['__name__']
Pablo Galindo293dd232019-11-19 21:34:03 +0000583 except (AttributeError, ValueError, KeyError):
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700584 pass
585 if module is None:
Ethan Furmanc314e602021-01-12 23:47:57 -0800586 _make_class_unpicklable(classdict)
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700587 else:
Ethan Furmanc314e602021-01-12 23:47:57 -0800588 classdict['__module__'] = module
Ethan Furmanca1b7942014-02-08 11:36:27 -0800589 if qualname is not None:
Ethan Furmanc314e602021-01-12 23:47:57 -0800590 classdict['__qualname__'] = qualname
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700591
Ethan Furmanc314e602021-01-12 23:47:57 -0800592 return metacls.__new__(metacls, class_name, bases, classdict)
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700593
orlnub1230fb9fad2018-09-12 20:28:53 +0300594 def _convert_(cls, name, module, filter, source=None):
595 """
596 Create a new Enum subclass that replaces a collection of global constants
597 """
598 # convert all constants from source (or module) that pass filter() to
599 # a new Enum called name, and export the enum and its members back to
600 # module;
601 # also, replace the __reduce_ex__ method so unpickling works in
602 # previous Python versions
603 module_globals = vars(sys.modules[module])
604 if source:
605 source = vars(source)
606 else:
607 source = module_globals
608 # _value2member_map_ is populated in the same order every time
609 # for a consistent reverse mapping of number to name when there
610 # are multiple names for the same number.
611 members = [
612 (name, value)
613 for name, value in source.items()
614 if filter(name)]
615 try:
616 # sort by value
617 members.sort(key=lambda t: (t[1], t[0]))
618 except TypeError:
619 # unless some values aren't comparable, in which case sort by name
620 members.sort(key=lambda t: t[0])
621 cls = cls(name, members, module=module)
622 cls.__reduce_ex__ = _reduce_ex_by_name
623 module_globals.update(cls.__members__)
624 module_globals[name] = cls
625 return cls
626
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700627 @staticmethod
Ethan Furman3064dbf2020-09-16 07:11:57 -0700628 def _check_for_existing_members(class_name, bases):
629 for chain in bases:
630 for base in chain.__mro__:
631 if issubclass(base, Enum) and base._member_names_:
Ethan Furman6d3dfee2020-12-08 12:26:56 -0800632 raise TypeError(
633 "%s: cannot extend enumeration %r"
634 % (class_name, base.__name__)
635 )
Ethan Furman3064dbf2020-09-16 07:11:57 -0700636
637 @staticmethod
638 def _get_mixins_(class_name, bases):
Ethan Furman6d3dfee2020-12-08 12:26:56 -0800639 """
640 Returns the type for creating enum members, and the first inherited
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700641 enum class.
642
643 bases: the tuple of bases that was given to __new__
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700644 """
645 if not bases:
646 return object, Enum
647
Ethan Furman5bdab642018-09-21 19:03:09 -0700648 def _find_data_type(bases):
Ethan Furmanbff01f32020-09-15 15:56:26 -0700649 data_types = []
Ethan Furman5bdab642018-09-21 19:03:09 -0700650 for chain in bases:
Ethan Furmanbff01f32020-09-15 15:56:26 -0700651 candidate = None
Ethan Furman5bdab642018-09-21 19:03:09 -0700652 for base in chain.__mro__:
653 if base is object:
654 continue
Ethan Furmanc2667362020-12-07 00:17:31 -0800655 elif issubclass(base, Enum):
656 if base._member_type_ is not object:
657 data_types.append(base._member_type_)
658 break
Ethan Furman5bdab642018-09-21 19:03:09 -0700659 elif '__new__' in base.__dict__:
Ethan Furmancd453852018-10-05 23:29:36 -0700660 if issubclass(base, Enum):
Ethan Furman5bdab642018-09-21 19:03:09 -0700661 continue
Ethan Furmanbff01f32020-09-15 15:56:26 -0700662 data_types.append(candidate or base)
663 break
Ethan Furmanc2667362020-12-07 00:17:31 -0800664 else:
Ethan Furmanbff01f32020-09-15 15:56:26 -0700665 candidate = base
666 if len(data_types) > 1:
Ethan Furman3064dbf2020-09-16 07:11:57 -0700667 raise TypeError('%r: too many data types: %r' % (class_name, data_types))
Ethan Furmanbff01f32020-09-15 15:56:26 -0700668 elif data_types:
669 return data_types[0]
670 else:
671 return None
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700672
Ethan Furman5bdab642018-09-21 19:03:09 -0700673 # ensure final parent class is an Enum derivative, find any concrete
674 # data type, and check that Enum has no members
675 first_enum = bases[-1]
676 if not issubclass(first_enum, Enum):
677 raise TypeError("new enumerations should be created as "
678 "`EnumName([mixin_type, ...] [data_type,] enum_type)`")
679 member_type = _find_data_type(bases) or object
680 if first_enum._member_names_:
681 raise TypeError("Cannot extend enumerations")
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700682 return member_type, first_enum
683
684 @staticmethod
685 def _find_new_(classdict, member_type, first_enum):
Ethan Furman6d3dfee2020-12-08 12:26:56 -0800686 """
687 Returns the __new__ to be used for creating the enum members.
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700688
689 classdict: the class dictionary given to __new__
690 member_type: the data type whose __new__ will be used by default
691 first_enum: enumeration to check for an overriding __new__
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700692 """
693 # now find the correct __new__, checking to see of one was defined
694 # by the user; also check earlier enum classes in case a __new__ was
695 # saved as __new_member__
696 __new__ = classdict.get('__new__', None)
697
698 # should __new__ be saved as __new_member__ later?
699 save_new = __new__ is not None
700
701 if __new__ is None:
702 # check all possibles for __new_member__ before falling back to
703 # __new__
704 for method in ('__new_member__', '__new__'):
705 for possible in (member_type, first_enum):
706 target = getattr(possible, method, None)
707 if target not in {
708 None,
709 None.__new__,
710 object.__new__,
711 Enum.__new__,
712 }:
713 __new__ = target
714 break
715 if __new__ is not None:
716 break
717 else:
718 __new__ = object.__new__
719
720 # if a non-object.__new__ is used then whatever value/tuple was
721 # assigned to the enum member name will be passed to __new__ and to the
722 # new enum member's __init__
723 if __new__ is object.__new__:
724 use_args = False
725 else:
726 use_args = True
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700727 return __new__, save_new, use_args
728
729
730class Enum(metaclass=EnumMeta):
Ethan Furman6d3dfee2020-12-08 12:26:56 -0800731 """
732 Generic enumeration.
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700733
734 Derive from this class to define new enumerations.
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700735 """
736 def __new__(cls, value):
737 # all enum instances are actually created during class construction
738 # without calling this method; this method is called by the metaclass'
739 # __call__ (i.e. Color(3) ), and by pickle
740 if type(value) is cls:
Ethan Furman23bb6f42016-11-21 09:22:05 -0800741 # For lookups like Color(Color.RED)
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700742 return value
743 # by-value search for a matching enum member
744 # see if it's in the reverse mapping (for hashable values)
Ethan Furman2aa27322013-07-19 19:35:56 -0700745 try:
Andrew Svetlov34ae04f2018-12-26 20:45:33 +0200746 return cls._value2member_map_[value]
747 except KeyError:
748 # Not found, no need to do long O(n) search
749 pass
Ethan Furman2aa27322013-07-19 19:35:56 -0700750 except TypeError:
751 # not there, now do long search -- O(n) behavior
Ethan Furman520ad572013-07-19 19:47:21 -0700752 for member in cls._member_map_.values():
Ethan Furman0081f232014-09-16 17:31:23 -0700753 if member._value_ == value:
Ethan Furman2aa27322013-07-19 19:35:56 -0700754 return member
Ethan Furmanee47e5c2016-08-31 00:12:15 -0700755 # still not found -- try _missing_ hook
Ethan Furman019f0a02018-09-12 11:43:34 -0700756 try:
757 exc = None
758 result = cls._missing_(value)
759 except Exception as e:
760 exc = e
761 result = None
762 if isinstance(result, cls):
763 return result
764 else:
Walter Dörwald323842c2019-07-18 20:37:13 +0200765 ve_exc = ValueError("%r is not a valid %s" % (value, cls.__qualname__))
Ethan Furman019f0a02018-09-12 11:43:34 -0700766 if result is None and exc is None:
767 raise ve_exc
768 elif exc is None:
769 exc = TypeError(
770 'error in %s._missing_: returned %r instead of None or a valid member'
771 % (cls.__name__, result)
772 )
773 exc.__context__ = ve_exc
774 raise exc
Ethan Furmanee47e5c2016-08-31 00:12:15 -0700775
Ethan Furmanc16595e2016-09-10 23:36:59 -0700776 def _generate_next_value_(name, start, count, last_values):
Ethan Furman6d3dfee2020-12-08 12:26:56 -0800777 """
778 Generate the next value when not given.
779
780 name: the name of the member
781 start: the initial start value or None
782 count: the number of existing members
783 last_value: the last value assigned or None
784 """
Ethan Furmanc16595e2016-09-10 23:36:59 -0700785 for last_value in reversed(last_values):
786 try:
787 return last_value + 1
788 except TypeError:
789 pass
790 else:
Ethan Furmanee47e5c2016-08-31 00:12:15 -0700791 return start
Ethan Furmanc16595e2016-09-10 23:36:59 -0700792
Ethan Furmanee47e5c2016-08-31 00:12:15 -0700793 @classmethod
794 def _missing_(cls, value):
Ethan Furmanc95ad7a2020-09-16 10:26:50 -0700795 return None
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700796
797 def __repr__(self):
798 return "<%s.%s: %r>" % (
Ethan Furman520ad572013-07-19 19:47:21 -0700799 self.__class__.__name__, self._name_, self._value_)
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700800
801 def __str__(self):
Ethan Furman520ad572013-07-19 19:47:21 -0700802 return "%s.%s" % (self.__class__.__name__, self._name_)
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700803
Ethan Furman388a3922013-08-12 06:51:41 -0700804 def __dir__(self):
Ethan Furman6d3dfee2020-12-08 12:26:56 -0800805 """
806 Returns all members and all public methods
807 """
Ethan Furman0ae550b2014-10-14 08:58:32 -0700808 added_behavior = [
809 m
810 for cls in self.__class__.mro()
811 for m in cls.__dict__
Ethan Furman354ecf12015-03-11 08:43:12 -0700812 if m[0] != '_' and m not in self._member_map_
Angelin BOOZ68526fe2020-09-21 15:11:06 +0200813 ] + [m for m in self.__dict__ if m[0] != '_']
Ethan Furmanec5f8eb2014-10-21 13:40:35 -0700814 return (['__class__', '__doc__', '__module__'] + added_behavior)
Ethan Furman388a3922013-08-12 06:51:41 -0700815
Ethan Furmanec15a822013-08-31 19:17:41 -0700816 def __format__(self, format_spec):
Ethan Furman6d3dfee2020-12-08 12:26:56 -0800817 """
818 Returns format using actual value type unless __str__ has been overridden.
819 """
Ethan Furmanec15a822013-08-31 19:17:41 -0700820 # mixed-in Enums should use the mixed-in type's __format__, otherwise
821 # we can get strange results with the Enum name showing up instead of
822 # the value
823
thatneat2f19e822019-07-04 11:28:37 -0700824 # pure Enum branch, or branch with __str__ explicitly overridden
Ethan Furman37440ee2020-12-08 11:14:10 -0800825 str_overridden = type(self).__str__ not in (Enum.__str__, Flag.__str__)
thatneat2f19e822019-07-04 11:28:37 -0700826 if self._member_type_ is object or str_overridden:
Ethan Furmanec15a822013-08-31 19:17:41 -0700827 cls = str
828 val = str(self)
829 # mix-in branch
830 else:
831 cls = self._member_type_
Ethan Furman0081f232014-09-16 17:31:23 -0700832 val = self._value_
Ethan Furmanec15a822013-08-31 19:17:41 -0700833 return cls.__format__(val, format_spec)
834
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700835 def __hash__(self):
Ethan Furman520ad572013-07-19 19:47:21 -0700836 return hash(self._name_)
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700837
Ethan Furmanca1b7942014-02-08 11:36:27 -0800838 def __reduce_ex__(self, proto):
Ethan Furmandc870522014-02-18 12:37:12 -0800839 return self.__class__, (self._value_, )
Ethan Furmanca1b7942014-02-08 11:36:27 -0800840
Ethan Furmanc314e602021-01-12 23:47:57 -0800841 # enum.property is used to provide access to the `name` and
842 # `value` attributes of enum members while keeping some measure of
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700843 # protection from modification, while still allowing for an enumeration
844 # to have members named `name` and `value`. This works because enumeration
Ethan Furmanc314e602021-01-12 23:47:57 -0800845 # members are not set directly on the enum class; they are kept in a
846 # separate structure, _member_map_, which is where enum.property looks for
847 # them
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700848
Ethan Furmanc314e602021-01-12 23:47:57 -0800849 @property
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700850 def name(self):
Ethan Furmanc850f342013-09-15 16:59:35 -0700851 """The name of the Enum member."""
Ethan Furman520ad572013-07-19 19:47:21 -0700852 return self._name_
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700853
Ethan Furmanc314e602021-01-12 23:47:57 -0800854 @property
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700855 def value(self):
Ethan Furmanc850f342013-09-15 16:59:35 -0700856 """The value of the Enum member."""
Ethan Furman520ad572013-07-19 19:47:21 -0700857 return self._value_
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700858
859
860class IntEnum(int, Enum):
Ethan Furman0063ff42020-09-21 17:23:13 -0700861 """
862 Enum where members are also (and must be) ints
863 """
864
865
866class StrEnum(str, Enum):
867 """
868 Enum where members are also (and must be) strings
869 """
870
871 def __new__(cls, *values):
872 if len(values) > 3:
873 raise TypeError('too many arguments for str(): %r' % (values, ))
874 if len(values) == 1:
875 # it must be a string
876 if not isinstance(values[0], str):
877 raise TypeError('%r is not a string' % (values[0], ))
878 if len(values) > 1:
879 # check that encoding argument is a string
880 if not isinstance(values[1], str):
881 raise TypeError('encoding must be a string, not %r' % (values[1], ))
882 if len(values) > 2:
883 # check that errors argument is a string
884 if not isinstance(values[2], str):
885 raise TypeError('errors must be a string, not %r' % (values[2], ))
886 value = str(*values)
887 member = str.__new__(cls, value)
888 member._value_ = value
889 return member
Ethan Furmanf24bb352013-07-18 17:05:39 -0700890
Ethan Furmand986d162020-09-22 13:00:07 -0700891 __str__ = str.__str__
892
Ethan Furmanefb13be2020-12-10 12:20:06 -0800893 def _generate_next_value_(name, start, count, last_values):
894 """
895 Return the lower-cased version of the member name.
896 """
897 return name.lower()
898
Ethan Furmanf24bb352013-07-18 17:05:39 -0700899
Ethan Furman24e837f2015-03-18 17:27:57 -0700900def _reduce_ex_by_name(self, proto):
901 return self.name
902
Ethan Furman65a5a472016-09-01 23:55:19 -0700903class Flag(Enum):
Ethan Furman6d3dfee2020-12-08 12:26:56 -0800904 """
905 Support for flags
906 """
Ethan Furmanc16595e2016-09-10 23:36:59 -0700907
908 def _generate_next_value_(name, start, count, last_values):
Ethan Furmanee47e5c2016-08-31 00:12:15 -0700909 """
910 Generate the next value when not given.
911
912 name: the name of the member
HongWeipengbb16fb22019-09-21 13:22:54 +0800913 start: the initial start value or None
Ethan Furmanee47e5c2016-08-31 00:12:15 -0700914 count: the number of existing members
915 last_value: the last value assigned or None
916 """
917 if not count:
918 return start if start is not None else 1
Ethan Furmanc16595e2016-09-10 23:36:59 -0700919 for last_value in reversed(last_values):
920 try:
921 high_bit = _high_bit(last_value)
922 break
Ethan Furman3515dcc2016-09-18 13:15:41 -0700923 except Exception:
Ethan Furmanc16595e2016-09-10 23:36:59 -0700924 raise TypeError('Invalid Flag value: %r' % last_value) from None
Ethan Furmanee47e5c2016-08-31 00:12:15 -0700925 return 2 ** (high_bit+1)
926
927 @classmethod
928 def _missing_(cls, value):
Ethan Furman6d3dfee2020-12-08 12:26:56 -0800929 """
930 Returns member (possibly creating it) if one can be found for value.
931 """
Ethan Furmanee47e5c2016-08-31 00:12:15 -0700932 original_value = value
933 if value < 0:
934 value = ~value
935 possible_member = cls._create_pseudo_member_(value)
Ethan Furmanee47e5c2016-08-31 00:12:15 -0700936 if original_value < 0:
937 possible_member = ~possible_member
938 return possible_member
939
940 @classmethod
941 def _create_pseudo_member_(cls, value):
Ethan Furman3515dcc2016-09-18 13:15:41 -0700942 """
943 Create a composite member iff value contains only members.
944 """
Ethan Furmanee47e5c2016-08-31 00:12:15 -0700945 pseudo_member = cls._value2member_map_.get(value, None)
946 if pseudo_member is None:
Ethan Furman3515dcc2016-09-18 13:15:41 -0700947 # verify all bits are accounted for
948 _, extra_flags = _decompose(cls, value)
949 if extra_flags:
Walter Dörwald323842c2019-07-18 20:37:13 +0200950 raise ValueError("%r is not a valid %s" % (value, cls.__qualname__))
Ethan Furman3515dcc2016-09-18 13:15:41 -0700951 # construct a singleton enum pseudo-member
Ethan Furmanee47e5c2016-08-31 00:12:15 -0700952 pseudo_member = object.__new__(cls)
953 pseudo_member._name_ = None
954 pseudo_member._value_ = value
Ethan Furman28cf6632017-01-24 12:12:06 -0800955 # use setdefault in case another thread already created a composite
956 # with this value
957 pseudo_member = cls._value2member_map_.setdefault(value, pseudo_member)
Ethan Furmanee47e5c2016-08-31 00:12:15 -0700958 return pseudo_member
959
Ethan Furmanee47e5c2016-08-31 00:12:15 -0700960 def __contains__(self, other):
Ethan Furman6d3dfee2020-12-08 12:26:56 -0800961 """
962 Returns True if self has at least the same flags set as other.
963 """
Ethan Furmanee47e5c2016-08-31 00:12:15 -0700964 if not isinstance(other, self.__class__):
Rahul Jha94306522018-09-10 23:51:04 +0530965 raise TypeError(
966 "unsupported operand type(s) for 'in': '%s' and '%s'" % (
967 type(other).__qualname__, self.__class__.__qualname__))
Ethan Furmanee47e5c2016-08-31 00:12:15 -0700968 return other._value_ & self._value_ == other._value_
969
Ethan Furman7219e272020-09-16 13:01:00 -0700970 def __iter__(self):
Ethan Furman6d3dfee2020-12-08 12:26:56 -0800971 """
972 Returns flags in decreasing value order.
973 """
Ethan Furman7219e272020-09-16 13:01:00 -0700974 members, extra_flags = _decompose(self.__class__, self.value)
975 return (m for m in members if m._value_ != 0)
976
Ethan Furmanee47e5c2016-08-31 00:12:15 -0700977 def __repr__(self):
978 cls = self.__class__
979 if self._name_ is not None:
980 return '<%s.%s: %r>' % (cls.__name__, self._name_, self._value_)
Ethan Furman3515dcc2016-09-18 13:15:41 -0700981 members, uncovered = _decompose(cls, self._value_)
Ethan Furman27682d22016-09-04 11:39:01 -0700982 return '<%s.%s: %r>' % (
983 cls.__name__,
984 '|'.join([str(m._name_ or m._value_) for m in members]),
985 self._value_,
986 )
Ethan Furmanee47e5c2016-08-31 00:12:15 -0700987
988 def __str__(self):
989 cls = self.__class__
990 if self._name_ is not None:
991 return '%s.%s' % (cls.__name__, self._name_)
Ethan Furman3515dcc2016-09-18 13:15:41 -0700992 members, uncovered = _decompose(cls, self._value_)
Ethan Furmanee47e5c2016-08-31 00:12:15 -0700993 if len(members) == 1 and members[0]._name_ is None:
994 return '%s.%r' % (cls.__name__, members[0]._value_)
995 else:
996 return '%s.%s' % (
997 cls.__name__,
998 '|'.join([str(m._name_ or m._value_) for m in members]),
999 )
1000
Ethan Furman25d94bb2016-09-02 16:32:32 -07001001 def __bool__(self):
1002 return bool(self._value_)
1003
Ethan Furmanee47e5c2016-08-31 00:12:15 -07001004 def __or__(self, other):
1005 if not isinstance(other, self.__class__):
1006 return NotImplemented
1007 return self.__class__(self._value_ | other._value_)
1008
1009 def __and__(self, other):
1010 if not isinstance(other, self.__class__):
1011 return NotImplemented
1012 return self.__class__(self._value_ & other._value_)
1013
1014 def __xor__(self, other):
1015 if not isinstance(other, self.__class__):
1016 return NotImplemented
1017 return self.__class__(self._value_ ^ other._value_)
1018
1019 def __invert__(self):
Ethan Furman3515dcc2016-09-18 13:15:41 -07001020 members, uncovered = _decompose(self.__class__, self._value_)
Serhiy Storchaka81108372017-09-26 00:55:55 +03001021 inverted = self.__class__(0)
1022 for m in self.__class__:
1023 if m not in members and not (m._value_ & self._value_):
1024 inverted = inverted | m
Ethan Furmanee47e5c2016-08-31 00:12:15 -07001025 return self.__class__(inverted)
1026
1027
Ethan Furman65a5a472016-09-01 23:55:19 -07001028class IntFlag(int, Flag):
Ethan Furman6d3dfee2020-12-08 12:26:56 -08001029 """
1030 Support for integer-based Flags
1031 """
Ethan Furmanee47e5c2016-08-31 00:12:15 -07001032
1033 @classmethod
Ethan Furman3515dcc2016-09-18 13:15:41 -07001034 def _missing_(cls, value):
Ethan Furman6d3dfee2020-12-08 12:26:56 -08001035 """
1036 Returns member (possibly creating it) if one can be found for value.
1037 """
Ethan Furman3515dcc2016-09-18 13:15:41 -07001038 if not isinstance(value, int):
Walter Dörwald323842c2019-07-18 20:37:13 +02001039 raise ValueError("%r is not a valid %s" % (value, cls.__qualname__))
Ethan Furman3515dcc2016-09-18 13:15:41 -07001040 new_member = cls._create_pseudo_member_(value)
1041 return new_member
1042
1043 @classmethod
Ethan Furmanee47e5c2016-08-31 00:12:15 -07001044 def _create_pseudo_member_(cls, value):
Ethan Furman6d3dfee2020-12-08 12:26:56 -08001045 """
1046 Create a composite member iff value contains only members.
1047 """
Ethan Furmanee47e5c2016-08-31 00:12:15 -07001048 pseudo_member = cls._value2member_map_.get(value, None)
1049 if pseudo_member is None:
Ethan Furman3515dcc2016-09-18 13:15:41 -07001050 need_to_create = [value]
1051 # get unaccounted for bits
1052 _, extra_flags = _decompose(cls, value)
1053 # timer = 10
1054 while extra_flags:
1055 # timer -= 1
1056 bit = _high_bit(extra_flags)
1057 flag_value = 2 ** bit
1058 if (flag_value not in cls._value2member_map_ and
1059 flag_value not in need_to_create
1060 ):
1061 need_to_create.append(flag_value)
1062 if extra_flags == -flag_value:
1063 extra_flags = 0
1064 else:
1065 extra_flags ^= flag_value
1066 for value in reversed(need_to_create):
1067 # construct singleton pseudo-members
1068 pseudo_member = int.__new__(cls, value)
1069 pseudo_member._name_ = None
1070 pseudo_member._value_ = value
Ethan Furman28cf6632017-01-24 12:12:06 -08001071 # use setdefault in case another thread already created a composite
1072 # with this value
1073 pseudo_member = cls._value2member_map_.setdefault(value, pseudo_member)
Ethan Furmanee47e5c2016-08-31 00:12:15 -07001074 return pseudo_member
1075
Ethan Furmanee47e5c2016-08-31 00:12:15 -07001076 def __or__(self, other):
1077 if not isinstance(other, (self.__class__, int)):
1078 return NotImplemented
Ethan Furman3515dcc2016-09-18 13:15:41 -07001079 result = self.__class__(self._value_ | self.__class__(other)._value_)
1080 return result
Ethan Furmanee47e5c2016-08-31 00:12:15 -07001081
1082 def __and__(self, other):
1083 if not isinstance(other, (self.__class__, int)):
1084 return NotImplemented
1085 return self.__class__(self._value_ & self.__class__(other)._value_)
1086
1087 def __xor__(self, other):
1088 if not isinstance(other, (self.__class__, int)):
1089 return NotImplemented
1090 return self.__class__(self._value_ ^ self.__class__(other)._value_)
1091
1092 __ror__ = __or__
1093 __rand__ = __and__
1094 __rxor__ = __xor__
1095
1096 def __invert__(self):
Ethan Furman3515dcc2016-09-18 13:15:41 -07001097 result = self.__class__(~self._value_)
1098 return result
Ethan Furmanee47e5c2016-08-31 00:12:15 -07001099
1100
1101def _high_bit(value):
Ethan Furman6d3dfee2020-12-08 12:26:56 -08001102 """
1103 returns index of highest bit, or -1 if value is zero or negative
1104 """
Ethan Furman3515dcc2016-09-18 13:15:41 -07001105 return value.bit_length() - 1
Ethan Furmanee47e5c2016-08-31 00:12:15 -07001106
Ethan Furmanf24bb352013-07-18 17:05:39 -07001107def unique(enumeration):
Ethan Furman6d3dfee2020-12-08 12:26:56 -08001108 """
1109 Class decorator for enumerations ensuring unique member values.
1110 """
Ethan Furmanf24bb352013-07-18 17:05:39 -07001111 duplicates = []
1112 for name, member in enumeration.__members__.items():
1113 if name != member.name:
1114 duplicates.append((name, member.name))
1115 if duplicates:
1116 alias_details = ', '.join(
1117 ["%s -> %s" % (alias, name) for (alias, name) in duplicates])
1118 raise ValueError('duplicate values found in %r: %s' %
1119 (enumeration, alias_details))
1120 return enumeration
Ethan Furman3515dcc2016-09-18 13:15:41 -07001121
1122def _decompose(flag, value):
Ethan Furman6d3dfee2020-12-08 12:26:56 -08001123 """
1124 Extract all members from the value.
1125 """
Ethan Furman3515dcc2016-09-18 13:15:41 -07001126 # _decompose is only called if the value is not named
1127 not_covered = value
1128 negative = value < 0
Ethan Furman3515dcc2016-09-18 13:15:41 -07001129 members = []
HongWeipeng0b41a922019-11-27 06:36:02 +08001130 for member in flag:
1131 member_value = member.value
Ethan Furman3515dcc2016-09-18 13:15:41 -07001132 if member_value and member_value & value == member_value:
1133 members.append(member)
1134 not_covered &= ~member_value
HongWeipeng0b41a922019-11-27 06:36:02 +08001135 if not negative:
1136 tmp = not_covered
1137 while tmp:
1138 flag_value = 2 ** _high_bit(tmp)
1139 if flag_value in flag._value2member_map_:
1140 members.append(flag._value2member_map_[flag_value])
1141 not_covered &= ~flag_value
1142 tmp &= ~flag_value
Ethan Furman3515dcc2016-09-18 13:15:41 -07001143 if not members and value in flag._value2member_map_:
1144 members.append(flag._value2member_map_[value])
1145 members.sort(key=lambda m: m._value_, reverse=True)
1146 if len(members) > 1 and members[0].value == value:
1147 # we have the breakdown, don't need the value member itself
1148 members.pop(0)
1149 return members, not_covered