blob: f31779baa0d65a5beeb4a8472bbb20358d1aa8a2 [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__ = [
Ethan Furmanb7751062021-03-30 21:17:26 -07007 'EnumType', '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 Furmanb7751062021-03-30 21:17:26 -070012 'global_flag_repr', 'global_enum_repr', 'global_enum',
Ethan Furmanc16595e2016-09-10 23:36:59 -070013 ]
Ethan Furman6b3d64a2013-06-14 16:55:46 -070014
15
Ethan Furman7aaeb2a2021-01-25 14:26:19 -080016# Dummy value for Enum and Flag as there are explicit checks for them
17# before they have been created.
Ethan Furmanb7751062021-03-30 21:17:26 -070018# This is also why there are checks in EnumType like `if Enum is not None`
Ethan Furman7aaeb2a2021-01-25 14:26:19 -080019Enum = Flag = EJECT = None
20
Ethan Furman101e0742013-09-15 12:34:36 -070021def _is_descriptor(obj):
Ethan Furman6d3dfee2020-12-08 12:26:56 -080022 """
23 Returns True if obj is a descriptor, False otherwise.
24 """
Ethan Furman101e0742013-09-15 12:34:36 -070025 return (
26 hasattr(obj, '__get__') or
27 hasattr(obj, '__set__') or
Ethan Furman6d3dfee2020-12-08 12:26:56 -080028 hasattr(obj, '__delete__')
29 )
Ethan Furman101e0742013-09-15 12:34:36 -070030
Ethan Furman6b3d64a2013-06-14 16:55:46 -070031def _is_dunder(name):
Ethan Furman6d3dfee2020-12-08 12:26:56 -080032 """
33 Returns True if a __dunder__ name, False otherwise.
34 """
35 return (
36 len(name) > 4 and
Brennan D Baraban8b914d22019-03-03 14:09:11 -080037 name[:2] == name[-2:] == '__' and
38 name[2] != '_' and
Ethan Furman6d3dfee2020-12-08 12:26:56 -080039 name[-3] != '_'
40 )
Ethan Furman6b3d64a2013-06-14 16:55:46 -070041
42def _is_sunder(name):
Ethan Furman6d3dfee2020-12-08 12:26:56 -080043 """
44 Returns True if a _sunder_ name, False otherwise.
45 """
46 return (
47 len(name) > 2 and
Brennan D Baraban8b914d22019-03-03 14:09:11 -080048 name[0] == name[-1] == '_' and
Ethan Furman6b3d64a2013-06-14 16:55:46 -070049 name[1:2] != '_' and
Ethan Furman6d3dfee2020-12-08 12:26:56 -080050 name[-2:-1] != '_'
51 )
Ethan Furman6b3d64a2013-06-14 16:55:46 -070052
Ethan Furman7cf0aad2020-12-09 17:12:11 -080053def _is_private(cls_name, name):
54 # do not use `re` as `re` imports `enum`
55 pattern = '_%s__' % (cls_name, )
56 if (
57 len(name) >= 5
58 and name.startswith(pattern)
59 and name[len(pattern)] != '_'
60 and (name[-1] != '_' or name[-2] != '_')
61 ):
62 return True
63 else:
64 return False
65
Ethan Furman7aaeb2a2021-01-25 14:26:19 -080066def _is_single_bit(num):
67 """
68 True if only one bit set in num (should be an int)
69 """
70 if num == 0:
71 return False
72 num &= num - 1
73 return num == 0
74
Ethan Furmanc314e602021-01-12 23:47:57 -080075def _make_class_unpicklable(obj):
Ethan Furman6d3dfee2020-12-08 12:26:56 -080076 """
Ethan Furmanc314e602021-01-12 23:47:57 -080077 Make the given obj un-picklable.
78
79 obj should be either a dictionary, on an Enum
Ethan Furman6d3dfee2020-12-08 12:26:56 -080080 """
Ethan Furmanca1b7942014-02-08 11:36:27 -080081 def _break_on_call_reduce(self, proto):
Ethan Furman6b3d64a2013-06-14 16:55:46 -070082 raise TypeError('%r cannot be pickled' % self)
Ethan Furmanc314e602021-01-12 23:47:57 -080083 if isinstance(obj, dict):
84 obj['__reduce_ex__'] = _break_on_call_reduce
85 obj['__module__'] = '<unknown>'
86 else:
87 setattr(obj, '__reduce_ex__', _break_on_call_reduce)
88 setattr(obj, '__module__', '<unknown>')
Ethan Furman6b3d64a2013-06-14 16:55:46 -070089
Ethan Furman7aaeb2a2021-01-25 14:26:19 -080090def _iter_bits_lsb(num):
91 while num:
92 b = num & (~num + 1)
93 yield b
94 num ^= b
95
96def bin(num, max_bits=None):
97 """
98 Like built-in bin(), except negative values are represented in
99 twos-compliment, and the leading bit always indicates sign
100 (0=positive, 1=negative).
101
102 >>> bin(10)
103 '0b0 1010'
104 >>> bin(~10) # ~10 is -11
105 '0b1 0101'
106 """
107
108 ceiling = 2 ** (num).bit_length()
109 if num >= 0:
110 s = _bltin_bin(num + ceiling).replace('1', '0', 1)
111 else:
112 s = _bltin_bin(~num ^ (ceiling - 1) + ceiling)
113 sign = s[:3]
114 digits = s[3:]
115 if max_bits is not None:
116 if len(digits) < max_bits:
117 digits = (sign[-1] * max_bits + digits)[-max_bits:]
118 return "%s %s" % (sign, digits)
119
120
Ethan Furman3515dcc2016-09-18 13:15:41 -0700121_auto_null = object()
Ethan Furmanc16595e2016-09-10 23:36:59 -0700122class auto:
123 """
124 Instances are replaced with an appropriate value in Enum class suites.
125 """
Ethan Furman3515dcc2016-09-18 13:15:41 -0700126 value = _auto_null
Ethan Furmanc16595e2016-09-10 23:36:59 -0700127
Ethan Furmanc314e602021-01-12 23:47:57 -0800128class property(DynamicClassAttribute):
129 """
130 This is a descriptor, used to define attributes that act differently
131 when accessed through an enum member and through an enum class.
132 Instance access is the same as property(), but access to an attribute
133 through the enum class will instead look in the class' _member_map_ for
134 a corresponding enum member.
135 """
136
137 def __get__(self, instance, ownerclass=None):
138 if instance is None:
139 try:
140 return ownerclass._member_map_[self.name]
141 except KeyError:
Ethan Furman7aaeb2a2021-01-25 14:26:19 -0800142 raise AttributeError(
Ethan Furmand65b9032021-02-08 17:32:38 -0800143 '%s: no class attribute %r' % (ownerclass.__name__, self.name)
Ethan Furman7aaeb2a2021-01-25 14:26:19 -0800144 )
Ethan Furmanc314e602021-01-12 23:47:57 -0800145 else:
146 if self.fget is None:
Ethan Furmand65b9032021-02-08 17:32:38 -0800147 # check for member
148 if self.name in ownerclass._member_map_:
149 import warnings
150 warnings.warn(
151 "accessing one member from another is not supported, "
Ethan Furman44e580f2021-03-03 09:54:30 -0800152 " and will be disabled in 3.12",
Ethan Furmand65b9032021-02-08 17:32:38 -0800153 DeprecationWarning,
154 stacklevel=2,
155 )
156 return ownerclass._member_map_[self.name]
Ethan Furman7aaeb2a2021-01-25 14:26:19 -0800157 raise AttributeError(
Ethan Furmand65b9032021-02-08 17:32:38 -0800158 '%s: no instance attribute %r' % (ownerclass.__name__, self.name)
Ethan Furman7aaeb2a2021-01-25 14:26:19 -0800159 )
Ethan Furmanc314e602021-01-12 23:47:57 -0800160 else:
161 return self.fget(instance)
162
163 def __set__(self, instance, value):
164 if self.fset is None:
Ethan Furman7aaeb2a2021-01-25 14:26:19 -0800165 raise AttributeError(
Ethan Furmand65b9032021-02-08 17:32:38 -0800166 "%s: cannot set instance attribute %r" % (self.clsname, self.name)
Ethan Furman7aaeb2a2021-01-25 14:26:19 -0800167 )
Ethan Furmanc314e602021-01-12 23:47:57 -0800168 else:
169 return self.fset(instance, value)
170
171 def __delete__(self, instance):
172 if self.fdel is None:
Ethan Furman7aaeb2a2021-01-25 14:26:19 -0800173 raise AttributeError(
Ethan Furmand65b9032021-02-08 17:32:38 -0800174 "%s: cannot delete instance attribute %r" % (self.clsname, self.name)
Ethan Furman7aaeb2a2021-01-25 14:26:19 -0800175 )
Ethan Furmanc314e602021-01-12 23:47:57 -0800176 else:
177 return self.fdel(instance)
178
179 def __set_name__(self, ownerclass, name):
180 self.name = name
181 self.clsname = ownerclass.__name__
182
183
184class _proto_member:
185 """
186 intermediate step for enum members between class execution and final creation
187 """
188
189 def __init__(self, value):
190 self.value = value
191
192 def __set_name__(self, enum_class, member_name):
193 """
194 convert each quasi-member into an instance of the new enum class
195 """
196 # first step: remove ourself from enum_class
197 delattr(enum_class, member_name)
198 # second step: create member based on enum_class
199 value = self.value
200 if not isinstance(value, tuple):
201 args = (value, )
202 else:
203 args = value
204 if enum_class._member_type_ is tuple: # special case for tuple enums
205 args = (args, ) # wrap it one more time
206 if not enum_class._use_args_:
207 enum_member = enum_class._new_member_(enum_class)
208 if not hasattr(enum_member, '_value_'):
209 enum_member._value_ = value
210 else:
211 enum_member = enum_class._new_member_(enum_class, *args)
212 if not hasattr(enum_member, '_value_'):
213 if enum_class._member_type_ is object:
214 enum_member._value_ = value
215 else:
Ethan Furman7aaeb2a2021-01-25 14:26:19 -0800216 try:
217 enum_member._value_ = enum_class._member_type_(*args)
218 except Exception as exc:
219 raise TypeError(
220 '_value_ not set in __new__, unable to create it'
221 ) from None
Ethan Furmanc314e602021-01-12 23:47:57 -0800222 value = enum_member._value_
223 enum_member._name_ = member_name
224 enum_member.__objclass__ = enum_class
225 enum_member.__init__(*args)
Ethan Furman7aaeb2a2021-01-25 14:26:19 -0800226 enum_member._sort_order_ = len(enum_class._member_names_)
Ethan Furmanc314e602021-01-12 23:47:57 -0800227 # If another member with the same value was already defined, the
228 # new member becomes an alias to the existing one.
229 for name, canonical_member in enum_class._member_map_.items():
230 if canonical_member._value_ == enum_member._value_:
231 enum_member = canonical_member
232 break
233 else:
Ethan Furman7aaeb2a2021-01-25 14:26:19 -0800234 # this could still be an alias if the value is multi-bit and the
235 # class is a flag class
236 if (
237 Flag is None
238 or not issubclass(enum_class, Flag)
239 ):
240 # no other instances found, record this member in _member_names_
241 enum_class._member_names_.append(member_name)
242 elif (
243 Flag is not None
244 and issubclass(enum_class, Flag)
245 and _is_single_bit(value)
246 ):
247 # no other instances found, record this member in _member_names_
248 enum_class._member_names_.append(member_name)
Ethan Furmanc314e602021-01-12 23:47:57 -0800249 # get redirect in place before adding to _member_map_
250 # but check for other instances in parent classes first
251 need_override = False
252 descriptor = None
253 for base in enum_class.__mro__[1:]:
254 descriptor = base.__dict__.get(member_name)
255 if descriptor is not None:
256 if isinstance(descriptor, (property, DynamicClassAttribute)):
257 break
258 else:
259 need_override = True
260 # keep looking for an enum.property
261 if descriptor and not need_override:
262 # previous enum.property found, no further action needed
263 pass
264 else:
265 redirect = property()
266 redirect.__set_name__(enum_class, member_name)
267 if descriptor and need_override:
268 # previous enum.property found, but some other inherited attribute
269 # is in the way; copy fget, fset, fdel to this one
270 redirect.fget = descriptor.fget
271 redirect.fset = descriptor.fset
272 redirect.fdel = descriptor.fdel
273 setattr(enum_class, member_name, redirect)
274 # now add to _member_map_ (even aliases)
275 enum_class._member_map_[member_name] = enum_member
276 try:
277 # This may fail if value is not hashable. We can't add the value
278 # to the map, and by-value lookups for this value will be
279 # linear.
Ethan Furman7aaeb2a2021-01-25 14:26:19 -0800280 enum_class._value2member_map_.setdefault(value, enum_member)
Ethan Furmanc314e602021-01-12 23:47:57 -0800281 except TypeError:
282 pass
283
Ethan Furman101e0742013-09-15 12:34:36 -0700284
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700285class _EnumDict(dict):
Ethan Furman6d3dfee2020-12-08 12:26:56 -0800286 """
287 Track enum member order and ensure member names are not reused.
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700288
Ethan Furmanb7751062021-03-30 21:17:26 -0700289 EnumType will use the names found in self._member_names as the
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700290 enumeration member names.
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700291 """
292 def __init__(self):
293 super().__init__()
294 self._member_names = []
Ethan Furmanc16595e2016-09-10 23:36:59 -0700295 self._last_values = []
Ethan Furmana4b1bb42018-01-22 07:56:37 -0800296 self._ignore = []
Ethan Onstottd9a43e22020-04-28 13:20:55 -0400297 self._auto_called = False
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700298
299 def __setitem__(self, key, value):
Ethan Furman6d3dfee2020-12-08 12:26:56 -0800300 """
301 Changes anything not dundered or not a descriptor.
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700302
303 If an enum member name is used twice, an error is raised; duplicate
304 values are not checked for.
305
306 Single underscore (sunder) names are reserved.
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700307 """
Ethan Furman7cf0aad2020-12-09 17:12:11 -0800308 if _is_private(self._cls_name, key):
309 # do nothing, name will be a normal attribute
310 pass
311 elif _is_sunder(key):
Ethan Furmanee47e5c2016-08-31 00:12:15 -0700312 if key not in (
Ethan Furman3515dcc2016-09-18 13:15:41 -0700313 '_order_', '_create_pseudo_member_',
Ethan Furmana4b1bb42018-01-22 07:56:37 -0800314 '_generate_next_value_', '_missing_', '_ignore_',
Ethan Furman7aaeb2a2021-01-25 14:26:19 -0800315 '_iter_member_', '_iter_member_by_value_', '_iter_member_by_def_',
Ethan Furmanee47e5c2016-08-31 00:12:15 -0700316 ):
Ethan Furman6d3dfee2020-12-08 12:26:56 -0800317 raise ValueError(
318 '_sunder_ names, such as %r, are reserved for future Enum use'
319 % (key, )
320 )
Ethan Furmanc16595e2016-09-10 23:36:59 -0700321 if key == '_generate_next_value_':
Ethan Onstottd9a43e22020-04-28 13:20:55 -0400322 # check if members already defined as auto()
323 if self._auto_called:
324 raise TypeError("_generate_next_value_ must be defined before members")
Ethan Furmanb7751062021-03-30 21:17:26 -0700325 _gnv = value.__func__ if isinstance(value, staticmethod) else value
326 setattr(self, '_generate_next_value', _gnv)
Ethan Furmana4b1bb42018-01-22 07:56:37 -0800327 elif key == '_ignore_':
328 if isinstance(value, str):
329 value = value.replace(',',' ').split()
330 else:
331 value = list(value)
332 self._ignore = value
333 already = set(value) & set(self._member_names)
334 if already:
Ethan Furman6d3dfee2020-12-08 12:26:56 -0800335 raise ValueError(
336 '_ignore_ cannot specify already set names: %r'
337 % (already, )
338 )
Ethan Furman101e0742013-09-15 12:34:36 -0700339 elif _is_dunder(key):
Ethan Furmane8e61272016-08-20 07:19:31 -0700340 if key == '__order__':
341 key = '_order_'
Ethan Furman101e0742013-09-15 12:34:36 -0700342 elif key in self._member_names:
343 # descriptor overwriting an enum?
Ethan Furmana6582872020-12-10 13:07:00 -0800344 raise TypeError('%r already defined as: %r' % (key, self[key]))
Ethan Furmana4b1bb42018-01-22 07:56:37 -0800345 elif key in self._ignore:
346 pass
Ethan Furman101e0742013-09-15 12:34:36 -0700347 elif not _is_descriptor(value):
348 if key in self:
349 # enum overwriting a descriptor?
Ethan Furmanee47e5c2016-08-31 00:12:15 -0700350 raise TypeError('%r already defined as: %r' % (key, self[key]))
Ethan Furmanc16595e2016-09-10 23:36:59 -0700351 if isinstance(value, auto):
Ethan Furman3515dcc2016-09-18 13:15:41 -0700352 if value.value == _auto_null:
Ethan Furman6d3dfee2020-12-08 12:26:56 -0800353 value.value = self._generate_next_value(
Ethan Furman7aaeb2a2021-01-25 14:26:19 -0800354 key, 1, len(self._member_names), self._last_values[:],
Ethan Furman6d3dfee2020-12-08 12:26:56 -0800355 )
Ethan Furmanfc23a942020-09-16 12:37:54 -0700356 self._auto_called = True
Ethan Furman3515dcc2016-09-18 13:15:41 -0700357 value = value.value
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700358 self._member_names.append(key)
Ethan Furmanc16595e2016-09-10 23:36:59 -0700359 self._last_values.append(value)
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700360 super().__setitem__(key, value)
361
Ethan Furmana6582872020-12-10 13:07:00 -0800362 def update(self, members, **more_members):
363 try:
364 for name in members.keys():
365 self[name] = members[name]
366 except AttributeError:
367 for name, value in members:
368 self[name] = value
369 for name, value in more_members.items():
370 self[name] = value
371
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700372
Ethan Furmanb7751062021-03-30 21:17:26 -0700373class EnumType(type):
Ethan Furman6d3dfee2020-12-08 12:26:56 -0800374 """
375 Metaclass for Enum
376 """
Ethan Furman7aaeb2a2021-01-25 14:26:19 -0800377
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700378 @classmethod
Ethan Furman6ec0ade2020-12-24 10:05:02 -0800379 def __prepare__(metacls, cls, bases, **kwds):
Ethan Furman3064dbf2020-09-16 07:11:57 -0700380 # check that previous enum members do not exist
381 metacls._check_for_existing_members(cls, bases)
Ethan Furmanee47e5c2016-08-31 00:12:15 -0700382 # create the namespace dict
383 enum_dict = _EnumDict()
Ethan Furman7cf0aad2020-12-09 17:12:11 -0800384 enum_dict._cls_name = cls
Ethan Furmanee47e5c2016-08-31 00:12:15 -0700385 # inherit previous flags and _generate_next_value_ function
Ethan Furman3064dbf2020-09-16 07:11:57 -0700386 member_type, first_enum = metacls._get_mixins_(cls, bases)
Ethan Furmanee47e5c2016-08-31 00:12:15 -0700387 if first_enum is not None:
Ethan Furman6d3dfee2020-12-08 12:26:56 -0800388 enum_dict['_generate_next_value_'] = getattr(
389 first_enum, '_generate_next_value_', None,
390 )
Ethan Furmanee47e5c2016-08-31 00:12:15 -0700391 return enum_dict
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700392
Ethan Furman7aaeb2a2021-01-25 14:26:19 -0800393 def __new__(metacls, cls, bases, classdict, boundary=None, **kwds):
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700394 # an Enum class is final once enumeration items have been defined; it
395 # cannot be mixed with other types (int, float, etc.) if it has an
396 # inherited __new__ unless a new __new__ is defined (or the resulting
397 # class will fail).
Ethan Furmana4b1bb42018-01-22 07:56:37 -0800398 #
399 # remove any keys listed in _ignore_
400 classdict.setdefault('_ignore_', []).append('_ignore_')
401 ignore = classdict['_ignore_']
402 for key in ignore:
403 classdict.pop(key, None)
Ethan Furmanc314e602021-01-12 23:47:57 -0800404 #
405 # grab member names
406 member_names = classdict._member_names
407 #
408 # check for illegal enum names (any others?)
409 invalid_names = set(member_names) & {'mro', ''}
410 if invalid_names:
411 raise ValueError('Invalid enum member name: {0}'.format(
412 ','.join(invalid_names)))
413 #
414 # adjust the sunders
415 _order_ = classdict.pop('_order_', None)
416 # convert to normal dict
417 classdict = dict(classdict.items())
418 #
419 # data type of member and the controlling Enum class
Ethan Furman3064dbf2020-09-16 07:11:57 -0700420 member_type, first_enum = metacls._get_mixins_(cls, bases)
Ethan Furmanc2667362020-12-07 00:17:31 -0800421 __new__, save_new, use_args = metacls._find_new_(
422 classdict, member_type, first_enum,
423 )
Ethan Furmanc314e602021-01-12 23:47:57 -0800424 classdict['_new_member_'] = __new__
425 classdict['_use_args_'] = use_args
426 #
427 # convert future enum members into temporary _proto_members
Ethan Furman7aaeb2a2021-01-25 14:26:19 -0800428 # and record integer values in case this will be a Flag
429 flag_mask = 0
Ethan Furmanc314e602021-01-12 23:47:57 -0800430 for name in member_names:
Ethan Furman7aaeb2a2021-01-25 14:26:19 -0800431 value = classdict[name]
432 if isinstance(value, int):
433 flag_mask |= value
434 classdict[name] = _proto_member(value)
Ethan Furmanc314e602021-01-12 23:47:57 -0800435 #
Ethan Furman7aaeb2a2021-01-25 14:26:19 -0800436 # house-keeping structures
Ethan Furmanc314e602021-01-12 23:47:57 -0800437 classdict['_member_names_'] = []
438 classdict['_member_map_'] = {}
439 classdict['_value2member_map_'] = {}
440 classdict['_member_type_'] = member_type
441 #
Ethan Furman7aaeb2a2021-01-25 14:26:19 -0800442 # Flag structures (will be removed if final class is not a Flag
443 classdict['_boundary_'] = (
444 boundary
445 or getattr(first_enum, '_boundary_', None)
446 )
447 classdict['_flag_mask_'] = flag_mask
448 classdict['_all_bits_'] = 2 ** ((flag_mask).bit_length()) - 1
449 classdict['_inverted_'] = None
450 #
Ethan Furman2da95042014-03-03 12:42:52 -0800451 # If a custom type is mixed into the Enum, and it does not know how
452 # to pickle itself, pickle.dumps will succeed but pickle.loads will
453 # fail. Rather than have the error show up later and possibly far
454 # from the source, sabotage the pickle protocol for this class so
455 # that pickle.dumps also fails.
456 #
457 # However, if the new class implements its own __reduce_ex__, do not
458 # sabotage -- it's on them to make sure it works correctly. We use
459 # __reduce_ex__ instead of any of the others as it is preferred by
460 # pickle over __reduce__, and it handles all pickle protocols.
461 if '__reduce_ex__' not in classdict:
Ethan Furmandc870522014-02-18 12:37:12 -0800462 if member_type is not object:
463 methods = ('__getnewargs_ex__', '__getnewargs__',
464 '__reduce_ex__', '__reduce__')
Ethan Furman2da95042014-03-03 12:42:52 -0800465 if not any(m in member_type.__dict__ for m in methods):
Ethan Furmanc314e602021-01-12 23:47:57 -0800466 _make_class_unpicklable(classdict)
467 #
468 # create a default docstring if one has not been provided
469 if '__doc__' not in classdict:
470 classdict['__doc__'] = 'An enumeration.'
471 try:
472 exc = None
473 enum_class = super().__new__(metacls, cls, bases, classdict, **kwds)
474 except RuntimeError as e:
475 # any exceptions raised by member.__new__ will get converted to a
476 # RuntimeError, so get that original exception back and raise it instead
477 exc = e.__cause__ or e
478 if exc is not None:
479 raise exc
480 #
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700481 # double check that repr and friends are not the mixin's or various
482 # things break (such as pickle)
Ethan Furman22415ad2020-09-15 16:28:25 -0700483 # however, if the method is defined in the Enum itself, don't replace
484 # it
Ethan Furmandc870522014-02-18 12:37:12 -0800485 for name in ('__repr__', '__str__', '__format__', '__reduce_ex__'):
Ethan Furman22415ad2020-09-15 16:28:25 -0700486 if name in classdict:
487 continue
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700488 class_method = getattr(enum_class, name)
489 obj_method = getattr(member_type, name, None)
490 enum_method = getattr(first_enum, name, None)
491 if obj_method is not None and obj_method is class_method:
492 setattr(enum_class, name, enum_method)
Ethan Furmanc314e602021-01-12 23:47:57 -0800493 #
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700494 # replace any other __new__ with our own (as long as Enum is not None,
495 # anyway) -- again, this is to support pickle
496 if Enum is not None:
497 # if the user defined their own __new__, save it before it gets
498 # clobbered in case they subclass later
499 if save_new:
500 enum_class.__new_member__ = __new__
501 enum_class.__new__ = Enum.__new__
Ethan Furmanc314e602021-01-12 23:47:57 -0800502 #
Ethan Furmane8e61272016-08-20 07:19:31 -0700503 # py3 support for definition order (helps keep py2/py3 code in sync)
Ethan Furman7aaeb2a2021-01-25 14:26:19 -0800504 #
505 # _order_ checking is spread out into three/four steps
506 # - if enum_class is a Flag:
507 # - remove any non-single-bit flags from _order_
508 # - remove any aliases from _order_
509 # - check that _order_ and _member_names_ match
510 #
511 # step 1: ensure we have a list
Ethan Furmane8e61272016-08-20 07:19:31 -0700512 if _order_ is not None:
513 if isinstance(_order_, str):
514 _order_ = _order_.replace(',', ' ').split()
Ethan Furman7aaeb2a2021-01-25 14:26:19 -0800515 #
516 # remove Flag structures if final class is not a Flag
517 if (
518 Flag is None and cls != 'Flag'
519 or Flag is not None and not issubclass(enum_class, Flag)
520 ):
521 delattr(enum_class, '_boundary_')
522 delattr(enum_class, '_flag_mask_')
523 delattr(enum_class, '_all_bits_')
524 delattr(enum_class, '_inverted_')
525 elif Flag is not None and issubclass(enum_class, Flag):
526 # ensure _all_bits_ is correct and there are no missing flags
527 single_bit_total = 0
528 multi_bit_total = 0
529 for flag in enum_class._member_map_.values():
530 flag_value = flag._value_
531 if _is_single_bit(flag_value):
532 single_bit_total |= flag_value
533 else:
534 # multi-bit flags are considered aliases
535 multi_bit_total |= flag_value
536 if enum_class._boundary_ is not KEEP:
537 missed = list(_iter_bits_lsb(multi_bit_total & ~single_bit_total))
538 if missed:
539 raise TypeError(
540 'invalid Flag %r -- missing values: %s'
541 % (cls, ', '.join((str(i) for i in missed)))
542 )
543 enum_class._flag_mask_ = single_bit_total
544 #
545 # set correct __iter__
546 member_list = [m._value_ for m in enum_class]
547 if member_list != sorted(member_list):
548 enum_class._iter_member_ = enum_class._iter_member_by_def_
549 if _order_:
550 # _order_ step 2: remove any items from _order_ that are not single-bit
551 _order_ = [
552 o
553 for o in _order_
554 if o not in enum_class._member_map_ or _is_single_bit(enum_class[o]._value_)
555 ]
556 #
557 if _order_:
558 # _order_ step 3: remove aliases from _order_
559 _order_ = [
560 o
561 for o in _order_
562 if (
563 o not in enum_class._member_map_
564 or
565 (o in enum_class._member_map_ and o in enum_class._member_names_)
566 )]
567 # _order_ step 4: verify that _order_ and _member_names_ match
Ethan Furmane8e61272016-08-20 07:19:31 -0700568 if _order_ != enum_class._member_names_:
Ethan Furman7aaeb2a2021-01-25 14:26:19 -0800569 raise TypeError(
570 'member order does not match _order_:\n%r\n%r'
571 % (enum_class._member_names_, _order_)
572 )
Ethan Furmanc314e602021-01-12 23:47:57 -0800573 #
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700574 return enum_class
575
Ethan Furman5de67b12016-04-13 23:52:09 -0700576 def __bool__(self):
577 """
578 classes/types should always be True.
579 """
580 return True
581
Ethan Furman7aaeb2a2021-01-25 14:26:19 -0800582 def __call__(cls, value, names=None, *, module=None, qualname=None, type=None, start=1, boundary=None):
Ethan Furman6d3dfee2020-12-08 12:26:56 -0800583 """
584 Either returns an existing member, or creates a new enum class.
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700585
586 This method is used both when an enum class is given a value to match
587 to an enumeration member (i.e. Color(3)) and for the functional API
Ethan Furman23bb6f42016-11-21 09:22:05 -0800588 (i.e. Color = Enum('Color', names='RED GREEN BLUE')).
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700589
Ethan Furman2da95042014-03-03 12:42:52 -0800590 When used for the functional API:
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700591
Ethan Furman2da95042014-03-03 12:42:52 -0800592 `value` will be the name of the new class.
593
594 `names` should be either a string of white-space/comma delimited names
Ethan Furmand9925a12014-09-16 20:35:55 -0700595 (values will start at `start`), or an iterator/mapping of name, value pairs.
Ethan Furman2da95042014-03-03 12:42:52 -0800596
597 `module` should be set to the module this class is being created in;
598 if it is not set, an attempt to find that module will be made, but if
599 it fails the class will not be picklable.
600
601 `qualname` should be set to the actual location this class can be found
602 at in its module; by default it is set to the global scope. If this is
603 not correct, unpickling will fail in some circumstances.
604
605 `type`, if set, will be mixed in as the first base class.
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700606 """
607 if names is None: # simple value lookup
608 return cls.__new__(cls, value)
609 # otherwise, functional API: we're creating a new Enum type
Ethan Furman6d3dfee2020-12-08 12:26:56 -0800610 return cls._create_(
611 value,
612 names,
613 module=module,
614 qualname=qualname,
615 type=type,
616 start=start,
Ethan Furman7aaeb2a2021-01-25 14:26:19 -0800617 boundary=boundary,
Ethan Furman6d3dfee2020-12-08 12:26:56 -0800618 )
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700619
620 def __contains__(cls, member):
Rahul Jha94306522018-09-10 23:51:04 +0530621 if not isinstance(member, Enum):
622 raise TypeError(
623 "unsupported operand type(s) for 'in': '%s' and '%s'" % (
624 type(member).__qualname__, cls.__class__.__qualname__))
Ethan Furman0081f232014-09-16 17:31:23 -0700625 return isinstance(member, cls) and member._name_ in cls._member_map_
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700626
Ethan Furman64a99722013-09-22 16:18:19 -0700627 def __delattr__(cls, attr):
628 # nicer error message when someone tries to delete an attribute
629 # (see issue19025).
630 if attr in cls._member_map_:
Ethan Furman6d3dfee2020-12-08 12:26:56 -0800631 raise AttributeError("%s: cannot delete Enum member %r." % (cls.__name__, attr))
Ethan Furman64a99722013-09-22 16:18:19 -0700632 super().__delattr__(attr)
633
Ethan Furman388a3922013-08-12 06:51:41 -0700634 def __dir__(self):
Ethan Furman6d3dfee2020-12-08 12:26:56 -0800635 return (
636 ['__class__', '__doc__', '__members__', '__module__']
637 + self._member_names_
638 )
Ethan Furman388a3922013-08-12 06:51:41 -0700639
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700640 def __getattr__(cls, name):
Ethan Furman6d3dfee2020-12-08 12:26:56 -0800641 """
642 Return the enum member matching `name`
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700643
644 We use __getattr__ instead of descriptors or inserting into the enum
645 class' __dict__ in order to support `name` and `value` being both
646 properties for enum members (which live in the class' __dict__) and
647 enum members themselves.
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700648 """
649 if _is_dunder(name):
650 raise AttributeError(name)
651 try:
Ethan Furman520ad572013-07-19 19:47:21 -0700652 return cls._member_map_[name]
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700653 except KeyError:
654 raise AttributeError(name) from None
655
656 def __getitem__(cls, name):
Ethan Furman520ad572013-07-19 19:47:21 -0700657 return cls._member_map_[name]
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700658
659 def __iter__(cls):
Ethan Furman6d3dfee2020-12-08 12:26:56 -0800660 """
661 Returns members in definition order.
662 """
Ethan Furman520ad572013-07-19 19:47:21 -0700663 return (cls._member_map_[name] for name in cls._member_names_)
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700664
665 def __len__(cls):
Ethan Furman520ad572013-07-19 19:47:21 -0700666 return len(cls._member_names_)
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700667
Ethan Furmanc314e602021-01-12 23:47:57 -0800668 @_bltin_property
Ethan Furman2131a4a2013-09-14 18:11:24 -0700669 def __members__(cls):
Ethan Furman6d3dfee2020-12-08 12:26:56 -0800670 """
671 Returns a mapping of member name->value.
Ethan Furman2131a4a2013-09-14 18:11:24 -0700672
673 This mapping lists all enum members, including aliases. Note that this
674 is a read-only view of the internal mapping.
Ethan Furman2131a4a2013-09-14 18:11:24 -0700675 """
676 return MappingProxyType(cls._member_map_)
677
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700678 def __repr__(cls):
679 return "<enum %r>" % cls.__name__
680
Ethan Furman2131a4a2013-09-14 18:11:24 -0700681 def __reversed__(cls):
Ethan Furman6d3dfee2020-12-08 12:26:56 -0800682 """
683 Returns members in reverse definition order.
684 """
Ethan Furman2131a4a2013-09-14 18:11:24 -0700685 return (cls._member_map_[name] for name in reversed(cls._member_names_))
686
Ethan Furmanf203f2d2013-09-06 07:16:48 -0700687 def __setattr__(cls, name, value):
Ethan Furman6d3dfee2020-12-08 12:26:56 -0800688 """
689 Block attempts to reassign Enum members.
Ethan Furmanf203f2d2013-09-06 07:16:48 -0700690
691 A simple assignment to the class namespace only changes one of the
692 several possible ways to get an Enum member from the Enum class,
693 resulting in an inconsistent Enumeration.
Ethan Furmanf203f2d2013-09-06 07:16:48 -0700694 """
695 member_map = cls.__dict__.get('_member_map_', {})
696 if name in member_map:
697 raise AttributeError('Cannot reassign members.')
698 super().__setattr__(name, value)
699
Ethan Furman7aaeb2a2021-01-25 14:26:19 -0800700 def _create_(cls, class_name, names, *, module=None, qualname=None, type=None, start=1, boundary=None):
Ethan Furman6d3dfee2020-12-08 12:26:56 -0800701 """
702 Convenience method to create a new Enum class.
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700703
704 `names` can be:
705
706 * A string containing member names, separated either with spaces or
Ethan Furmand9925a12014-09-16 20:35:55 -0700707 commas. Values are incremented by 1 from `start`.
708 * An iterable of member names. Values are incremented by 1 from `start`.
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700709 * An iterable of (member name, value) pairs.
Ethan Furmand9925a12014-09-16 20:35:55 -0700710 * A mapping of member name -> value pairs.
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700711 """
712 metacls = cls.__class__
713 bases = (cls, ) if type is None else (type, cls)
Ethan Furman3064dbf2020-09-16 07:11:57 -0700714 _, first_enum = cls._get_mixins_(cls, bases)
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700715 classdict = metacls.__prepare__(class_name, bases)
716
717 # special processing needed for names?
718 if isinstance(names, str):
719 names = names.replace(',', ' ').split()
Dong-hee Nadcc8ce42017-06-22 01:52:32 +0900720 if isinstance(names, (tuple, list)) and names and isinstance(names[0], str):
Ethan Furmanee47e5c2016-08-31 00:12:15 -0700721 original_names, names = names, []
Ethan Furmanc16595e2016-09-10 23:36:59 -0700722 last_values = []
Ethan Furmanee47e5c2016-08-31 00:12:15 -0700723 for count, name in enumerate(original_names):
Ethan Furmanc16595e2016-09-10 23:36:59 -0700724 value = first_enum._generate_next_value_(name, start, count, last_values[:])
725 last_values.append(value)
726 names.append((name, value))
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700727
728 # Here, names is either an iterable of (name, value) or a mapping.
729 for item in names:
730 if isinstance(item, str):
731 member_name, member_value = item, names[item]
732 else:
733 member_name, member_value = item
734 classdict[member_name] = member_value
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700735
736 # TODO: replace the frame hack if a blessed way to know the calling
737 # module is ever developed
738 if module is None:
739 try:
740 module = sys._getframe(2).f_globals['__name__']
Pablo Galindo293dd232019-11-19 21:34:03 +0000741 except (AttributeError, ValueError, KeyError):
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700742 pass
743 if module is None:
Ethan Furmanc314e602021-01-12 23:47:57 -0800744 _make_class_unpicklable(classdict)
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700745 else:
Ethan Furmanc314e602021-01-12 23:47:57 -0800746 classdict['__module__'] = module
Ethan Furmanca1b7942014-02-08 11:36:27 -0800747 if qualname is not None:
Ethan Furmanc314e602021-01-12 23:47:57 -0800748 classdict['__qualname__'] = qualname
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700749
Ethan Furman7aaeb2a2021-01-25 14:26:19 -0800750 return metacls.__new__(metacls, class_name, bases, classdict, boundary=boundary)
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700751
Ethan Furman7aaeb2a2021-01-25 14:26:19 -0800752 def _convert_(cls, name, module, filter, source=None, boundary=None):
orlnub1230fb9fad2018-09-12 20:28:53 +0300753 """
754 Create a new Enum subclass that replaces a collection of global constants
755 """
756 # convert all constants from source (or module) that pass filter() to
757 # a new Enum called name, and export the enum and its members back to
758 # module;
759 # also, replace the __reduce_ex__ method so unpickling works in
760 # previous Python versions
Ethan Furmanb7751062021-03-30 21:17:26 -0700761 module_globals = sys.modules[module].__dict__
orlnub1230fb9fad2018-09-12 20:28:53 +0300762 if source:
Ethan Furmanb7751062021-03-30 21:17:26 -0700763 source = source.__dict__
orlnub1230fb9fad2018-09-12 20:28:53 +0300764 else:
765 source = module_globals
766 # _value2member_map_ is populated in the same order every time
767 # for a consistent reverse mapping of number to name when there
768 # are multiple names for the same number.
769 members = [
770 (name, value)
771 for name, value in source.items()
772 if filter(name)]
773 try:
774 # sort by value
775 members.sort(key=lambda t: (t[1], t[0]))
776 except TypeError:
777 # unless some values aren't comparable, in which case sort by name
778 members.sort(key=lambda t: t[0])
Ethan Furman7aaeb2a2021-01-25 14:26:19 -0800779 cls = cls(name, members, module=module, boundary=boundary or KEEP)
orlnub1230fb9fad2018-09-12 20:28:53 +0300780 cls.__reduce_ex__ = _reduce_ex_by_name
Ethan Furmanb7751062021-03-30 21:17:26 -0700781 global_enum(cls)
orlnub1230fb9fad2018-09-12 20:28:53 +0300782 module_globals[name] = cls
783 return cls
784
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700785 @staticmethod
Ethan Furman3064dbf2020-09-16 07:11:57 -0700786 def _check_for_existing_members(class_name, bases):
787 for chain in bases:
788 for base in chain.__mro__:
789 if issubclass(base, Enum) and base._member_names_:
Ethan Furman6d3dfee2020-12-08 12:26:56 -0800790 raise TypeError(
791 "%s: cannot extend enumeration %r"
792 % (class_name, base.__name__)
793 )
Ethan Furman3064dbf2020-09-16 07:11:57 -0700794
795 @staticmethod
796 def _get_mixins_(class_name, bases):
Ethan Furman6d3dfee2020-12-08 12:26:56 -0800797 """
798 Returns the type for creating enum members, and the first inherited
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700799 enum class.
800
801 bases: the tuple of bases that was given to __new__
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700802 """
803 if not bases:
804 return object, Enum
805
Ethan Furman5bdab642018-09-21 19:03:09 -0700806 def _find_data_type(bases):
Ethan Furmanbff01f32020-09-15 15:56:26 -0700807 data_types = []
Ethan Furman5bdab642018-09-21 19:03:09 -0700808 for chain in bases:
Ethan Furmanbff01f32020-09-15 15:56:26 -0700809 candidate = None
Ethan Furman5bdab642018-09-21 19:03:09 -0700810 for base in chain.__mro__:
811 if base is object:
812 continue
Ethan Furmanc2667362020-12-07 00:17:31 -0800813 elif issubclass(base, Enum):
814 if base._member_type_ is not object:
815 data_types.append(base._member_type_)
816 break
Ethan Furman5bdab642018-09-21 19:03:09 -0700817 elif '__new__' in base.__dict__:
Ethan Furmancd453852018-10-05 23:29:36 -0700818 if issubclass(base, Enum):
Ethan Furman5bdab642018-09-21 19:03:09 -0700819 continue
Ethan Furmanbff01f32020-09-15 15:56:26 -0700820 data_types.append(candidate or base)
821 break
Ethan Furmanc2667362020-12-07 00:17:31 -0800822 else:
Ethan Furmanbff01f32020-09-15 15:56:26 -0700823 candidate = base
824 if len(data_types) > 1:
Ethan Furman3064dbf2020-09-16 07:11:57 -0700825 raise TypeError('%r: too many data types: %r' % (class_name, data_types))
Ethan Furmanbff01f32020-09-15 15:56:26 -0700826 elif data_types:
827 return data_types[0]
828 else:
829 return None
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700830
Ethan Furman5bdab642018-09-21 19:03:09 -0700831 # ensure final parent class is an Enum derivative, find any concrete
832 # data type, and check that Enum has no members
833 first_enum = bases[-1]
834 if not issubclass(first_enum, Enum):
835 raise TypeError("new enumerations should be created as "
836 "`EnumName([mixin_type, ...] [data_type,] enum_type)`")
837 member_type = _find_data_type(bases) or object
838 if first_enum._member_names_:
839 raise TypeError("Cannot extend enumerations")
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700840 return member_type, first_enum
841
842 @staticmethod
843 def _find_new_(classdict, member_type, first_enum):
Ethan Furman6d3dfee2020-12-08 12:26:56 -0800844 """
845 Returns the __new__ to be used for creating the enum members.
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700846
847 classdict: the class dictionary given to __new__
848 member_type: the data type whose __new__ will be used by default
849 first_enum: enumeration to check for an overriding __new__
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700850 """
851 # now find the correct __new__, checking to see of one was defined
852 # by the user; also check earlier enum classes in case a __new__ was
853 # saved as __new_member__
854 __new__ = classdict.get('__new__', None)
855
856 # should __new__ be saved as __new_member__ later?
857 save_new = __new__ is not None
858
859 if __new__ is None:
860 # check all possibles for __new_member__ before falling back to
861 # __new__
862 for method in ('__new_member__', '__new__'):
863 for possible in (member_type, first_enum):
864 target = getattr(possible, method, None)
865 if target not in {
866 None,
867 None.__new__,
868 object.__new__,
869 Enum.__new__,
870 }:
871 __new__ = target
872 break
873 if __new__ is not None:
874 break
875 else:
876 __new__ = object.__new__
877
878 # if a non-object.__new__ is used then whatever value/tuple was
879 # assigned to the enum member name will be passed to __new__ and to the
880 # new enum member's __init__
881 if __new__ is object.__new__:
882 use_args = False
883 else:
884 use_args = True
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700885 return __new__, save_new, use_args
Ethan Furmanb7751062021-03-30 21:17:26 -0700886EnumMeta = EnumType
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700887
888
Ethan Furmanb7751062021-03-30 21:17:26 -0700889class Enum(metaclass=EnumType):
Ethan Furman6d3dfee2020-12-08 12:26:56 -0800890 """
891 Generic enumeration.
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700892
893 Derive from this class to define new enumerations.
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700894 """
Ethan Furman7aaeb2a2021-01-25 14:26:19 -0800895
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700896 def __new__(cls, value):
897 # all enum instances are actually created during class construction
898 # without calling this method; this method is called by the metaclass'
899 # __call__ (i.e. Color(3) ), and by pickle
900 if type(value) is cls:
Ethan Furman23bb6f42016-11-21 09:22:05 -0800901 # For lookups like Color(Color.RED)
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700902 return value
903 # by-value search for a matching enum member
904 # see if it's in the reverse mapping (for hashable values)
Ethan Furman2aa27322013-07-19 19:35:56 -0700905 try:
Andrew Svetlov34ae04f2018-12-26 20:45:33 +0200906 return cls._value2member_map_[value]
907 except KeyError:
908 # Not found, no need to do long O(n) search
909 pass
Ethan Furman2aa27322013-07-19 19:35:56 -0700910 except TypeError:
911 # not there, now do long search -- O(n) behavior
Ethan Furman520ad572013-07-19 19:47:21 -0700912 for member in cls._member_map_.values():
Ethan Furman0081f232014-09-16 17:31:23 -0700913 if member._value_ == value:
Ethan Furman2aa27322013-07-19 19:35:56 -0700914 return member
Ethan Furmanee47e5c2016-08-31 00:12:15 -0700915 # still not found -- try _missing_ hook
Ethan Furman019f0a02018-09-12 11:43:34 -0700916 try:
917 exc = None
918 result = cls._missing_(value)
919 except Exception as e:
920 exc = e
921 result = None
922 if isinstance(result, cls):
923 return result
Ethan Furman7aaeb2a2021-01-25 14:26:19 -0800924 elif (
925 Flag is not None and issubclass(cls, Flag)
926 and cls._boundary_ is EJECT and isinstance(result, int)
927 ):
928 return result
Ethan Furman019f0a02018-09-12 11:43:34 -0700929 else:
Walter Dörwald323842c2019-07-18 20:37:13 +0200930 ve_exc = ValueError("%r is not a valid %s" % (value, cls.__qualname__))
Ethan Furman019f0a02018-09-12 11:43:34 -0700931 if result is None and exc is None:
932 raise ve_exc
933 elif exc is None:
934 exc = TypeError(
935 'error in %s._missing_: returned %r instead of None or a valid member'
936 % (cls.__name__, result)
937 )
Ethan Furman7aaeb2a2021-01-25 14:26:19 -0800938 if not isinstance(exc, ValueError):
939 exc.__context__ = ve_exc
Ethan Furman019f0a02018-09-12 11:43:34 -0700940 raise exc
Ethan Furmanee47e5c2016-08-31 00:12:15 -0700941
Ethan Furmanc16595e2016-09-10 23:36:59 -0700942 def _generate_next_value_(name, start, count, last_values):
Ethan Furman6d3dfee2020-12-08 12:26:56 -0800943 """
944 Generate the next value when not given.
945
946 name: the name of the member
947 start: the initial start value or None
948 count: the number of existing members
949 last_value: the last value assigned or None
950 """
Ethan Furmanc16595e2016-09-10 23:36:59 -0700951 for last_value in reversed(last_values):
952 try:
953 return last_value + 1
954 except TypeError:
955 pass
956 else:
Ethan Furmanee47e5c2016-08-31 00:12:15 -0700957 return start
Ethan Furmanc16595e2016-09-10 23:36:59 -0700958
Ethan Furmanee47e5c2016-08-31 00:12:15 -0700959 @classmethod
960 def _missing_(cls, value):
Ethan Furmanc95ad7a2020-09-16 10:26:50 -0700961 return None
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700962
963 def __repr__(self):
Ethan Furmanb7751062021-03-30 21:17:26 -0700964 return "%s.%s" % ( self.__class__.__name__, self._name_)
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700965
966 def __str__(self):
Ethan Furmanb7751062021-03-30 21:17:26 -0700967 return "%s" % (self._name_, )
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700968
Ethan Furman388a3922013-08-12 06:51:41 -0700969 def __dir__(self):
Ethan Furman6d3dfee2020-12-08 12:26:56 -0800970 """
971 Returns all members and all public methods
972 """
Ethan Furman0ae550b2014-10-14 08:58:32 -0700973 added_behavior = [
974 m
975 for cls in self.__class__.mro()
976 for m in cls.__dict__
Ethan Furman354ecf12015-03-11 08:43:12 -0700977 if m[0] != '_' and m not in self._member_map_
Angelin BOOZ68526fe2020-09-21 15:11:06 +0200978 ] + [m for m in self.__dict__ if m[0] != '_']
Ethan Furmanec5f8eb2014-10-21 13:40:35 -0700979 return (['__class__', '__doc__', '__module__'] + added_behavior)
Ethan Furman388a3922013-08-12 06:51:41 -0700980
Ethan Furmanec15a822013-08-31 19:17:41 -0700981 def __format__(self, format_spec):
Ethan Furman6d3dfee2020-12-08 12:26:56 -0800982 """
983 Returns format using actual value type unless __str__ has been overridden.
984 """
Ethan Furmanec15a822013-08-31 19:17:41 -0700985 # mixed-in Enums should use the mixed-in type's __format__, otherwise
986 # we can get strange results with the Enum name showing up instead of
987 # the value
988
thatneat2f19e822019-07-04 11:28:37 -0700989 # pure Enum branch, or branch with __str__ explicitly overridden
Ethan Furman37440ee2020-12-08 11:14:10 -0800990 str_overridden = type(self).__str__ not in (Enum.__str__, Flag.__str__)
thatneat2f19e822019-07-04 11:28:37 -0700991 if self._member_type_ is object or str_overridden:
Ethan Furmanec15a822013-08-31 19:17:41 -0700992 cls = str
993 val = str(self)
994 # mix-in branch
995 else:
996 cls = self._member_type_
Ethan Furman0081f232014-09-16 17:31:23 -0700997 val = self._value_
Ethan Furmanec15a822013-08-31 19:17:41 -0700998 return cls.__format__(val, format_spec)
999
Ethan Furman6b3d64a2013-06-14 16:55:46 -07001000 def __hash__(self):
Ethan Furman520ad572013-07-19 19:47:21 -07001001 return hash(self._name_)
Ethan Furman6b3d64a2013-06-14 16:55:46 -07001002
Ethan Furmanca1b7942014-02-08 11:36:27 -08001003 def __reduce_ex__(self, proto):
Ethan Furmandc870522014-02-18 12:37:12 -08001004 return self.__class__, (self._value_, )
Ethan Furmanca1b7942014-02-08 11:36:27 -08001005
Ethan Furmanc314e602021-01-12 23:47:57 -08001006 # enum.property is used to provide access to the `name` and
1007 # `value` attributes of enum members while keeping some measure of
Ethan Furman6b3d64a2013-06-14 16:55:46 -07001008 # protection from modification, while still allowing for an enumeration
1009 # to have members named `name` and `value`. This works because enumeration
Ethan Furmanc314e602021-01-12 23:47:57 -08001010 # members are not set directly on the enum class; they are kept in a
1011 # separate structure, _member_map_, which is where enum.property looks for
1012 # them
Ethan Furman6b3d64a2013-06-14 16:55:46 -07001013
Ethan Furmanc314e602021-01-12 23:47:57 -08001014 @property
Ethan Furman6b3d64a2013-06-14 16:55:46 -07001015 def name(self):
Ethan Furmanc850f342013-09-15 16:59:35 -07001016 """The name of the Enum member."""
Ethan Furman520ad572013-07-19 19:47:21 -07001017 return self._name_
Ethan Furman6b3d64a2013-06-14 16:55:46 -07001018
Ethan Furmanc314e602021-01-12 23:47:57 -08001019 @property
Ethan Furman6b3d64a2013-06-14 16:55:46 -07001020 def value(self):
Ethan Furmanc850f342013-09-15 16:59:35 -07001021 """The value of the Enum member."""
Ethan Furman520ad572013-07-19 19:47:21 -07001022 return self._value_
Ethan Furman6b3d64a2013-06-14 16:55:46 -07001023
1024
1025class IntEnum(int, Enum):
Ethan Furman0063ff42020-09-21 17:23:13 -07001026 """
1027 Enum where members are also (and must be) ints
1028 """
1029
1030
1031class StrEnum(str, Enum):
1032 """
1033 Enum where members are also (and must be) strings
1034 """
1035
1036 def __new__(cls, *values):
1037 if len(values) > 3:
1038 raise TypeError('too many arguments for str(): %r' % (values, ))
1039 if len(values) == 1:
1040 # it must be a string
1041 if not isinstance(values[0], str):
1042 raise TypeError('%r is not a string' % (values[0], ))
Ethan Furman7aaeb2a2021-01-25 14:26:19 -08001043 if len(values) >= 2:
Ethan Furman0063ff42020-09-21 17:23:13 -07001044 # check that encoding argument is a string
1045 if not isinstance(values[1], str):
1046 raise TypeError('encoding must be a string, not %r' % (values[1], ))
Ethan Furman7aaeb2a2021-01-25 14:26:19 -08001047 if len(values) == 3:
1048 # check that errors argument is a string
1049 if not isinstance(values[2], str):
1050 raise TypeError('errors must be a string, not %r' % (values[2]))
Ethan Furman0063ff42020-09-21 17:23:13 -07001051 value = str(*values)
1052 member = str.__new__(cls, value)
1053 member._value_ = value
1054 return member
Ethan Furmanf24bb352013-07-18 17:05:39 -07001055
Ethan Furmand986d162020-09-22 13:00:07 -07001056 __str__ = str.__str__
1057
Ethan Furmanefb13be2020-12-10 12:20:06 -08001058 def _generate_next_value_(name, start, count, last_values):
1059 """
1060 Return the lower-cased version of the member name.
1061 """
1062 return name.lower()
1063
Ethan Furmanf24bb352013-07-18 17:05:39 -07001064
Ethan Furman24e837f2015-03-18 17:27:57 -07001065def _reduce_ex_by_name(self, proto):
1066 return self.name
1067
Ethan Furman7aaeb2a2021-01-25 14:26:19 -08001068class FlagBoundary(StrEnum):
1069 """
1070 control how out of range values are handled
1071 "strict" -> error is raised [default for Flag]
1072 "conform" -> extra bits are discarded
1073 "eject" -> lose flag status [default for IntFlag]
1074 "keep" -> keep flag status and all bits
1075 """
1076 STRICT = auto()
1077 CONFORM = auto()
1078 EJECT = auto()
1079 KEEP = auto()
1080STRICT, CONFORM, EJECT, KEEP = FlagBoundary
1081
1082
1083class Flag(Enum, boundary=STRICT):
Ethan Furman6d3dfee2020-12-08 12:26:56 -08001084 """
1085 Support for flags
1086 """
Ethan Furmanc16595e2016-09-10 23:36:59 -07001087
1088 def _generate_next_value_(name, start, count, last_values):
Ethan Furmanee47e5c2016-08-31 00:12:15 -07001089 """
1090 Generate the next value when not given.
1091
1092 name: the name of the member
HongWeipengbb16fb22019-09-21 13:22:54 +08001093 start: the initial start value or None
Ethan Furmanee47e5c2016-08-31 00:12:15 -07001094 count: the number of existing members
1095 last_value: the last value assigned or None
1096 """
1097 if not count:
1098 return start if start is not None else 1
Ethan Furman7aaeb2a2021-01-25 14:26:19 -08001099 last_value = max(last_values)
1100 try:
1101 high_bit = _high_bit(last_value)
1102 except Exception:
1103 raise TypeError('Invalid Flag value: %r' % last_value) from None
Ethan Furmanee47e5c2016-08-31 00:12:15 -07001104 return 2 ** (high_bit+1)
1105
1106 @classmethod
Ethan Furman7aaeb2a2021-01-25 14:26:19 -08001107 def _iter_member_by_value_(cls, value):
1108 """
1109 Extract all members from the value in definition (i.e. increasing value) order.
1110 """
1111 for val in _iter_bits_lsb(value & cls._flag_mask_):
1112 yield cls._value2member_map_.get(val)
1113
1114 _iter_member_ = _iter_member_by_value_
1115
1116 @classmethod
1117 def _iter_member_by_def_(cls, value):
1118 """
1119 Extract all members from the value in definition order.
1120 """
1121 yield from sorted(
1122 cls._iter_member_by_value_(value),
1123 key=lambda m: m._sort_order_,
1124 )
1125
1126 @classmethod
Ethan Furmanee47e5c2016-08-31 00:12:15 -07001127 def _missing_(cls, value):
Ethan Furman6d3dfee2020-12-08 12:26:56 -08001128 """
Ethan Furman3515dcc2016-09-18 13:15:41 -07001129 Create a composite member iff value contains only members.
1130 """
Ethan Furman7aaeb2a2021-01-25 14:26:19 -08001131 if not isinstance(value, int):
1132 raise ValueError(
1133 "%r is not a valid %s" % (value, cls.__qualname__)
1134 )
1135 # check boundaries
1136 # - value must be in range (e.g. -16 <-> +15, i.e. ~15 <-> 15)
1137 # - value must not include any skipped flags (e.g. if bit 2 is not
1138 # defined, then 0d10 is invalid)
1139 flag_mask = cls._flag_mask_
1140 all_bits = cls._all_bits_
1141 neg_value = None
1142 if (
1143 not ~all_bits <= value <= all_bits
1144 or value & (all_bits ^ flag_mask)
1145 ):
1146 if cls._boundary_ is STRICT:
1147 max_bits = max(value.bit_length(), flag_mask.bit_length())
1148 raise ValueError(
1149 "%s: invalid value: %r\n given %s\n allowed %s" % (
1150 cls.__name__, value, bin(value, max_bits), bin(flag_mask, max_bits),
1151 ))
1152 elif cls._boundary_ is CONFORM:
1153 value = value & flag_mask
1154 elif cls._boundary_ is EJECT:
1155 return value
1156 elif cls._boundary_ is KEEP:
1157 if value < 0:
1158 value = (
1159 max(all_bits+1, 2**(value.bit_length()))
1160 + value
1161 )
1162 else:
1163 raise ValueError(
1164 'unknown flag boundary: %r' % (cls._boundary_, )
1165 )
1166 if value < 0:
1167 neg_value = value
1168 value = all_bits + 1 + value
1169 # get members and unknown
1170 unknown = value & ~flag_mask
1171 member_value = value & flag_mask
1172 if unknown and cls._boundary_ is not KEEP:
1173 raise ValueError(
1174 '%s(%r) --> unknown values %r [%s]'
1175 % (cls.__name__, value, unknown, bin(unknown))
1176 )
1177 # normal Flag?
1178 __new__ = getattr(cls, '__new_member__', None)
1179 if cls._member_type_ is object and not __new__:
Ethan Furman3515dcc2016-09-18 13:15:41 -07001180 # construct a singleton enum pseudo-member
Ethan Furmanee47e5c2016-08-31 00:12:15 -07001181 pseudo_member = object.__new__(cls)
Ethan Furman7aaeb2a2021-01-25 14:26:19 -08001182 else:
1183 pseudo_member = (__new__ or cls._member_type_.__new__)(cls, value)
1184 if not hasattr(pseudo_member, 'value'):
Ethan Furmanee47e5c2016-08-31 00:12:15 -07001185 pseudo_member._value_ = value
Ethan Furman7aaeb2a2021-01-25 14:26:19 -08001186 if member_value:
1187 pseudo_member._name_ = '|'.join([
1188 m._name_ for m in cls._iter_member_(member_value)
1189 ])
1190 if unknown:
1191 pseudo_member._name_ += '|0x%x' % unknown
1192 else:
1193 pseudo_member._name_ = None
1194 # use setdefault in case another thread already created a composite
1195 # with this value, but only if all members are known
1196 # note: zero is a special case -- add it
1197 if not unknown:
Ethan Furman28cf6632017-01-24 12:12:06 -08001198 pseudo_member = cls._value2member_map_.setdefault(value, pseudo_member)
Ethan Furman7aaeb2a2021-01-25 14:26:19 -08001199 if neg_value is not None:
1200 cls._value2member_map_[neg_value] = pseudo_member
Ethan Furmanee47e5c2016-08-31 00:12:15 -07001201 return pseudo_member
1202
Ethan Furmanee47e5c2016-08-31 00:12:15 -07001203 def __contains__(self, other):
Ethan Furman6d3dfee2020-12-08 12:26:56 -08001204 """
1205 Returns True if self has at least the same flags set as other.
1206 """
Ethan Furmanee47e5c2016-08-31 00:12:15 -07001207 if not isinstance(other, self.__class__):
Rahul Jha94306522018-09-10 23:51:04 +05301208 raise TypeError(
1209 "unsupported operand type(s) for 'in': '%s' and '%s'" % (
1210 type(other).__qualname__, self.__class__.__qualname__))
Ethan Furman7aaeb2a2021-01-25 14:26:19 -08001211 if other._value_ == 0 or self._value_ == 0:
1212 return False
Ethan Furmanee47e5c2016-08-31 00:12:15 -07001213 return other._value_ & self._value_ == other._value_
1214
Ethan Furman7219e272020-09-16 13:01:00 -07001215 def __iter__(self):
Ethan Furman6d3dfee2020-12-08 12:26:56 -08001216 """
Ethan Furman7aaeb2a2021-01-25 14:26:19 -08001217 Returns flags in definition order.
Ethan Furman6d3dfee2020-12-08 12:26:56 -08001218 """
Ethan Furman7aaeb2a2021-01-25 14:26:19 -08001219 yield from self._iter_member_(self._value_)
1220
1221 def __len__(self):
1222 return self._value_.bit_count()
Ethan Furman7219e272020-09-16 13:01:00 -07001223
Ethan Furmanee47e5c2016-08-31 00:12:15 -07001224 def __repr__(self):
Ethan Furmanb7751062021-03-30 21:17:26 -07001225 cls_name = self.__class__.__name__
1226 if self._name_ is None:
1227 return "0x%x" % (self._value_, )
1228 if _is_single_bit(self._value_):
1229 return '%s.%s' % (cls_name, self._name_)
1230 if self._boundary_ is not FlagBoundary.KEEP:
1231 return '%s.' % cls_name + ('|%s.' % cls_name).join(self.name.split('|'))
Ethan Furman7aaeb2a2021-01-25 14:26:19 -08001232 else:
Ethan Furmanb7751062021-03-30 21:17:26 -07001233 name = []
1234 for n in self._name_.split('|'):
1235 if n.startswith('0'):
1236 name.append(n)
1237 else:
1238 name.append('%s.%s' % (cls_name, n))
1239 return '|'.join(name)
Ethan Furmanee47e5c2016-08-31 00:12:15 -07001240
1241 def __str__(self):
1242 cls = self.__class__
Ethan Furmanb7751062021-03-30 21:17:26 -07001243 if self._name_ is None:
1244 return '%s(%x)' % (cls.__name__, self._value_)
Ethan Furmanee47e5c2016-08-31 00:12:15 -07001245 else:
Ethan Furmanb7751062021-03-30 21:17:26 -07001246 return self._name_
Ethan Furmanee47e5c2016-08-31 00:12:15 -07001247
Ethan Furman25d94bb2016-09-02 16:32:32 -07001248 def __bool__(self):
1249 return bool(self._value_)
1250
Ethan Furmanee47e5c2016-08-31 00:12:15 -07001251 def __or__(self, other):
1252 if not isinstance(other, self.__class__):
1253 return NotImplemented
1254 return self.__class__(self._value_ | other._value_)
1255
1256 def __and__(self, other):
1257 if not isinstance(other, self.__class__):
1258 return NotImplemented
1259 return self.__class__(self._value_ & other._value_)
1260
1261 def __xor__(self, other):
1262 if not isinstance(other, self.__class__):
1263 return NotImplemented
1264 return self.__class__(self._value_ ^ other._value_)
1265
1266 def __invert__(self):
Ethan Furman7aaeb2a2021-01-25 14:26:19 -08001267 if self._inverted_ is None:
1268 if self._boundary_ is KEEP:
1269 # use all bits
1270 self._inverted_ = self.__class__(~self._value_)
1271 else:
1272 # calculate flags not in this member
1273 self._inverted_ = self.__class__(self._flag_mask_ ^ self._value_)
1274 self._inverted_._inverted_ = self
1275 return self._inverted_
Ethan Furmanee47e5c2016-08-31 00:12:15 -07001276
1277
Ethan Furman7aaeb2a2021-01-25 14:26:19 -08001278class IntFlag(int, Flag, boundary=EJECT):
Ethan Furman6d3dfee2020-12-08 12:26:56 -08001279 """
1280 Support for integer-based Flags
1281 """
Ethan Furmanee47e5c2016-08-31 00:12:15 -07001282
Ethan Furmanee47e5c2016-08-31 00:12:15 -07001283 def __or__(self, other):
Ethan Furman7aaeb2a2021-01-25 14:26:19 -08001284 if isinstance(other, self.__class__):
1285 other = other._value_
1286 elif isinstance(other, int):
1287 other = other
1288 else:
Ethan Furmanee47e5c2016-08-31 00:12:15 -07001289 return NotImplemented
Ethan Furman7aaeb2a2021-01-25 14:26:19 -08001290 value = self._value_
1291 return self.__class__(value | other)
Ethan Furmanee47e5c2016-08-31 00:12:15 -07001292
1293 def __and__(self, other):
Ethan Furman7aaeb2a2021-01-25 14:26:19 -08001294 if isinstance(other, self.__class__):
1295 other = other._value_
1296 elif isinstance(other, int):
1297 other = other
1298 else:
Ethan Furmanee47e5c2016-08-31 00:12:15 -07001299 return NotImplemented
Ethan Furman7aaeb2a2021-01-25 14:26:19 -08001300 value = self._value_
1301 return self.__class__(value & other)
Ethan Furmanee47e5c2016-08-31 00:12:15 -07001302
1303 def __xor__(self, other):
Ethan Furman7aaeb2a2021-01-25 14:26:19 -08001304 if isinstance(other, self.__class__):
1305 other = other._value_
1306 elif isinstance(other, int):
1307 other = other
1308 else:
Ethan Furmanee47e5c2016-08-31 00:12:15 -07001309 return NotImplemented
Ethan Furman7aaeb2a2021-01-25 14:26:19 -08001310 value = self._value_
1311 return self.__class__(value ^ other)
Ethan Furmanee47e5c2016-08-31 00:12:15 -07001312
1313 __ror__ = __or__
1314 __rand__ = __and__
1315 __rxor__ = __xor__
Ethan Furman7aaeb2a2021-01-25 14:26:19 -08001316 __invert__ = Flag.__invert__
Ethan Furmanee47e5c2016-08-31 00:12:15 -07001317
1318def _high_bit(value):
Ethan Furman6d3dfee2020-12-08 12:26:56 -08001319 """
1320 returns index of highest bit, or -1 if value is zero or negative
1321 """
Ethan Furman3515dcc2016-09-18 13:15:41 -07001322 return value.bit_length() - 1
Ethan Furmanee47e5c2016-08-31 00:12:15 -07001323
Ethan Furmanf24bb352013-07-18 17:05:39 -07001324def unique(enumeration):
Ethan Furman6d3dfee2020-12-08 12:26:56 -08001325 """
1326 Class decorator for enumerations ensuring unique member values.
1327 """
Ethan Furmanf24bb352013-07-18 17:05:39 -07001328 duplicates = []
1329 for name, member in enumeration.__members__.items():
1330 if name != member.name:
1331 duplicates.append((name, member.name))
1332 if duplicates:
1333 alias_details = ', '.join(
1334 ["%s -> %s" % (alias, name) for (alias, name) in duplicates])
1335 raise ValueError('duplicate values found in %r: %s' %
1336 (enumeration, alias_details))
1337 return enumeration
Ethan Furman3515dcc2016-09-18 13:15:41 -07001338
Ethan Furman7aaeb2a2021-01-25 14:26:19 -08001339def _power_of_two(value):
1340 if value < 1:
1341 return False
1342 return value == 2 ** _high_bit(value)
Ethan Furmanb7751062021-03-30 21:17:26 -07001343
1344def global_enum_repr(self):
1345 return '%s.%s' % (self.__class__.__module__, self._name_)
1346
1347def global_flag_repr(self):
1348 module = self.__class__.__module__
1349 cls_name = self.__class__.__name__
1350 if self._name_ is None:
1351 return "%x" % (module, cls_name, self._value_)
1352 if _is_single_bit(self):
1353 return '%s.%s' % (module, self._name_)
1354 if self._boundary_ is not FlagBoundary.KEEP:
1355 return module + module.join(self.name.split('|'))
1356 else:
1357 name = []
1358 for n in self._name_.split('|'):
1359 if n.startswith('0'):
1360 name.append(n)
1361 else:
1362 name.append('%s.%s' % (module, n))
1363 return '|'.join(name)
1364
1365
1366def global_enum(cls):
1367 """
1368 decorator that makes the repr() of an enum member reference its module
1369 instead of its class; also exports all members to the enum's module's
1370 global namespace
1371 """
1372 if issubclass(cls, Flag):
1373 cls.__repr__ = global_flag_repr
1374 else:
1375 cls.__repr__ = global_enum_repr
1376 sys.modules[cls.__module__].__dict__.update(cls.__members__)
1377 return cls