blob: d4b11521ab27f38d28f2ed73c3fdca7205b7fffa [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 Furman7aaeb2a2021-01-25 14:26:19 -08003from builtins import property as _bltin_property, bin as _bltin_bin
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 Furman7aaeb2a2021-01-25 14:26:19 -080011 'FlagBoundary', 'STRICT', 'CONFORM', 'EJECT', 'KEEP',
Ethan Furmanc16595e2016-09-10 23:36:59 -070012 ]
Ethan Furman6b3d64a2013-06-14 16:55:46 -070013
14
Ethan Furman7aaeb2a2021-01-25 14:26:19 -080015# Dummy value for Enum and Flag as there are explicit checks for them
16# before they have been created.
17# This is also why there are checks in EnumMeta like `if Enum is not None`
18Enum = Flag = EJECT = None
19
Ethan Furman101e0742013-09-15 12:34:36 -070020def _is_descriptor(obj):
Ethan Furman6d3dfee2020-12-08 12:26:56 -080021 """
22 Returns True if obj is a descriptor, False otherwise.
23 """
Ethan Furman101e0742013-09-15 12:34:36 -070024 return (
25 hasattr(obj, '__get__') or
26 hasattr(obj, '__set__') or
Ethan Furman6d3dfee2020-12-08 12:26:56 -080027 hasattr(obj, '__delete__')
28 )
Ethan Furman101e0742013-09-15 12:34:36 -070029
Ethan Furman6b3d64a2013-06-14 16:55:46 -070030def _is_dunder(name):
Ethan Furman6d3dfee2020-12-08 12:26:56 -080031 """
32 Returns True if a __dunder__ name, False otherwise.
33 """
34 return (
35 len(name) > 4 and
Brennan D Baraban8b914d22019-03-03 14:09:11 -080036 name[:2] == name[-2:] == '__' and
37 name[2] != '_' and
Ethan Furman6d3dfee2020-12-08 12:26:56 -080038 name[-3] != '_'
39 )
Ethan Furman6b3d64a2013-06-14 16:55:46 -070040
41def _is_sunder(name):
Ethan Furman6d3dfee2020-12-08 12:26:56 -080042 """
43 Returns True if a _sunder_ name, False otherwise.
44 """
45 return (
46 len(name) > 2 and
Brennan D Baraban8b914d22019-03-03 14:09:11 -080047 name[0] == name[-1] == '_' and
Ethan Furman6b3d64a2013-06-14 16:55:46 -070048 name[1:2] != '_' and
Ethan Furman6d3dfee2020-12-08 12:26:56 -080049 name[-2:-1] != '_'
50 )
Ethan Furman6b3d64a2013-06-14 16:55:46 -070051
Ethan Furman7cf0aad2020-12-09 17:12:11 -080052def _is_private(cls_name, name):
53 # do not use `re` as `re` imports `enum`
54 pattern = '_%s__' % (cls_name, )
55 if (
56 len(name) >= 5
57 and name.startswith(pattern)
58 and name[len(pattern)] != '_'
59 and (name[-1] != '_' or name[-2] != '_')
60 ):
61 return True
62 else:
63 return False
64
Ethan Furman7aaeb2a2021-01-25 14:26:19 -080065def _is_single_bit(num):
66 """
67 True if only one bit set in num (should be an int)
68 """
69 if num == 0:
70 return False
71 num &= num - 1
72 return num == 0
73
Ethan Furmanc314e602021-01-12 23:47:57 -080074def _make_class_unpicklable(obj):
Ethan Furman6d3dfee2020-12-08 12:26:56 -080075 """
Ethan Furmanc314e602021-01-12 23:47:57 -080076 Make the given obj un-picklable.
77
78 obj should be either a dictionary, on an Enum
Ethan Furman6d3dfee2020-12-08 12:26:56 -080079 """
Ethan Furmanca1b7942014-02-08 11:36:27 -080080 def _break_on_call_reduce(self, proto):
Ethan Furman6b3d64a2013-06-14 16:55:46 -070081 raise TypeError('%r cannot be pickled' % self)
Ethan Furmanc314e602021-01-12 23:47:57 -080082 if isinstance(obj, dict):
83 obj['__reduce_ex__'] = _break_on_call_reduce
84 obj['__module__'] = '<unknown>'
85 else:
86 setattr(obj, '__reduce_ex__', _break_on_call_reduce)
87 setattr(obj, '__module__', '<unknown>')
Ethan Furman6b3d64a2013-06-14 16:55:46 -070088
Ethan Furman7aaeb2a2021-01-25 14:26:19 -080089def _iter_bits_lsb(num):
90 while num:
91 b = num & (~num + 1)
92 yield b
93 num ^= b
94
95def bin(num, max_bits=None):
96 """
97 Like built-in bin(), except negative values are represented in
98 twos-compliment, and the leading bit always indicates sign
99 (0=positive, 1=negative).
100
101 >>> bin(10)
102 '0b0 1010'
103 >>> bin(~10) # ~10 is -11
104 '0b1 0101'
105 """
106
107 ceiling = 2 ** (num).bit_length()
108 if num >= 0:
109 s = _bltin_bin(num + ceiling).replace('1', '0', 1)
110 else:
111 s = _bltin_bin(~num ^ (ceiling - 1) + ceiling)
112 sign = s[:3]
113 digits = s[3:]
114 if max_bits is not None:
115 if len(digits) < max_bits:
116 digits = (sign[-1] * max_bits + digits)[-max_bits:]
117 return "%s %s" % (sign, digits)
118
119
Ethan Furman3515dcc2016-09-18 13:15:41 -0700120_auto_null = object()
Ethan Furmanc16595e2016-09-10 23:36:59 -0700121class auto:
122 """
123 Instances are replaced with an appropriate value in Enum class suites.
124 """
Ethan Furman3515dcc2016-09-18 13:15:41 -0700125 value = _auto_null
Ethan Furmanc16595e2016-09-10 23:36:59 -0700126
Ethan Furmanc314e602021-01-12 23:47:57 -0800127class property(DynamicClassAttribute):
128 """
129 This is a descriptor, used to define attributes that act differently
130 when accessed through an enum member and through an enum class.
131 Instance access is the same as property(), but access to an attribute
132 through the enum class will instead look in the class' _member_map_ for
133 a corresponding enum member.
134 """
135
136 def __get__(self, instance, ownerclass=None):
137 if instance is None:
138 try:
139 return ownerclass._member_map_[self.name]
140 except KeyError:
Ethan Furman7aaeb2a2021-01-25 14:26:19 -0800141 raise AttributeError(
142 '%s: no attribute %r' % (ownerclass.__name__, self.name)
143 )
Ethan Furmanc314e602021-01-12 23:47:57 -0800144 else:
145 if self.fget is None:
Ethan Furman7aaeb2a2021-01-25 14:26:19 -0800146 raise AttributeError(
147 '%s: no attribute %r' % (ownerclass.__name__, self.name)
148 )
Ethan Furmanc314e602021-01-12 23:47:57 -0800149 else:
150 return self.fget(instance)
151
152 def __set__(self, instance, value):
153 if self.fset is None:
Ethan Furman7aaeb2a2021-01-25 14:26:19 -0800154 raise AttributeError(
155 "%s: cannot set attribute %r" % (self.clsname, self.name)
156 )
Ethan Furmanc314e602021-01-12 23:47:57 -0800157 else:
158 return self.fset(instance, value)
159
160 def __delete__(self, instance):
161 if self.fdel is None:
Ethan Furman7aaeb2a2021-01-25 14:26:19 -0800162 raise AttributeError(
163 "%s: cannot delete attribute %r" % (self.clsname, self.name)
164 )
Ethan Furmanc314e602021-01-12 23:47:57 -0800165 else:
166 return self.fdel(instance)
167
168 def __set_name__(self, ownerclass, name):
169 self.name = name
170 self.clsname = ownerclass.__name__
171
172
173class _proto_member:
174 """
175 intermediate step for enum members between class execution and final creation
176 """
177
178 def __init__(self, value):
179 self.value = value
180
181 def __set_name__(self, enum_class, member_name):
182 """
183 convert each quasi-member into an instance of the new enum class
184 """
185 # first step: remove ourself from enum_class
186 delattr(enum_class, member_name)
187 # second step: create member based on enum_class
188 value = self.value
189 if not isinstance(value, tuple):
190 args = (value, )
191 else:
192 args = value
193 if enum_class._member_type_ is tuple: # special case for tuple enums
194 args = (args, ) # wrap it one more time
195 if not enum_class._use_args_:
196 enum_member = enum_class._new_member_(enum_class)
197 if not hasattr(enum_member, '_value_'):
198 enum_member._value_ = value
199 else:
200 enum_member = enum_class._new_member_(enum_class, *args)
201 if not hasattr(enum_member, '_value_'):
202 if enum_class._member_type_ is object:
203 enum_member._value_ = value
204 else:
Ethan Furman7aaeb2a2021-01-25 14:26:19 -0800205 try:
206 enum_member._value_ = enum_class._member_type_(*args)
207 except Exception as exc:
208 raise TypeError(
209 '_value_ not set in __new__, unable to create it'
210 ) from None
Ethan Furmanc314e602021-01-12 23:47:57 -0800211 value = enum_member._value_
212 enum_member._name_ = member_name
213 enum_member.__objclass__ = enum_class
214 enum_member.__init__(*args)
Ethan Furman7aaeb2a2021-01-25 14:26:19 -0800215 enum_member._sort_order_ = len(enum_class._member_names_)
Ethan Furmanc314e602021-01-12 23:47:57 -0800216 # If another member with the same value was already defined, the
217 # new member becomes an alias to the existing one.
218 for name, canonical_member in enum_class._member_map_.items():
219 if canonical_member._value_ == enum_member._value_:
220 enum_member = canonical_member
221 break
222 else:
Ethan Furman7aaeb2a2021-01-25 14:26:19 -0800223 # this could still be an alias if the value is multi-bit and the
224 # class is a flag class
225 if (
226 Flag is None
227 or not issubclass(enum_class, Flag)
228 ):
229 # no other instances found, record this member in _member_names_
230 enum_class._member_names_.append(member_name)
231 elif (
232 Flag is not None
233 and issubclass(enum_class, Flag)
234 and _is_single_bit(value)
235 ):
236 # no other instances found, record this member in _member_names_
237 enum_class._member_names_.append(member_name)
Ethan Furmanc314e602021-01-12 23:47:57 -0800238 # get redirect in place before adding to _member_map_
239 # but check for other instances in parent classes first
240 need_override = False
241 descriptor = None
242 for base in enum_class.__mro__[1:]:
243 descriptor = base.__dict__.get(member_name)
244 if descriptor is not None:
245 if isinstance(descriptor, (property, DynamicClassAttribute)):
246 break
247 else:
248 need_override = True
249 # keep looking for an enum.property
250 if descriptor and not need_override:
251 # previous enum.property found, no further action needed
252 pass
253 else:
254 redirect = property()
255 redirect.__set_name__(enum_class, member_name)
256 if descriptor and need_override:
257 # previous enum.property found, but some other inherited attribute
258 # is in the way; copy fget, fset, fdel to this one
259 redirect.fget = descriptor.fget
260 redirect.fset = descriptor.fset
261 redirect.fdel = descriptor.fdel
262 setattr(enum_class, member_name, redirect)
263 # now add to _member_map_ (even aliases)
264 enum_class._member_map_[member_name] = enum_member
265 try:
266 # This may fail if value is not hashable. We can't add the value
267 # to the map, and by-value lookups for this value will be
268 # linear.
Ethan Furman7aaeb2a2021-01-25 14:26:19 -0800269 enum_class._value2member_map_.setdefault(value, enum_member)
Ethan Furmanc314e602021-01-12 23:47:57 -0800270 except TypeError:
271 pass
272
Ethan Furman101e0742013-09-15 12:34:36 -0700273
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700274class _EnumDict(dict):
Ethan Furman6d3dfee2020-12-08 12:26:56 -0800275 """
276 Track enum member order and ensure member names are not reused.
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700277
278 EnumMeta will use the names found in self._member_names as the
279 enumeration member names.
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700280 """
281 def __init__(self):
282 super().__init__()
283 self._member_names = []
Ethan Furmanc16595e2016-09-10 23:36:59 -0700284 self._last_values = []
Ethan Furmana4b1bb42018-01-22 07:56:37 -0800285 self._ignore = []
Ethan Onstottd9a43e22020-04-28 13:20:55 -0400286 self._auto_called = False
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700287
288 def __setitem__(self, key, value):
Ethan Furman6d3dfee2020-12-08 12:26:56 -0800289 """
290 Changes anything not dundered or not a descriptor.
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700291
292 If an enum member name is used twice, an error is raised; duplicate
293 values are not checked for.
294
295 Single underscore (sunder) names are reserved.
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700296 """
Ethan Furman7cf0aad2020-12-09 17:12:11 -0800297 if _is_private(self._cls_name, key):
298 # do nothing, name will be a normal attribute
299 pass
300 elif _is_sunder(key):
Ethan Furmanee47e5c2016-08-31 00:12:15 -0700301 if key not in (
Ethan Furman3515dcc2016-09-18 13:15:41 -0700302 '_order_', '_create_pseudo_member_',
Ethan Furmana4b1bb42018-01-22 07:56:37 -0800303 '_generate_next_value_', '_missing_', '_ignore_',
Ethan Furman7aaeb2a2021-01-25 14:26:19 -0800304 '_iter_member_', '_iter_member_by_value_', '_iter_member_by_def_',
Ethan Furmanee47e5c2016-08-31 00:12:15 -0700305 ):
Ethan Furman6d3dfee2020-12-08 12:26:56 -0800306 raise ValueError(
307 '_sunder_ names, such as %r, are reserved for future Enum use'
308 % (key, )
309 )
Ethan Furmanc16595e2016-09-10 23:36:59 -0700310 if key == '_generate_next_value_':
Ethan Onstottd9a43e22020-04-28 13:20:55 -0400311 # check if members already defined as auto()
312 if self._auto_called:
313 raise TypeError("_generate_next_value_ must be defined before members")
Ethan Furmanc16595e2016-09-10 23:36:59 -0700314 setattr(self, '_generate_next_value', value)
Ethan Furmana4b1bb42018-01-22 07:56:37 -0800315 elif key == '_ignore_':
316 if isinstance(value, str):
317 value = value.replace(',',' ').split()
318 else:
319 value = list(value)
320 self._ignore = value
321 already = set(value) & set(self._member_names)
322 if already:
Ethan Furman6d3dfee2020-12-08 12:26:56 -0800323 raise ValueError(
324 '_ignore_ cannot specify already set names: %r'
325 % (already, )
326 )
Ethan Furman101e0742013-09-15 12:34:36 -0700327 elif _is_dunder(key):
Ethan Furmane8e61272016-08-20 07:19:31 -0700328 if key == '__order__':
329 key = '_order_'
Ethan Furman101e0742013-09-15 12:34:36 -0700330 elif key in self._member_names:
331 # descriptor overwriting an enum?
Ethan Furmana6582872020-12-10 13:07:00 -0800332 raise TypeError('%r already defined as: %r' % (key, self[key]))
Ethan Furmana4b1bb42018-01-22 07:56:37 -0800333 elif key in self._ignore:
334 pass
Ethan Furman101e0742013-09-15 12:34:36 -0700335 elif not _is_descriptor(value):
336 if key in self:
337 # enum overwriting a descriptor?
Ethan Furmanee47e5c2016-08-31 00:12:15 -0700338 raise TypeError('%r already defined as: %r' % (key, self[key]))
Ethan Furmanc16595e2016-09-10 23:36:59 -0700339 if isinstance(value, auto):
Ethan Furman3515dcc2016-09-18 13:15:41 -0700340 if value.value == _auto_null:
Ethan Furman6d3dfee2020-12-08 12:26:56 -0800341 value.value = self._generate_next_value(
Ethan Furman7aaeb2a2021-01-25 14:26:19 -0800342 key, 1, len(self._member_names), self._last_values[:],
Ethan Furman6d3dfee2020-12-08 12:26:56 -0800343 )
Ethan Furmanfc23a942020-09-16 12:37:54 -0700344 self._auto_called = True
Ethan Furman3515dcc2016-09-18 13:15:41 -0700345 value = value.value
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700346 self._member_names.append(key)
Ethan Furmanc16595e2016-09-10 23:36:59 -0700347 self._last_values.append(value)
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700348 super().__setitem__(key, value)
349
Ethan Furmana6582872020-12-10 13:07:00 -0800350 def update(self, members, **more_members):
351 try:
352 for name in members.keys():
353 self[name] = members[name]
354 except AttributeError:
355 for name, value in members:
356 self[name] = value
357 for name, value in more_members.items():
358 self[name] = value
359
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700360
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700361class EnumMeta(type):
Ethan Furman6d3dfee2020-12-08 12:26:56 -0800362 """
363 Metaclass for Enum
364 """
Ethan Furman7aaeb2a2021-01-25 14:26:19 -0800365
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700366 @classmethod
Ethan Furman6ec0ade2020-12-24 10:05:02 -0800367 def __prepare__(metacls, cls, bases, **kwds):
Ethan Furman3064dbf2020-09-16 07:11:57 -0700368 # check that previous enum members do not exist
369 metacls._check_for_existing_members(cls, bases)
Ethan Furmanee47e5c2016-08-31 00:12:15 -0700370 # create the namespace dict
371 enum_dict = _EnumDict()
Ethan Furman7cf0aad2020-12-09 17:12:11 -0800372 enum_dict._cls_name = cls
Ethan Furmanee47e5c2016-08-31 00:12:15 -0700373 # inherit previous flags and _generate_next_value_ function
Ethan Furman3064dbf2020-09-16 07:11:57 -0700374 member_type, first_enum = metacls._get_mixins_(cls, bases)
Ethan Furmanee47e5c2016-08-31 00:12:15 -0700375 if first_enum is not None:
Ethan Furman6d3dfee2020-12-08 12:26:56 -0800376 enum_dict['_generate_next_value_'] = getattr(
377 first_enum, '_generate_next_value_', None,
378 )
Ethan Furmanee47e5c2016-08-31 00:12:15 -0700379 return enum_dict
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700380
Ethan Furman7aaeb2a2021-01-25 14:26:19 -0800381 def __new__(metacls, cls, bases, classdict, boundary=None, **kwds):
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700382 # an Enum class is final once enumeration items have been defined; it
383 # cannot be mixed with other types (int, float, etc.) if it has an
384 # inherited __new__ unless a new __new__ is defined (or the resulting
385 # class will fail).
Ethan Furmana4b1bb42018-01-22 07:56:37 -0800386 #
387 # remove any keys listed in _ignore_
388 classdict.setdefault('_ignore_', []).append('_ignore_')
389 ignore = classdict['_ignore_']
390 for key in ignore:
391 classdict.pop(key, None)
Ethan Furmanc314e602021-01-12 23:47:57 -0800392 #
393 # grab member names
394 member_names = classdict._member_names
395 #
396 # check for illegal enum names (any others?)
397 invalid_names = set(member_names) & {'mro', ''}
398 if invalid_names:
399 raise ValueError('Invalid enum member name: {0}'.format(
400 ','.join(invalid_names)))
401 #
402 # adjust the sunders
403 _order_ = classdict.pop('_order_', None)
404 # convert to normal dict
405 classdict = dict(classdict.items())
406 #
407 # data type of member and the controlling Enum class
Ethan Furman3064dbf2020-09-16 07:11:57 -0700408 member_type, first_enum = metacls._get_mixins_(cls, bases)
Ethan Furmanc2667362020-12-07 00:17:31 -0800409 __new__, save_new, use_args = metacls._find_new_(
410 classdict, member_type, first_enum,
411 )
Ethan Furmanc314e602021-01-12 23:47:57 -0800412 classdict['_new_member_'] = __new__
413 classdict['_use_args_'] = use_args
414 #
415 # convert future enum members into temporary _proto_members
Ethan Furman7aaeb2a2021-01-25 14:26:19 -0800416 # and record integer values in case this will be a Flag
417 flag_mask = 0
Ethan Furmanc314e602021-01-12 23:47:57 -0800418 for name in member_names:
Ethan Furman7aaeb2a2021-01-25 14:26:19 -0800419 value = classdict[name]
420 if isinstance(value, int):
421 flag_mask |= value
422 classdict[name] = _proto_member(value)
Ethan Furmanc314e602021-01-12 23:47:57 -0800423 #
Ethan Furman7aaeb2a2021-01-25 14:26:19 -0800424 # house-keeping structures
Ethan Furmanc314e602021-01-12 23:47:57 -0800425 classdict['_member_names_'] = []
426 classdict['_member_map_'] = {}
427 classdict['_value2member_map_'] = {}
428 classdict['_member_type_'] = member_type
429 #
Ethan Furman7aaeb2a2021-01-25 14:26:19 -0800430 # Flag structures (will be removed if final class is not a Flag
431 classdict['_boundary_'] = (
432 boundary
433 or getattr(first_enum, '_boundary_', None)
434 )
435 classdict['_flag_mask_'] = flag_mask
436 classdict['_all_bits_'] = 2 ** ((flag_mask).bit_length()) - 1
437 classdict['_inverted_'] = None
438 #
Ethan Furman2da95042014-03-03 12:42:52 -0800439 # If a custom type is mixed into the Enum, and it does not know how
440 # to pickle itself, pickle.dumps will succeed but pickle.loads will
441 # fail. Rather than have the error show up later and possibly far
442 # from the source, sabotage the pickle protocol for this class so
443 # that pickle.dumps also fails.
444 #
445 # However, if the new class implements its own __reduce_ex__, do not
446 # sabotage -- it's on them to make sure it works correctly. We use
447 # __reduce_ex__ instead of any of the others as it is preferred by
448 # pickle over __reduce__, and it handles all pickle protocols.
449 if '__reduce_ex__' not in classdict:
Ethan Furmandc870522014-02-18 12:37:12 -0800450 if member_type is not object:
451 methods = ('__getnewargs_ex__', '__getnewargs__',
452 '__reduce_ex__', '__reduce__')
Ethan Furman2da95042014-03-03 12:42:52 -0800453 if not any(m in member_type.__dict__ for m in methods):
Ethan Furmanc314e602021-01-12 23:47:57 -0800454 _make_class_unpicklable(classdict)
455 #
456 # create a default docstring if one has not been provided
457 if '__doc__' not in classdict:
458 classdict['__doc__'] = 'An enumeration.'
459 try:
460 exc = None
461 enum_class = super().__new__(metacls, cls, bases, classdict, **kwds)
462 except RuntimeError as e:
463 # any exceptions raised by member.__new__ will get converted to a
464 # RuntimeError, so get that original exception back and raise it instead
465 exc = e.__cause__ or e
466 if exc is not None:
467 raise exc
468 #
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700469 # double check that repr and friends are not the mixin's or various
470 # things break (such as pickle)
Ethan Furman22415ad2020-09-15 16:28:25 -0700471 # however, if the method is defined in the Enum itself, don't replace
472 # it
Ethan Furmandc870522014-02-18 12:37:12 -0800473 for name in ('__repr__', '__str__', '__format__', '__reduce_ex__'):
Ethan Furman22415ad2020-09-15 16:28:25 -0700474 if name in classdict:
475 continue
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700476 class_method = getattr(enum_class, name)
477 obj_method = getattr(member_type, name, None)
478 enum_method = getattr(first_enum, name, None)
479 if obj_method is not None and obj_method is class_method:
480 setattr(enum_class, name, enum_method)
Ethan Furmanc314e602021-01-12 23:47:57 -0800481 #
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700482 # replace any other __new__ with our own (as long as Enum is not None,
483 # anyway) -- again, this is to support pickle
484 if Enum is not None:
485 # if the user defined their own __new__, save it before it gets
486 # clobbered in case they subclass later
487 if save_new:
488 enum_class.__new_member__ = __new__
489 enum_class.__new__ = Enum.__new__
Ethan Furmanc314e602021-01-12 23:47:57 -0800490 #
Ethan Furmane8e61272016-08-20 07:19:31 -0700491 # py3 support for definition order (helps keep py2/py3 code in sync)
Ethan Furman7aaeb2a2021-01-25 14:26:19 -0800492 #
493 # _order_ checking is spread out into three/four steps
494 # - if enum_class is a Flag:
495 # - remove any non-single-bit flags from _order_
496 # - remove any aliases from _order_
497 # - check that _order_ and _member_names_ match
498 #
499 # step 1: ensure we have a list
Ethan Furmane8e61272016-08-20 07:19:31 -0700500 if _order_ is not None:
501 if isinstance(_order_, str):
502 _order_ = _order_.replace(',', ' ').split()
Ethan Furman7aaeb2a2021-01-25 14:26:19 -0800503 #
504 # remove Flag structures if final class is not a Flag
505 if (
506 Flag is None and cls != 'Flag'
507 or Flag is not None and not issubclass(enum_class, Flag)
508 ):
509 delattr(enum_class, '_boundary_')
510 delattr(enum_class, '_flag_mask_')
511 delattr(enum_class, '_all_bits_')
512 delattr(enum_class, '_inverted_')
513 elif Flag is not None and issubclass(enum_class, Flag):
514 # ensure _all_bits_ is correct and there are no missing flags
515 single_bit_total = 0
516 multi_bit_total = 0
517 for flag in enum_class._member_map_.values():
518 flag_value = flag._value_
519 if _is_single_bit(flag_value):
520 single_bit_total |= flag_value
521 else:
522 # multi-bit flags are considered aliases
523 multi_bit_total |= flag_value
524 if enum_class._boundary_ is not KEEP:
525 missed = list(_iter_bits_lsb(multi_bit_total & ~single_bit_total))
526 if missed:
527 raise TypeError(
528 'invalid Flag %r -- missing values: %s'
529 % (cls, ', '.join((str(i) for i in missed)))
530 )
531 enum_class._flag_mask_ = single_bit_total
532 #
533 # set correct __iter__
534 member_list = [m._value_ for m in enum_class]
535 if member_list != sorted(member_list):
536 enum_class._iter_member_ = enum_class._iter_member_by_def_
537 if _order_:
538 # _order_ step 2: remove any items from _order_ that are not single-bit
539 _order_ = [
540 o
541 for o in _order_
542 if o not in enum_class._member_map_ or _is_single_bit(enum_class[o]._value_)
543 ]
544 #
545 if _order_:
546 # _order_ step 3: remove aliases from _order_
547 _order_ = [
548 o
549 for o in _order_
550 if (
551 o not in enum_class._member_map_
552 or
553 (o in enum_class._member_map_ and o in enum_class._member_names_)
554 )]
555 # _order_ step 4: verify that _order_ and _member_names_ match
Ethan Furmane8e61272016-08-20 07:19:31 -0700556 if _order_ != enum_class._member_names_:
Ethan Furman7aaeb2a2021-01-25 14:26:19 -0800557 raise TypeError(
558 'member order does not match _order_:\n%r\n%r'
559 % (enum_class._member_names_, _order_)
560 )
Ethan Furmanc314e602021-01-12 23:47:57 -0800561 #
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700562 return enum_class
563
Ethan Furman5de67b12016-04-13 23:52:09 -0700564 def __bool__(self):
565 """
566 classes/types should always be True.
567 """
568 return True
569
Ethan Furman7aaeb2a2021-01-25 14:26:19 -0800570 def __call__(cls, value, names=None, *, module=None, qualname=None, type=None, start=1, boundary=None):
Ethan Furman6d3dfee2020-12-08 12:26:56 -0800571 """
572 Either returns an existing member, or creates a new enum class.
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700573
574 This method is used both when an enum class is given a value to match
575 to an enumeration member (i.e. Color(3)) and for the functional API
Ethan Furman23bb6f42016-11-21 09:22:05 -0800576 (i.e. Color = Enum('Color', names='RED GREEN BLUE')).
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700577
Ethan Furman2da95042014-03-03 12:42:52 -0800578 When used for the functional API:
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700579
Ethan Furman2da95042014-03-03 12:42:52 -0800580 `value` will be the name of the new class.
581
582 `names` should be either a string of white-space/comma delimited names
Ethan Furmand9925a12014-09-16 20:35:55 -0700583 (values will start at `start`), or an iterator/mapping of name, value pairs.
Ethan Furman2da95042014-03-03 12:42:52 -0800584
585 `module` should be set to the module this class is being created in;
586 if it is not set, an attempt to find that module will be made, but if
587 it fails the class will not be picklable.
588
589 `qualname` should be set to the actual location this class can be found
590 at in its module; by default it is set to the global scope. If this is
591 not correct, unpickling will fail in some circumstances.
592
593 `type`, if set, will be mixed in as the first base class.
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700594 """
595 if names is None: # simple value lookup
596 return cls.__new__(cls, value)
597 # otherwise, functional API: we're creating a new Enum type
Ethan Furman6d3dfee2020-12-08 12:26:56 -0800598 return cls._create_(
599 value,
600 names,
601 module=module,
602 qualname=qualname,
603 type=type,
604 start=start,
Ethan Furman7aaeb2a2021-01-25 14:26:19 -0800605 boundary=boundary,
Ethan Furman6d3dfee2020-12-08 12:26:56 -0800606 )
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700607
608 def __contains__(cls, member):
Rahul Jha94306522018-09-10 23:51:04 +0530609 if not isinstance(member, Enum):
610 raise TypeError(
611 "unsupported operand type(s) for 'in': '%s' and '%s'" % (
612 type(member).__qualname__, cls.__class__.__qualname__))
Ethan Furman0081f232014-09-16 17:31:23 -0700613 return isinstance(member, cls) and member._name_ in cls._member_map_
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700614
Ethan Furman64a99722013-09-22 16:18:19 -0700615 def __delattr__(cls, attr):
616 # nicer error message when someone tries to delete an attribute
617 # (see issue19025).
618 if attr in cls._member_map_:
Ethan Furman6d3dfee2020-12-08 12:26:56 -0800619 raise AttributeError("%s: cannot delete Enum member %r." % (cls.__name__, attr))
Ethan Furman64a99722013-09-22 16:18:19 -0700620 super().__delattr__(attr)
621
Ethan Furman388a3922013-08-12 06:51:41 -0700622 def __dir__(self):
Ethan Furman6d3dfee2020-12-08 12:26:56 -0800623 return (
624 ['__class__', '__doc__', '__members__', '__module__']
625 + self._member_names_
626 )
Ethan Furman388a3922013-08-12 06:51:41 -0700627
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700628 def __getattr__(cls, name):
Ethan Furman6d3dfee2020-12-08 12:26:56 -0800629 """
630 Return the enum member matching `name`
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700631
632 We use __getattr__ instead of descriptors or inserting into the enum
633 class' __dict__ in order to support `name` and `value` being both
634 properties for enum members (which live in the class' __dict__) and
635 enum members themselves.
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700636 """
637 if _is_dunder(name):
638 raise AttributeError(name)
639 try:
Ethan Furman520ad572013-07-19 19:47:21 -0700640 return cls._member_map_[name]
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700641 except KeyError:
642 raise AttributeError(name) from None
643
644 def __getitem__(cls, name):
Ethan Furman520ad572013-07-19 19:47:21 -0700645 return cls._member_map_[name]
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700646
647 def __iter__(cls):
Ethan Furman6d3dfee2020-12-08 12:26:56 -0800648 """
649 Returns members in definition order.
650 """
Ethan Furman520ad572013-07-19 19:47:21 -0700651 return (cls._member_map_[name] for name in cls._member_names_)
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700652
653 def __len__(cls):
Ethan Furman520ad572013-07-19 19:47:21 -0700654 return len(cls._member_names_)
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700655
Ethan Furmanc314e602021-01-12 23:47:57 -0800656 @_bltin_property
Ethan Furman2131a4a2013-09-14 18:11:24 -0700657 def __members__(cls):
Ethan Furman6d3dfee2020-12-08 12:26:56 -0800658 """
659 Returns a mapping of member name->value.
Ethan Furman2131a4a2013-09-14 18:11:24 -0700660
661 This mapping lists all enum members, including aliases. Note that this
662 is a read-only view of the internal mapping.
Ethan Furman2131a4a2013-09-14 18:11:24 -0700663 """
664 return MappingProxyType(cls._member_map_)
665
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700666 def __repr__(cls):
667 return "<enum %r>" % cls.__name__
668
Ethan Furman2131a4a2013-09-14 18:11:24 -0700669 def __reversed__(cls):
Ethan Furman6d3dfee2020-12-08 12:26:56 -0800670 """
671 Returns members in reverse definition order.
672 """
Ethan Furman2131a4a2013-09-14 18:11:24 -0700673 return (cls._member_map_[name] for name in reversed(cls._member_names_))
674
Ethan Furmanf203f2d2013-09-06 07:16:48 -0700675 def __setattr__(cls, name, value):
Ethan Furman6d3dfee2020-12-08 12:26:56 -0800676 """
677 Block attempts to reassign Enum members.
Ethan Furmanf203f2d2013-09-06 07:16:48 -0700678
679 A simple assignment to the class namespace only changes one of the
680 several possible ways to get an Enum member from the Enum class,
681 resulting in an inconsistent Enumeration.
Ethan Furmanf203f2d2013-09-06 07:16:48 -0700682 """
683 member_map = cls.__dict__.get('_member_map_', {})
684 if name in member_map:
685 raise AttributeError('Cannot reassign members.')
686 super().__setattr__(name, value)
687
Ethan Furman7aaeb2a2021-01-25 14:26:19 -0800688 def _create_(cls, class_name, names, *, module=None, qualname=None, type=None, start=1, boundary=None):
Ethan Furman6d3dfee2020-12-08 12:26:56 -0800689 """
690 Convenience method to create a new Enum class.
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700691
692 `names` can be:
693
694 * A string containing member names, separated either with spaces or
Ethan Furmand9925a12014-09-16 20:35:55 -0700695 commas. Values are incremented by 1 from `start`.
696 * An iterable of member names. Values are incremented by 1 from `start`.
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700697 * An iterable of (member name, value) pairs.
Ethan Furmand9925a12014-09-16 20:35:55 -0700698 * A mapping of member name -> value pairs.
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700699 """
700 metacls = cls.__class__
701 bases = (cls, ) if type is None else (type, cls)
Ethan Furman3064dbf2020-09-16 07:11:57 -0700702 _, first_enum = cls._get_mixins_(cls, bases)
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700703 classdict = metacls.__prepare__(class_name, bases)
704
705 # special processing needed for names?
706 if isinstance(names, str):
707 names = names.replace(',', ' ').split()
Dong-hee Nadcc8ce42017-06-22 01:52:32 +0900708 if isinstance(names, (tuple, list)) and names and isinstance(names[0], str):
Ethan Furmanee47e5c2016-08-31 00:12:15 -0700709 original_names, names = names, []
Ethan Furmanc16595e2016-09-10 23:36:59 -0700710 last_values = []
Ethan Furmanee47e5c2016-08-31 00:12:15 -0700711 for count, name in enumerate(original_names):
Ethan Furmanc16595e2016-09-10 23:36:59 -0700712 value = first_enum._generate_next_value_(name, start, count, last_values[:])
713 last_values.append(value)
714 names.append((name, value))
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700715
716 # Here, names is either an iterable of (name, value) or a mapping.
717 for item in names:
718 if isinstance(item, str):
719 member_name, member_value = item, names[item]
720 else:
721 member_name, member_value = item
722 classdict[member_name] = member_value
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700723
724 # TODO: replace the frame hack if a blessed way to know the calling
725 # module is ever developed
726 if module is None:
727 try:
728 module = sys._getframe(2).f_globals['__name__']
Pablo Galindo293dd232019-11-19 21:34:03 +0000729 except (AttributeError, ValueError, KeyError):
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700730 pass
731 if module is None:
Ethan Furmanc314e602021-01-12 23:47:57 -0800732 _make_class_unpicklable(classdict)
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700733 else:
Ethan Furmanc314e602021-01-12 23:47:57 -0800734 classdict['__module__'] = module
Ethan Furmanca1b7942014-02-08 11:36:27 -0800735 if qualname is not None:
Ethan Furmanc314e602021-01-12 23:47:57 -0800736 classdict['__qualname__'] = qualname
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700737
Ethan Furman7aaeb2a2021-01-25 14:26:19 -0800738 return metacls.__new__(metacls, class_name, bases, classdict, boundary=boundary)
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700739
Ethan Furman7aaeb2a2021-01-25 14:26:19 -0800740 def _convert_(cls, name, module, filter, source=None, boundary=None):
orlnub1230fb9fad2018-09-12 20:28:53 +0300741 """
742 Create a new Enum subclass that replaces a collection of global constants
743 """
744 # convert all constants from source (or module) that pass filter() to
745 # a new Enum called name, and export the enum and its members back to
746 # module;
747 # also, replace the __reduce_ex__ method so unpickling works in
748 # previous Python versions
749 module_globals = vars(sys.modules[module])
750 if source:
751 source = vars(source)
752 else:
753 source = module_globals
754 # _value2member_map_ is populated in the same order every time
755 # for a consistent reverse mapping of number to name when there
756 # are multiple names for the same number.
757 members = [
758 (name, value)
759 for name, value in source.items()
760 if filter(name)]
761 try:
762 # sort by value
763 members.sort(key=lambda t: (t[1], t[0]))
764 except TypeError:
765 # unless some values aren't comparable, in which case sort by name
766 members.sort(key=lambda t: t[0])
Ethan Furman7aaeb2a2021-01-25 14:26:19 -0800767 cls = cls(name, members, module=module, boundary=boundary or KEEP)
orlnub1230fb9fad2018-09-12 20:28:53 +0300768 cls.__reduce_ex__ = _reduce_ex_by_name
769 module_globals.update(cls.__members__)
770 module_globals[name] = cls
771 return cls
772
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700773 @staticmethod
Ethan Furman3064dbf2020-09-16 07:11:57 -0700774 def _check_for_existing_members(class_name, bases):
775 for chain in bases:
776 for base in chain.__mro__:
777 if issubclass(base, Enum) and base._member_names_:
Ethan Furman6d3dfee2020-12-08 12:26:56 -0800778 raise TypeError(
779 "%s: cannot extend enumeration %r"
780 % (class_name, base.__name__)
781 )
Ethan Furman3064dbf2020-09-16 07:11:57 -0700782
783 @staticmethod
784 def _get_mixins_(class_name, bases):
Ethan Furman6d3dfee2020-12-08 12:26:56 -0800785 """
786 Returns the type for creating enum members, and the first inherited
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700787 enum class.
788
789 bases: the tuple of bases that was given to __new__
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700790 """
791 if not bases:
792 return object, Enum
793
Ethan Furman5bdab642018-09-21 19:03:09 -0700794 def _find_data_type(bases):
Ethan Furmanbff01f32020-09-15 15:56:26 -0700795 data_types = []
Ethan Furman5bdab642018-09-21 19:03:09 -0700796 for chain in bases:
Ethan Furmanbff01f32020-09-15 15:56:26 -0700797 candidate = None
Ethan Furman5bdab642018-09-21 19:03:09 -0700798 for base in chain.__mro__:
799 if base is object:
800 continue
Ethan Furmanc2667362020-12-07 00:17:31 -0800801 elif issubclass(base, Enum):
802 if base._member_type_ is not object:
803 data_types.append(base._member_type_)
804 break
Ethan Furman5bdab642018-09-21 19:03:09 -0700805 elif '__new__' in base.__dict__:
Ethan Furmancd453852018-10-05 23:29:36 -0700806 if issubclass(base, Enum):
Ethan Furman5bdab642018-09-21 19:03:09 -0700807 continue
Ethan Furmanbff01f32020-09-15 15:56:26 -0700808 data_types.append(candidate or base)
809 break
Ethan Furmanc2667362020-12-07 00:17:31 -0800810 else:
Ethan Furmanbff01f32020-09-15 15:56:26 -0700811 candidate = base
812 if len(data_types) > 1:
Ethan Furman3064dbf2020-09-16 07:11:57 -0700813 raise TypeError('%r: too many data types: %r' % (class_name, data_types))
Ethan Furmanbff01f32020-09-15 15:56:26 -0700814 elif data_types:
815 return data_types[0]
816 else:
817 return None
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700818
Ethan Furman5bdab642018-09-21 19:03:09 -0700819 # ensure final parent class is an Enum derivative, find any concrete
820 # data type, and check that Enum has no members
821 first_enum = bases[-1]
822 if not issubclass(first_enum, Enum):
823 raise TypeError("new enumerations should be created as "
824 "`EnumName([mixin_type, ...] [data_type,] enum_type)`")
825 member_type = _find_data_type(bases) or object
826 if first_enum._member_names_:
827 raise TypeError("Cannot extend enumerations")
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700828 return member_type, first_enum
829
830 @staticmethod
831 def _find_new_(classdict, member_type, first_enum):
Ethan Furman6d3dfee2020-12-08 12:26:56 -0800832 """
833 Returns the __new__ to be used for creating the enum members.
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700834
835 classdict: the class dictionary given to __new__
836 member_type: the data type whose __new__ will be used by default
837 first_enum: enumeration to check for an overriding __new__
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700838 """
839 # now find the correct __new__, checking to see of one was defined
840 # by the user; also check earlier enum classes in case a __new__ was
841 # saved as __new_member__
842 __new__ = classdict.get('__new__', None)
843
844 # should __new__ be saved as __new_member__ later?
845 save_new = __new__ is not None
846
847 if __new__ is None:
848 # check all possibles for __new_member__ before falling back to
849 # __new__
850 for method in ('__new_member__', '__new__'):
851 for possible in (member_type, first_enum):
852 target = getattr(possible, method, None)
853 if target not in {
854 None,
855 None.__new__,
856 object.__new__,
857 Enum.__new__,
858 }:
859 __new__ = target
860 break
861 if __new__ is not None:
862 break
863 else:
864 __new__ = object.__new__
865
866 # if a non-object.__new__ is used then whatever value/tuple was
867 # assigned to the enum member name will be passed to __new__ and to the
868 # new enum member's __init__
869 if __new__ is object.__new__:
870 use_args = False
871 else:
872 use_args = True
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700873 return __new__, save_new, use_args
874
875
876class Enum(metaclass=EnumMeta):
Ethan Furman6d3dfee2020-12-08 12:26:56 -0800877 """
878 Generic enumeration.
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700879
880 Derive from this class to define new enumerations.
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700881 """
Ethan Furman7aaeb2a2021-01-25 14:26:19 -0800882
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700883 def __new__(cls, value):
884 # all enum instances are actually created during class construction
885 # without calling this method; this method is called by the metaclass'
886 # __call__ (i.e. Color(3) ), and by pickle
887 if type(value) is cls:
Ethan Furman23bb6f42016-11-21 09:22:05 -0800888 # For lookups like Color(Color.RED)
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700889 return value
890 # by-value search for a matching enum member
891 # see if it's in the reverse mapping (for hashable values)
Ethan Furman2aa27322013-07-19 19:35:56 -0700892 try:
Andrew Svetlov34ae04f2018-12-26 20:45:33 +0200893 return cls._value2member_map_[value]
894 except KeyError:
895 # Not found, no need to do long O(n) search
896 pass
Ethan Furman2aa27322013-07-19 19:35:56 -0700897 except TypeError:
898 # not there, now do long search -- O(n) behavior
Ethan Furman520ad572013-07-19 19:47:21 -0700899 for member in cls._member_map_.values():
Ethan Furman0081f232014-09-16 17:31:23 -0700900 if member._value_ == value:
Ethan Furman2aa27322013-07-19 19:35:56 -0700901 return member
Ethan Furmanee47e5c2016-08-31 00:12:15 -0700902 # still not found -- try _missing_ hook
Ethan Furman019f0a02018-09-12 11:43:34 -0700903 try:
904 exc = None
905 result = cls._missing_(value)
906 except Exception as e:
907 exc = e
908 result = None
909 if isinstance(result, cls):
910 return result
Ethan Furman7aaeb2a2021-01-25 14:26:19 -0800911 elif (
912 Flag is not None and issubclass(cls, Flag)
913 and cls._boundary_ is EJECT and isinstance(result, int)
914 ):
915 return result
Ethan Furman019f0a02018-09-12 11:43:34 -0700916 else:
Walter Dörwald323842c2019-07-18 20:37:13 +0200917 ve_exc = ValueError("%r is not a valid %s" % (value, cls.__qualname__))
Ethan Furman019f0a02018-09-12 11:43:34 -0700918 if result is None and exc is None:
919 raise ve_exc
920 elif exc is None:
921 exc = TypeError(
922 'error in %s._missing_: returned %r instead of None or a valid member'
923 % (cls.__name__, result)
924 )
Ethan Furman7aaeb2a2021-01-25 14:26:19 -0800925 if not isinstance(exc, ValueError):
926 exc.__context__ = ve_exc
Ethan Furman019f0a02018-09-12 11:43:34 -0700927 raise exc
Ethan Furmanee47e5c2016-08-31 00:12:15 -0700928
Ethan Furmanc16595e2016-09-10 23:36:59 -0700929 def _generate_next_value_(name, start, count, last_values):
Ethan Furman6d3dfee2020-12-08 12:26:56 -0800930 """
931 Generate the next value when not given.
932
933 name: the name of the member
934 start: the initial start value or None
935 count: the number of existing members
936 last_value: the last value assigned or None
937 """
Ethan Furmanc16595e2016-09-10 23:36:59 -0700938 for last_value in reversed(last_values):
939 try:
940 return last_value + 1
941 except TypeError:
942 pass
943 else:
Ethan Furmanee47e5c2016-08-31 00:12:15 -0700944 return start
Ethan Furmanc16595e2016-09-10 23:36:59 -0700945
Ethan Furmanee47e5c2016-08-31 00:12:15 -0700946 @classmethod
947 def _missing_(cls, value):
Ethan Furmanc95ad7a2020-09-16 10:26:50 -0700948 return None
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700949
950 def __repr__(self):
951 return "<%s.%s: %r>" % (
Ethan Furman520ad572013-07-19 19:47:21 -0700952 self.__class__.__name__, self._name_, self._value_)
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700953
954 def __str__(self):
Ethan Furman520ad572013-07-19 19:47:21 -0700955 return "%s.%s" % (self.__class__.__name__, self._name_)
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700956
Ethan Furman388a3922013-08-12 06:51:41 -0700957 def __dir__(self):
Ethan Furman6d3dfee2020-12-08 12:26:56 -0800958 """
959 Returns all members and all public methods
960 """
Ethan Furman0ae550b2014-10-14 08:58:32 -0700961 added_behavior = [
962 m
963 for cls in self.__class__.mro()
964 for m in cls.__dict__
Ethan Furman354ecf12015-03-11 08:43:12 -0700965 if m[0] != '_' and m not in self._member_map_
Angelin BOOZ68526fe2020-09-21 15:11:06 +0200966 ] + [m for m in self.__dict__ if m[0] != '_']
Ethan Furmanec5f8eb2014-10-21 13:40:35 -0700967 return (['__class__', '__doc__', '__module__'] + added_behavior)
Ethan Furman388a3922013-08-12 06:51:41 -0700968
Ethan Furmanec15a822013-08-31 19:17:41 -0700969 def __format__(self, format_spec):
Ethan Furman6d3dfee2020-12-08 12:26:56 -0800970 """
971 Returns format using actual value type unless __str__ has been overridden.
972 """
Ethan Furmanec15a822013-08-31 19:17:41 -0700973 # mixed-in Enums should use the mixed-in type's __format__, otherwise
974 # we can get strange results with the Enum name showing up instead of
975 # the value
976
thatneat2f19e822019-07-04 11:28:37 -0700977 # pure Enum branch, or branch with __str__ explicitly overridden
Ethan Furman37440ee2020-12-08 11:14:10 -0800978 str_overridden = type(self).__str__ not in (Enum.__str__, Flag.__str__)
thatneat2f19e822019-07-04 11:28:37 -0700979 if self._member_type_ is object or str_overridden:
Ethan Furmanec15a822013-08-31 19:17:41 -0700980 cls = str
981 val = str(self)
982 # mix-in branch
983 else:
984 cls = self._member_type_
Ethan Furman0081f232014-09-16 17:31:23 -0700985 val = self._value_
Ethan Furmanec15a822013-08-31 19:17:41 -0700986 return cls.__format__(val, format_spec)
987
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700988 def __hash__(self):
Ethan Furman520ad572013-07-19 19:47:21 -0700989 return hash(self._name_)
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700990
Ethan Furmanca1b7942014-02-08 11:36:27 -0800991 def __reduce_ex__(self, proto):
Ethan Furmandc870522014-02-18 12:37:12 -0800992 return self.__class__, (self._value_, )
Ethan Furmanca1b7942014-02-08 11:36:27 -0800993
Ethan Furmanc314e602021-01-12 23:47:57 -0800994 # enum.property is used to provide access to the `name` and
995 # `value` attributes of enum members while keeping some measure of
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700996 # protection from modification, while still allowing for an enumeration
997 # to have members named `name` and `value`. This works because enumeration
Ethan Furmanc314e602021-01-12 23:47:57 -0800998 # members are not set directly on the enum class; they are kept in a
999 # separate structure, _member_map_, which is where enum.property looks for
1000 # them
Ethan Furman6b3d64a2013-06-14 16:55:46 -07001001
Ethan Furmanc314e602021-01-12 23:47:57 -08001002 @property
Ethan Furman6b3d64a2013-06-14 16:55:46 -07001003 def name(self):
Ethan Furmanc850f342013-09-15 16:59:35 -07001004 """The name of the Enum member."""
Ethan Furman520ad572013-07-19 19:47:21 -07001005 return self._name_
Ethan Furman6b3d64a2013-06-14 16:55:46 -07001006
Ethan Furmanc314e602021-01-12 23:47:57 -08001007 @property
Ethan Furman6b3d64a2013-06-14 16:55:46 -07001008 def value(self):
Ethan Furmanc850f342013-09-15 16:59:35 -07001009 """The value of the Enum member."""
Ethan Furman520ad572013-07-19 19:47:21 -07001010 return self._value_
Ethan Furman6b3d64a2013-06-14 16:55:46 -07001011
1012
1013class IntEnum(int, Enum):
Ethan Furman0063ff42020-09-21 17:23:13 -07001014 """
1015 Enum where members are also (and must be) ints
1016 """
1017
1018
1019class StrEnum(str, Enum):
1020 """
1021 Enum where members are also (and must be) strings
1022 """
1023
1024 def __new__(cls, *values):
1025 if len(values) > 3:
1026 raise TypeError('too many arguments for str(): %r' % (values, ))
1027 if len(values) == 1:
1028 # it must be a string
1029 if not isinstance(values[0], str):
1030 raise TypeError('%r is not a string' % (values[0], ))
Ethan Furman7aaeb2a2021-01-25 14:26:19 -08001031 if len(values) >= 2:
Ethan Furman0063ff42020-09-21 17:23:13 -07001032 # check that encoding argument is a string
1033 if not isinstance(values[1], str):
1034 raise TypeError('encoding must be a string, not %r' % (values[1], ))
Ethan Furman7aaeb2a2021-01-25 14:26:19 -08001035 if len(values) == 3:
1036 # check that errors argument is a string
1037 if not isinstance(values[2], str):
1038 raise TypeError('errors must be a string, not %r' % (values[2]))
Ethan Furman0063ff42020-09-21 17:23:13 -07001039 value = str(*values)
1040 member = str.__new__(cls, value)
1041 member._value_ = value
1042 return member
Ethan Furmanf24bb352013-07-18 17:05:39 -07001043
Ethan Furmand986d162020-09-22 13:00:07 -07001044 __str__ = str.__str__
1045
Ethan Furmanefb13be2020-12-10 12:20:06 -08001046 def _generate_next_value_(name, start, count, last_values):
1047 """
1048 Return the lower-cased version of the member name.
1049 """
1050 return name.lower()
1051
Ethan Furmanf24bb352013-07-18 17:05:39 -07001052
Ethan Furman24e837f2015-03-18 17:27:57 -07001053def _reduce_ex_by_name(self, proto):
1054 return self.name
1055
Ethan Furman7aaeb2a2021-01-25 14:26:19 -08001056class FlagBoundary(StrEnum):
1057 """
1058 control how out of range values are handled
1059 "strict" -> error is raised [default for Flag]
1060 "conform" -> extra bits are discarded
1061 "eject" -> lose flag status [default for IntFlag]
1062 "keep" -> keep flag status and all bits
1063 """
1064 STRICT = auto()
1065 CONFORM = auto()
1066 EJECT = auto()
1067 KEEP = auto()
1068STRICT, CONFORM, EJECT, KEEP = FlagBoundary
1069
1070
1071class Flag(Enum, boundary=STRICT):
Ethan Furman6d3dfee2020-12-08 12:26:56 -08001072 """
1073 Support for flags
1074 """
Ethan Furmanc16595e2016-09-10 23:36:59 -07001075
1076 def _generate_next_value_(name, start, count, last_values):
Ethan Furmanee47e5c2016-08-31 00:12:15 -07001077 """
1078 Generate the next value when not given.
1079
1080 name: the name of the member
HongWeipengbb16fb22019-09-21 13:22:54 +08001081 start: the initial start value or None
Ethan Furmanee47e5c2016-08-31 00:12:15 -07001082 count: the number of existing members
1083 last_value: the last value assigned or None
1084 """
1085 if not count:
1086 return start if start is not None else 1
Ethan Furman7aaeb2a2021-01-25 14:26:19 -08001087 last_value = max(last_values)
1088 try:
1089 high_bit = _high_bit(last_value)
1090 except Exception:
1091 raise TypeError('Invalid Flag value: %r' % last_value) from None
Ethan Furmanee47e5c2016-08-31 00:12:15 -07001092 return 2 ** (high_bit+1)
1093
1094 @classmethod
Ethan Furman7aaeb2a2021-01-25 14:26:19 -08001095 def _iter_member_by_value_(cls, value):
1096 """
1097 Extract all members from the value in definition (i.e. increasing value) order.
1098 """
1099 for val in _iter_bits_lsb(value & cls._flag_mask_):
1100 yield cls._value2member_map_.get(val)
1101
1102 _iter_member_ = _iter_member_by_value_
1103
1104 @classmethod
1105 def _iter_member_by_def_(cls, value):
1106 """
1107 Extract all members from the value in definition order.
1108 """
1109 yield from sorted(
1110 cls._iter_member_by_value_(value),
1111 key=lambda m: m._sort_order_,
1112 )
1113
1114 @classmethod
Ethan Furmanee47e5c2016-08-31 00:12:15 -07001115 def _missing_(cls, value):
Ethan Furman6d3dfee2020-12-08 12:26:56 -08001116 """
Ethan Furman3515dcc2016-09-18 13:15:41 -07001117 Create a composite member iff value contains only members.
1118 """
Ethan Furman7aaeb2a2021-01-25 14:26:19 -08001119 if not isinstance(value, int):
1120 raise ValueError(
1121 "%r is not a valid %s" % (value, cls.__qualname__)
1122 )
1123 # check boundaries
1124 # - value must be in range (e.g. -16 <-> +15, i.e. ~15 <-> 15)
1125 # - value must not include any skipped flags (e.g. if bit 2 is not
1126 # defined, then 0d10 is invalid)
1127 flag_mask = cls._flag_mask_
1128 all_bits = cls._all_bits_
1129 neg_value = None
1130 if (
1131 not ~all_bits <= value <= all_bits
1132 or value & (all_bits ^ flag_mask)
1133 ):
1134 if cls._boundary_ is STRICT:
1135 max_bits = max(value.bit_length(), flag_mask.bit_length())
1136 raise ValueError(
1137 "%s: invalid value: %r\n given %s\n allowed %s" % (
1138 cls.__name__, value, bin(value, max_bits), bin(flag_mask, max_bits),
1139 ))
1140 elif cls._boundary_ is CONFORM:
1141 value = value & flag_mask
1142 elif cls._boundary_ is EJECT:
1143 return value
1144 elif cls._boundary_ is KEEP:
1145 if value < 0:
1146 value = (
1147 max(all_bits+1, 2**(value.bit_length()))
1148 + value
1149 )
1150 else:
1151 raise ValueError(
1152 'unknown flag boundary: %r' % (cls._boundary_, )
1153 )
1154 if value < 0:
1155 neg_value = value
1156 value = all_bits + 1 + value
1157 # get members and unknown
1158 unknown = value & ~flag_mask
1159 member_value = value & flag_mask
1160 if unknown and cls._boundary_ is not KEEP:
1161 raise ValueError(
1162 '%s(%r) --> unknown values %r [%s]'
1163 % (cls.__name__, value, unknown, bin(unknown))
1164 )
1165 # normal Flag?
1166 __new__ = getattr(cls, '__new_member__', None)
1167 if cls._member_type_ is object and not __new__:
Ethan Furman3515dcc2016-09-18 13:15:41 -07001168 # construct a singleton enum pseudo-member
Ethan Furmanee47e5c2016-08-31 00:12:15 -07001169 pseudo_member = object.__new__(cls)
Ethan Furman7aaeb2a2021-01-25 14:26:19 -08001170 else:
1171 pseudo_member = (__new__ or cls._member_type_.__new__)(cls, value)
1172 if not hasattr(pseudo_member, 'value'):
Ethan Furmanee47e5c2016-08-31 00:12:15 -07001173 pseudo_member._value_ = value
Ethan Furman7aaeb2a2021-01-25 14:26:19 -08001174 if member_value:
1175 pseudo_member._name_ = '|'.join([
1176 m._name_ for m in cls._iter_member_(member_value)
1177 ])
1178 if unknown:
1179 pseudo_member._name_ += '|0x%x' % unknown
1180 else:
1181 pseudo_member._name_ = None
1182 # use setdefault in case another thread already created a composite
1183 # with this value, but only if all members are known
1184 # note: zero is a special case -- add it
1185 if not unknown:
Ethan Furman28cf6632017-01-24 12:12:06 -08001186 pseudo_member = cls._value2member_map_.setdefault(value, pseudo_member)
Ethan Furman7aaeb2a2021-01-25 14:26:19 -08001187 if neg_value is not None:
1188 cls._value2member_map_[neg_value] = pseudo_member
Ethan Furmanee47e5c2016-08-31 00:12:15 -07001189 return pseudo_member
1190
Ethan Furmanee47e5c2016-08-31 00:12:15 -07001191 def __contains__(self, other):
Ethan Furman6d3dfee2020-12-08 12:26:56 -08001192 """
1193 Returns True if self has at least the same flags set as other.
1194 """
Ethan Furmanee47e5c2016-08-31 00:12:15 -07001195 if not isinstance(other, self.__class__):
Rahul Jha94306522018-09-10 23:51:04 +05301196 raise TypeError(
1197 "unsupported operand type(s) for 'in': '%s' and '%s'" % (
1198 type(other).__qualname__, self.__class__.__qualname__))
Ethan Furman7aaeb2a2021-01-25 14:26:19 -08001199 if other._value_ == 0 or self._value_ == 0:
1200 return False
Ethan Furmanee47e5c2016-08-31 00:12:15 -07001201 return other._value_ & self._value_ == other._value_
1202
Ethan Furman7219e272020-09-16 13:01:00 -07001203 def __iter__(self):
Ethan Furman6d3dfee2020-12-08 12:26:56 -08001204 """
Ethan Furman7aaeb2a2021-01-25 14:26:19 -08001205 Returns flags in definition order.
Ethan Furman6d3dfee2020-12-08 12:26:56 -08001206 """
Ethan Furman7aaeb2a2021-01-25 14:26:19 -08001207 yield from self._iter_member_(self._value_)
1208
1209 def __len__(self):
1210 return self._value_.bit_count()
Ethan Furman7219e272020-09-16 13:01:00 -07001211
Ethan Furmanee47e5c2016-08-31 00:12:15 -07001212 def __repr__(self):
1213 cls = self.__class__
1214 if self._name_ is not None:
1215 return '<%s.%s: %r>' % (cls.__name__, self._name_, self._value_)
Ethan Furman7aaeb2a2021-01-25 14:26:19 -08001216 else:
1217 # only zero is unnamed by default
1218 return '<%s: %r>' % (cls.__name__, self._value_)
Ethan Furmanee47e5c2016-08-31 00:12:15 -07001219
1220 def __str__(self):
1221 cls = self.__class__
1222 if self._name_ is not None:
1223 return '%s.%s' % (cls.__name__, self._name_)
Ethan Furmanee47e5c2016-08-31 00:12:15 -07001224 else:
Ethan Furman7aaeb2a2021-01-25 14:26:19 -08001225 return '%s(%s)' % (cls.__name__, self._value_)
Ethan Furmanee47e5c2016-08-31 00:12:15 -07001226
Ethan Furman25d94bb2016-09-02 16:32:32 -07001227 def __bool__(self):
1228 return bool(self._value_)
1229
Ethan Furmanee47e5c2016-08-31 00:12:15 -07001230 def __or__(self, other):
1231 if not isinstance(other, self.__class__):
1232 return NotImplemented
1233 return self.__class__(self._value_ | other._value_)
1234
1235 def __and__(self, other):
1236 if not isinstance(other, self.__class__):
1237 return NotImplemented
1238 return self.__class__(self._value_ & other._value_)
1239
1240 def __xor__(self, other):
1241 if not isinstance(other, self.__class__):
1242 return NotImplemented
1243 return self.__class__(self._value_ ^ other._value_)
1244
1245 def __invert__(self):
Ethan Furman7aaeb2a2021-01-25 14:26:19 -08001246 if self._inverted_ is None:
1247 if self._boundary_ is KEEP:
1248 # use all bits
1249 self._inverted_ = self.__class__(~self._value_)
1250 else:
1251 # calculate flags not in this member
1252 self._inverted_ = self.__class__(self._flag_mask_ ^ self._value_)
1253 self._inverted_._inverted_ = self
1254 return self._inverted_
Ethan Furmanee47e5c2016-08-31 00:12:15 -07001255
1256
Ethan Furman7aaeb2a2021-01-25 14:26:19 -08001257class IntFlag(int, Flag, boundary=EJECT):
Ethan Furman6d3dfee2020-12-08 12:26:56 -08001258 """
1259 Support for integer-based Flags
1260 """
Ethan Furmanee47e5c2016-08-31 00:12:15 -07001261
Ethan Furmanee47e5c2016-08-31 00:12:15 -07001262 def __or__(self, other):
Ethan Furman7aaeb2a2021-01-25 14:26:19 -08001263 if isinstance(other, self.__class__):
1264 other = other._value_
1265 elif isinstance(other, int):
1266 other = other
1267 else:
Ethan Furmanee47e5c2016-08-31 00:12:15 -07001268 return NotImplemented
Ethan Furman7aaeb2a2021-01-25 14:26:19 -08001269 value = self._value_
1270 return self.__class__(value | other)
Ethan Furmanee47e5c2016-08-31 00:12:15 -07001271
1272 def __and__(self, other):
Ethan Furman7aaeb2a2021-01-25 14:26:19 -08001273 if isinstance(other, self.__class__):
1274 other = other._value_
1275 elif isinstance(other, int):
1276 other = other
1277 else:
Ethan Furmanee47e5c2016-08-31 00:12:15 -07001278 return NotImplemented
Ethan Furman7aaeb2a2021-01-25 14:26:19 -08001279 value = self._value_
1280 return self.__class__(value & other)
Ethan Furmanee47e5c2016-08-31 00:12:15 -07001281
1282 def __xor__(self, other):
Ethan Furman7aaeb2a2021-01-25 14:26:19 -08001283 if isinstance(other, self.__class__):
1284 other = other._value_
1285 elif isinstance(other, int):
1286 other = other
1287 else:
Ethan Furmanee47e5c2016-08-31 00:12:15 -07001288 return NotImplemented
Ethan Furman7aaeb2a2021-01-25 14:26:19 -08001289 value = self._value_
1290 return self.__class__(value ^ other)
Ethan Furmanee47e5c2016-08-31 00:12:15 -07001291
1292 __ror__ = __or__
1293 __rand__ = __and__
1294 __rxor__ = __xor__
Ethan Furman7aaeb2a2021-01-25 14:26:19 -08001295 __invert__ = Flag.__invert__
Ethan Furmanee47e5c2016-08-31 00:12:15 -07001296
1297def _high_bit(value):
Ethan Furman6d3dfee2020-12-08 12:26:56 -08001298 """
1299 returns index of highest bit, or -1 if value is zero or negative
1300 """
Ethan Furman3515dcc2016-09-18 13:15:41 -07001301 return value.bit_length() - 1
Ethan Furmanee47e5c2016-08-31 00:12:15 -07001302
Ethan Furmanf24bb352013-07-18 17:05:39 -07001303def unique(enumeration):
Ethan Furman6d3dfee2020-12-08 12:26:56 -08001304 """
1305 Class decorator for enumerations ensuring unique member values.
1306 """
Ethan Furmanf24bb352013-07-18 17:05:39 -07001307 duplicates = []
1308 for name, member in enumeration.__members__.items():
1309 if name != member.name:
1310 duplicates.append((name, member.name))
1311 if duplicates:
1312 alias_details = ', '.join(
1313 ["%s -> %s" % (alias, name) for (alias, name) in duplicates])
1314 raise ValueError('duplicate values found in %r: %s' %
1315 (enumeration, alias_details))
1316 return enumeration
Ethan Furman3515dcc2016-09-18 13:15:41 -07001317
Ethan Furman7aaeb2a2021-01-25 14:26:19 -08001318def _power_of_two(value):
1319 if value < 1:
1320 return False
1321 return value == 2 ** _high_bit(value)