blob: b102a4e4cc78202b05b91eb08fcef6a2fdff5e09 [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
Ethan Furman8c14f5a2021-04-12 08:51:20 -0700922 try:
923 if isinstance(result, cls):
924 return result
925 elif (
926 Flag is not None and issubclass(cls, Flag)
927 and cls._boundary_ is EJECT and isinstance(result, int)
928 ):
929 return result
930 else:
931 ve_exc = ValueError("%r is not a valid %s" % (value, cls.__qualname__))
932 if result is None and exc is None:
933 raise ve_exc
934 elif exc is None:
935 exc = TypeError(
936 'error in %s._missing_: returned %r instead of None or a valid member'
937 % (cls.__name__, result)
938 )
939 if not isinstance(exc, ValueError):
940 exc.__context__ = ve_exc
941 raise exc
942 finally:
943 # ensure all variables that could hold an exception are destroyed
944 exc = None
945 ve_exc = None
Ethan Furmanee47e5c2016-08-31 00:12:15 -0700946
Ethan Furmanc16595e2016-09-10 23:36:59 -0700947 def _generate_next_value_(name, start, count, last_values):
Ethan Furman6d3dfee2020-12-08 12:26:56 -0800948 """
949 Generate the next value when not given.
950
951 name: the name of the member
952 start: the initial start value or None
953 count: the number of existing members
954 last_value: the last value assigned or None
955 """
Ethan Furmanc16595e2016-09-10 23:36:59 -0700956 for last_value in reversed(last_values):
957 try:
958 return last_value + 1
959 except TypeError:
960 pass
961 else:
Ethan Furmanee47e5c2016-08-31 00:12:15 -0700962 return start
Ethan Furmanc16595e2016-09-10 23:36:59 -0700963
Ethan Furmanee47e5c2016-08-31 00:12:15 -0700964 @classmethod
965 def _missing_(cls, value):
Ethan Furmanc95ad7a2020-09-16 10:26:50 -0700966 return None
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700967
968 def __repr__(self):
Ethan Furmanb7751062021-03-30 21:17:26 -0700969 return "%s.%s" % ( self.__class__.__name__, self._name_)
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700970
971 def __str__(self):
Ethan Furmanb7751062021-03-30 21:17:26 -0700972 return "%s" % (self._name_, )
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700973
Ethan Furman388a3922013-08-12 06:51:41 -0700974 def __dir__(self):
Ethan Furman6d3dfee2020-12-08 12:26:56 -0800975 """
976 Returns all members and all public methods
977 """
Ethan Furman0ae550b2014-10-14 08:58:32 -0700978 added_behavior = [
979 m
980 for cls in self.__class__.mro()
981 for m in cls.__dict__
Ethan Furman354ecf12015-03-11 08:43:12 -0700982 if m[0] != '_' and m not in self._member_map_
Angelin BOOZ68526fe2020-09-21 15:11:06 +0200983 ] + [m for m in self.__dict__ if m[0] != '_']
Ethan Furmanec5f8eb2014-10-21 13:40:35 -0700984 return (['__class__', '__doc__', '__module__'] + added_behavior)
Ethan Furman388a3922013-08-12 06:51:41 -0700985
Ethan Furmanec15a822013-08-31 19:17:41 -0700986 def __format__(self, format_spec):
Ethan Furman6d3dfee2020-12-08 12:26:56 -0800987 """
988 Returns format using actual value type unless __str__ has been overridden.
989 """
Ethan Furmanec15a822013-08-31 19:17:41 -0700990 # mixed-in Enums should use the mixed-in type's __format__, otherwise
991 # we can get strange results with the Enum name showing up instead of
992 # the value
993
thatneat2f19e822019-07-04 11:28:37 -0700994 # pure Enum branch, or branch with __str__ explicitly overridden
Ethan Furman37440ee2020-12-08 11:14:10 -0800995 str_overridden = type(self).__str__ not in (Enum.__str__, Flag.__str__)
thatneat2f19e822019-07-04 11:28:37 -0700996 if self._member_type_ is object or str_overridden:
Ethan Furmanec15a822013-08-31 19:17:41 -0700997 cls = str
998 val = str(self)
999 # mix-in branch
1000 else:
1001 cls = self._member_type_
Ethan Furman0081f232014-09-16 17:31:23 -07001002 val = self._value_
Ethan Furmanec15a822013-08-31 19:17:41 -07001003 return cls.__format__(val, format_spec)
1004
Ethan Furman6b3d64a2013-06-14 16:55:46 -07001005 def __hash__(self):
Ethan Furman520ad572013-07-19 19:47:21 -07001006 return hash(self._name_)
Ethan Furman6b3d64a2013-06-14 16:55:46 -07001007
Ethan Furmanca1b7942014-02-08 11:36:27 -08001008 def __reduce_ex__(self, proto):
Ethan Furmandc870522014-02-18 12:37:12 -08001009 return self.__class__, (self._value_, )
Ethan Furmanca1b7942014-02-08 11:36:27 -08001010
Ethan Furmanc314e602021-01-12 23:47:57 -08001011 # enum.property is used to provide access to the `name` and
1012 # `value` attributes of enum members while keeping some measure of
Ethan Furman6b3d64a2013-06-14 16:55:46 -07001013 # protection from modification, while still allowing for an enumeration
1014 # to have members named `name` and `value`. This works because enumeration
Ethan Furmanc314e602021-01-12 23:47:57 -08001015 # members are not set directly on the enum class; they are kept in a
1016 # separate structure, _member_map_, which is where enum.property looks for
1017 # them
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 name(self):
Ethan Furmanc850f342013-09-15 16:59:35 -07001021 """The name of the Enum member."""
Ethan Furman520ad572013-07-19 19:47:21 -07001022 return self._name_
Ethan Furman6b3d64a2013-06-14 16:55:46 -07001023
Ethan Furmanc314e602021-01-12 23:47:57 -08001024 @property
Ethan Furman6b3d64a2013-06-14 16:55:46 -07001025 def value(self):
Ethan Furmanc850f342013-09-15 16:59:35 -07001026 """The value of the Enum member."""
Ethan Furman520ad572013-07-19 19:47:21 -07001027 return self._value_
Ethan Furman6b3d64a2013-06-14 16:55:46 -07001028
1029
1030class IntEnum(int, Enum):
Ethan Furman0063ff42020-09-21 17:23:13 -07001031 """
1032 Enum where members are also (and must be) ints
1033 """
1034
1035
1036class StrEnum(str, Enum):
1037 """
1038 Enum where members are also (and must be) strings
1039 """
1040
1041 def __new__(cls, *values):
1042 if len(values) > 3:
1043 raise TypeError('too many arguments for str(): %r' % (values, ))
1044 if len(values) == 1:
1045 # it must be a string
1046 if not isinstance(values[0], str):
1047 raise TypeError('%r is not a string' % (values[0], ))
Ethan Furman7aaeb2a2021-01-25 14:26:19 -08001048 if len(values) >= 2:
Ethan Furman0063ff42020-09-21 17:23:13 -07001049 # check that encoding argument is a string
1050 if not isinstance(values[1], str):
1051 raise TypeError('encoding must be a string, not %r' % (values[1], ))
Ethan Furman7aaeb2a2021-01-25 14:26:19 -08001052 if len(values) == 3:
1053 # check that errors argument is a string
1054 if not isinstance(values[2], str):
1055 raise TypeError('errors must be a string, not %r' % (values[2]))
Ethan Furman0063ff42020-09-21 17:23:13 -07001056 value = str(*values)
1057 member = str.__new__(cls, value)
1058 member._value_ = value
1059 return member
Ethan Furmanf24bb352013-07-18 17:05:39 -07001060
Ethan Furmand986d162020-09-22 13:00:07 -07001061 __str__ = str.__str__
1062
Ethan Furmanefb13be2020-12-10 12:20:06 -08001063 def _generate_next_value_(name, start, count, last_values):
1064 """
1065 Return the lower-cased version of the member name.
1066 """
1067 return name.lower()
1068
Ethan Furmanf24bb352013-07-18 17:05:39 -07001069
Ethan Furman24e837f2015-03-18 17:27:57 -07001070def _reduce_ex_by_name(self, proto):
1071 return self.name
1072
Ethan Furman7aaeb2a2021-01-25 14:26:19 -08001073class FlagBoundary(StrEnum):
1074 """
1075 control how out of range values are handled
1076 "strict" -> error is raised [default for Flag]
1077 "conform" -> extra bits are discarded
1078 "eject" -> lose flag status [default for IntFlag]
1079 "keep" -> keep flag status and all bits
1080 """
1081 STRICT = auto()
1082 CONFORM = auto()
1083 EJECT = auto()
1084 KEEP = auto()
1085STRICT, CONFORM, EJECT, KEEP = FlagBoundary
1086
1087
1088class Flag(Enum, boundary=STRICT):
Ethan Furman6d3dfee2020-12-08 12:26:56 -08001089 """
1090 Support for flags
1091 """
Ethan Furmanc16595e2016-09-10 23:36:59 -07001092
1093 def _generate_next_value_(name, start, count, last_values):
Ethan Furmanee47e5c2016-08-31 00:12:15 -07001094 """
1095 Generate the next value when not given.
1096
1097 name: the name of the member
HongWeipengbb16fb22019-09-21 13:22:54 +08001098 start: the initial start value or None
Ethan Furmanee47e5c2016-08-31 00:12:15 -07001099 count: the number of existing members
1100 last_value: the last value assigned or None
1101 """
1102 if not count:
1103 return start if start is not None else 1
Ethan Furman7aaeb2a2021-01-25 14:26:19 -08001104 last_value = max(last_values)
1105 try:
1106 high_bit = _high_bit(last_value)
1107 except Exception:
1108 raise TypeError('Invalid Flag value: %r' % last_value) from None
Ethan Furmanee47e5c2016-08-31 00:12:15 -07001109 return 2 ** (high_bit+1)
1110
1111 @classmethod
Ethan Furman7aaeb2a2021-01-25 14:26:19 -08001112 def _iter_member_by_value_(cls, value):
1113 """
1114 Extract all members from the value in definition (i.e. increasing value) order.
1115 """
1116 for val in _iter_bits_lsb(value & cls._flag_mask_):
1117 yield cls._value2member_map_.get(val)
1118
1119 _iter_member_ = _iter_member_by_value_
1120
1121 @classmethod
1122 def _iter_member_by_def_(cls, value):
1123 """
1124 Extract all members from the value in definition order.
1125 """
1126 yield from sorted(
1127 cls._iter_member_by_value_(value),
1128 key=lambda m: m._sort_order_,
1129 )
1130
1131 @classmethod
Ethan Furmanee47e5c2016-08-31 00:12:15 -07001132 def _missing_(cls, value):
Ethan Furman6d3dfee2020-12-08 12:26:56 -08001133 """
Ethan Furman3515dcc2016-09-18 13:15:41 -07001134 Create a composite member iff value contains only members.
1135 """
Ethan Furman7aaeb2a2021-01-25 14:26:19 -08001136 if not isinstance(value, int):
1137 raise ValueError(
1138 "%r is not a valid %s" % (value, cls.__qualname__)
1139 )
1140 # check boundaries
1141 # - value must be in range (e.g. -16 <-> +15, i.e. ~15 <-> 15)
1142 # - value must not include any skipped flags (e.g. if bit 2 is not
1143 # defined, then 0d10 is invalid)
1144 flag_mask = cls._flag_mask_
1145 all_bits = cls._all_bits_
1146 neg_value = None
1147 if (
1148 not ~all_bits <= value <= all_bits
1149 or value & (all_bits ^ flag_mask)
1150 ):
1151 if cls._boundary_ is STRICT:
1152 max_bits = max(value.bit_length(), flag_mask.bit_length())
1153 raise ValueError(
1154 "%s: invalid value: %r\n given %s\n allowed %s" % (
1155 cls.__name__, value, bin(value, max_bits), bin(flag_mask, max_bits),
1156 ))
1157 elif cls._boundary_ is CONFORM:
1158 value = value & flag_mask
1159 elif cls._boundary_ is EJECT:
1160 return value
1161 elif cls._boundary_ is KEEP:
1162 if value < 0:
1163 value = (
1164 max(all_bits+1, 2**(value.bit_length()))
1165 + value
1166 )
1167 else:
1168 raise ValueError(
1169 'unknown flag boundary: %r' % (cls._boundary_, )
1170 )
1171 if value < 0:
1172 neg_value = value
1173 value = all_bits + 1 + value
1174 # get members and unknown
1175 unknown = value & ~flag_mask
1176 member_value = value & flag_mask
1177 if unknown and cls._boundary_ is not KEEP:
1178 raise ValueError(
1179 '%s(%r) --> unknown values %r [%s]'
1180 % (cls.__name__, value, unknown, bin(unknown))
1181 )
1182 # normal Flag?
1183 __new__ = getattr(cls, '__new_member__', None)
1184 if cls._member_type_ is object and not __new__:
Ethan Furman3515dcc2016-09-18 13:15:41 -07001185 # construct a singleton enum pseudo-member
Ethan Furmanee47e5c2016-08-31 00:12:15 -07001186 pseudo_member = object.__new__(cls)
Ethan Furman7aaeb2a2021-01-25 14:26:19 -08001187 else:
1188 pseudo_member = (__new__ or cls._member_type_.__new__)(cls, value)
1189 if not hasattr(pseudo_member, 'value'):
Ethan Furmanee47e5c2016-08-31 00:12:15 -07001190 pseudo_member._value_ = value
Ethan Furman7aaeb2a2021-01-25 14:26:19 -08001191 if member_value:
1192 pseudo_member._name_ = '|'.join([
1193 m._name_ for m in cls._iter_member_(member_value)
1194 ])
1195 if unknown:
1196 pseudo_member._name_ += '|0x%x' % unknown
1197 else:
1198 pseudo_member._name_ = None
1199 # use setdefault in case another thread already created a composite
1200 # with this value, but only if all members are known
1201 # note: zero is a special case -- add it
1202 if not unknown:
Ethan Furman28cf6632017-01-24 12:12:06 -08001203 pseudo_member = cls._value2member_map_.setdefault(value, pseudo_member)
Ethan Furman7aaeb2a2021-01-25 14:26:19 -08001204 if neg_value is not None:
1205 cls._value2member_map_[neg_value] = pseudo_member
Ethan Furmanee47e5c2016-08-31 00:12:15 -07001206 return pseudo_member
1207
Ethan Furmanee47e5c2016-08-31 00:12:15 -07001208 def __contains__(self, other):
Ethan Furman6d3dfee2020-12-08 12:26:56 -08001209 """
1210 Returns True if self has at least the same flags set as other.
1211 """
Ethan Furmanee47e5c2016-08-31 00:12:15 -07001212 if not isinstance(other, self.__class__):
Rahul Jha94306522018-09-10 23:51:04 +05301213 raise TypeError(
1214 "unsupported operand type(s) for 'in': '%s' and '%s'" % (
1215 type(other).__qualname__, self.__class__.__qualname__))
Ethan Furman7aaeb2a2021-01-25 14:26:19 -08001216 if other._value_ == 0 or self._value_ == 0:
1217 return False
Ethan Furmanee47e5c2016-08-31 00:12:15 -07001218 return other._value_ & self._value_ == other._value_
1219
Ethan Furman7219e272020-09-16 13:01:00 -07001220 def __iter__(self):
Ethan Furman6d3dfee2020-12-08 12:26:56 -08001221 """
Ethan Furman7aaeb2a2021-01-25 14:26:19 -08001222 Returns flags in definition order.
Ethan Furman6d3dfee2020-12-08 12:26:56 -08001223 """
Ethan Furman7aaeb2a2021-01-25 14:26:19 -08001224 yield from self._iter_member_(self._value_)
1225
1226 def __len__(self):
1227 return self._value_.bit_count()
Ethan Furman7219e272020-09-16 13:01:00 -07001228
Ethan Furmanee47e5c2016-08-31 00:12:15 -07001229 def __repr__(self):
Ethan Furmanb7751062021-03-30 21:17:26 -07001230 cls_name = self.__class__.__name__
1231 if self._name_ is None:
1232 return "0x%x" % (self._value_, )
1233 if _is_single_bit(self._value_):
1234 return '%s.%s' % (cls_name, self._name_)
1235 if self._boundary_ is not FlagBoundary.KEEP:
1236 return '%s.' % cls_name + ('|%s.' % cls_name).join(self.name.split('|'))
Ethan Furman7aaeb2a2021-01-25 14:26:19 -08001237 else:
Ethan Furmanb7751062021-03-30 21:17:26 -07001238 name = []
1239 for n in self._name_.split('|'):
1240 if n.startswith('0'):
1241 name.append(n)
1242 else:
1243 name.append('%s.%s' % (cls_name, n))
1244 return '|'.join(name)
Ethan Furmanee47e5c2016-08-31 00:12:15 -07001245
1246 def __str__(self):
1247 cls = self.__class__
Ethan Furmanb7751062021-03-30 21:17:26 -07001248 if self._name_ is None:
1249 return '%s(%x)' % (cls.__name__, self._value_)
Ethan Furmanee47e5c2016-08-31 00:12:15 -07001250 else:
Ethan Furmanb7751062021-03-30 21:17:26 -07001251 return self._name_
Ethan Furmanee47e5c2016-08-31 00:12:15 -07001252
Ethan Furman25d94bb2016-09-02 16:32:32 -07001253 def __bool__(self):
1254 return bool(self._value_)
1255
Ethan Furmanee47e5c2016-08-31 00:12:15 -07001256 def __or__(self, other):
1257 if not isinstance(other, self.__class__):
1258 return NotImplemented
1259 return self.__class__(self._value_ | other._value_)
1260
1261 def __and__(self, other):
1262 if not isinstance(other, self.__class__):
1263 return NotImplemented
1264 return self.__class__(self._value_ & other._value_)
1265
1266 def __xor__(self, other):
1267 if not isinstance(other, self.__class__):
1268 return NotImplemented
1269 return self.__class__(self._value_ ^ other._value_)
1270
1271 def __invert__(self):
Ethan Furman7aaeb2a2021-01-25 14:26:19 -08001272 if self._inverted_ is None:
1273 if self._boundary_ is KEEP:
1274 # use all bits
1275 self._inverted_ = self.__class__(~self._value_)
1276 else:
1277 # calculate flags not in this member
1278 self._inverted_ = self.__class__(self._flag_mask_ ^ self._value_)
1279 self._inverted_._inverted_ = self
1280 return self._inverted_
Ethan Furmanee47e5c2016-08-31 00:12:15 -07001281
1282
Ethan Furman7aaeb2a2021-01-25 14:26:19 -08001283class IntFlag(int, Flag, boundary=EJECT):
Ethan Furman6d3dfee2020-12-08 12:26:56 -08001284 """
1285 Support for integer-based Flags
1286 """
Ethan Furmanee47e5c2016-08-31 00:12:15 -07001287
Ethan Furmanee47e5c2016-08-31 00:12:15 -07001288 def __or__(self, other):
Ethan Furman7aaeb2a2021-01-25 14:26:19 -08001289 if isinstance(other, self.__class__):
1290 other = other._value_
1291 elif isinstance(other, int):
1292 other = other
1293 else:
Ethan Furmanee47e5c2016-08-31 00:12:15 -07001294 return NotImplemented
Ethan Furman7aaeb2a2021-01-25 14:26:19 -08001295 value = self._value_
1296 return self.__class__(value | other)
Ethan Furmanee47e5c2016-08-31 00:12:15 -07001297
1298 def __and__(self, other):
Ethan Furman7aaeb2a2021-01-25 14:26:19 -08001299 if isinstance(other, self.__class__):
1300 other = other._value_
1301 elif isinstance(other, int):
1302 other = other
1303 else:
Ethan Furmanee47e5c2016-08-31 00:12:15 -07001304 return NotImplemented
Ethan Furman7aaeb2a2021-01-25 14:26:19 -08001305 value = self._value_
1306 return self.__class__(value & other)
Ethan Furmanee47e5c2016-08-31 00:12:15 -07001307
1308 def __xor__(self, other):
Ethan Furman7aaeb2a2021-01-25 14:26:19 -08001309 if isinstance(other, self.__class__):
1310 other = other._value_
1311 elif isinstance(other, int):
1312 other = other
1313 else:
Ethan Furmanee47e5c2016-08-31 00:12:15 -07001314 return NotImplemented
Ethan Furman7aaeb2a2021-01-25 14:26:19 -08001315 value = self._value_
1316 return self.__class__(value ^ other)
Ethan Furmanee47e5c2016-08-31 00:12:15 -07001317
1318 __ror__ = __or__
1319 __rand__ = __and__
1320 __rxor__ = __xor__
Ethan Furman7aaeb2a2021-01-25 14:26:19 -08001321 __invert__ = Flag.__invert__
Ethan Furmanee47e5c2016-08-31 00:12:15 -07001322
1323def _high_bit(value):
Ethan Furman6d3dfee2020-12-08 12:26:56 -08001324 """
1325 returns index of highest bit, or -1 if value is zero or negative
1326 """
Ethan Furman3515dcc2016-09-18 13:15:41 -07001327 return value.bit_length() - 1
Ethan Furmanee47e5c2016-08-31 00:12:15 -07001328
Ethan Furmanf24bb352013-07-18 17:05:39 -07001329def unique(enumeration):
Ethan Furman6d3dfee2020-12-08 12:26:56 -08001330 """
1331 Class decorator for enumerations ensuring unique member values.
1332 """
Ethan Furmanf24bb352013-07-18 17:05:39 -07001333 duplicates = []
1334 for name, member in enumeration.__members__.items():
1335 if name != member.name:
1336 duplicates.append((name, member.name))
1337 if duplicates:
1338 alias_details = ', '.join(
1339 ["%s -> %s" % (alias, name) for (alias, name) in duplicates])
1340 raise ValueError('duplicate values found in %r: %s' %
1341 (enumeration, alias_details))
1342 return enumeration
Ethan Furman3515dcc2016-09-18 13:15:41 -07001343
Ethan Furman7aaeb2a2021-01-25 14:26:19 -08001344def _power_of_two(value):
1345 if value < 1:
1346 return False
1347 return value == 2 ** _high_bit(value)
Ethan Furmanb7751062021-03-30 21:17:26 -07001348
1349def global_enum_repr(self):
1350 return '%s.%s' % (self.__class__.__module__, self._name_)
1351
1352def global_flag_repr(self):
1353 module = self.__class__.__module__
1354 cls_name = self.__class__.__name__
1355 if self._name_ is None:
1356 return "%x" % (module, cls_name, self._value_)
1357 if _is_single_bit(self):
1358 return '%s.%s' % (module, self._name_)
1359 if self._boundary_ is not FlagBoundary.KEEP:
1360 return module + module.join(self.name.split('|'))
1361 else:
1362 name = []
1363 for n in self._name_.split('|'):
1364 if n.startswith('0'):
1365 name.append(n)
1366 else:
1367 name.append('%s.%s' % (module, n))
1368 return '|'.join(name)
1369
1370
1371def global_enum(cls):
1372 """
1373 decorator that makes the repr() of an enum member reference its module
1374 instead of its class; also exports all members to the enum's module's
1375 global namespace
1376 """
1377 if issubclass(cls, Flag):
1378 cls.__repr__ = global_flag_repr
1379 else:
1380 cls.__repr__ = global_enum_repr
1381 sys.modules[cls.__module__].__dict__.update(cls.__members__)
1382 return cls