blob: 17deb4b9c05ae70f36b18c8cbf7c495dd4ddcd5e [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, )
Ethan Furmanec099732021-04-15 06:58:33 -070056 pat_len = len(pattern)
Ethan Furman7cf0aad2020-12-09 17:12:11 -080057 if (
Ethan Furmanec099732021-04-15 06:58:33 -070058 len(name) > pat_len
Ethan Furman7cf0aad2020-12-09 17:12:11 -080059 and name.startswith(pattern)
Ethan Furmanec099732021-04-15 06:58:33 -070060 and name[pat_len:pat_len+1] != ['_']
Ethan Furman7cf0aad2020-12-09 17:12:11 -080061 and (name[-1] != '_' or name[-2] != '_')
62 ):
63 return True
64 else:
65 return False
66
Ethan Furman7aaeb2a2021-01-25 14:26:19 -080067def _is_single_bit(num):
68 """
69 True if only one bit set in num (should be an int)
70 """
71 if num == 0:
72 return False
73 num &= num - 1
74 return num == 0
75
Ethan Furmanc314e602021-01-12 23:47:57 -080076def _make_class_unpicklable(obj):
Ethan Furman6d3dfee2020-12-08 12:26:56 -080077 """
Ethan Furmanc314e602021-01-12 23:47:57 -080078 Make the given obj un-picklable.
79
80 obj should be either a dictionary, on an Enum
Ethan Furman6d3dfee2020-12-08 12:26:56 -080081 """
Ethan Furmanca1b7942014-02-08 11:36:27 -080082 def _break_on_call_reduce(self, proto):
Ethan Furman6b3d64a2013-06-14 16:55:46 -070083 raise TypeError('%r cannot be pickled' % self)
Ethan Furmanc314e602021-01-12 23:47:57 -080084 if isinstance(obj, dict):
85 obj['__reduce_ex__'] = _break_on_call_reduce
86 obj['__module__'] = '<unknown>'
87 else:
88 setattr(obj, '__reduce_ex__', _break_on_call_reduce)
89 setattr(obj, '__module__', '<unknown>')
Ethan Furman6b3d64a2013-06-14 16:55:46 -070090
Ethan Furman7aaeb2a2021-01-25 14:26:19 -080091def _iter_bits_lsb(num):
92 while num:
93 b = num & (~num + 1)
94 yield b
95 num ^= b
96
97def bin(num, max_bits=None):
98 """
99 Like built-in bin(), except negative values are represented in
100 twos-compliment, and the leading bit always indicates sign
101 (0=positive, 1=negative).
102
103 >>> bin(10)
104 '0b0 1010'
105 >>> bin(~10) # ~10 is -11
106 '0b1 0101'
107 """
108
109 ceiling = 2 ** (num).bit_length()
110 if num >= 0:
111 s = _bltin_bin(num + ceiling).replace('1', '0', 1)
112 else:
113 s = _bltin_bin(~num ^ (ceiling - 1) + ceiling)
114 sign = s[:3]
115 digits = s[3:]
116 if max_bits is not None:
117 if len(digits) < max_bits:
118 digits = (sign[-1] * max_bits + digits)[-max_bits:]
119 return "%s %s" % (sign, digits)
120
121
Ethan Furman3515dcc2016-09-18 13:15:41 -0700122_auto_null = object()
Ethan Furmanc16595e2016-09-10 23:36:59 -0700123class auto:
124 """
125 Instances are replaced with an appropriate value in Enum class suites.
126 """
Ethan Furman3515dcc2016-09-18 13:15:41 -0700127 value = _auto_null
Ethan Furmanc16595e2016-09-10 23:36:59 -0700128
Ethan Furmanc314e602021-01-12 23:47:57 -0800129class property(DynamicClassAttribute):
130 """
131 This is a descriptor, used to define attributes that act differently
132 when accessed through an enum member and through an enum class.
133 Instance access is the same as property(), but access to an attribute
134 through the enum class will instead look in the class' _member_map_ for
135 a corresponding enum member.
136 """
137
138 def __get__(self, instance, ownerclass=None):
139 if instance is None:
140 try:
141 return ownerclass._member_map_[self.name]
142 except KeyError:
Ethan Furman7aaeb2a2021-01-25 14:26:19 -0800143 raise AttributeError(
Ethan Furmand65b9032021-02-08 17:32:38 -0800144 '%s: no class attribute %r' % (ownerclass.__name__, self.name)
Ethan Furman7aaeb2a2021-01-25 14:26:19 -0800145 )
Ethan Furmanc314e602021-01-12 23:47:57 -0800146 else:
147 if self.fget is None:
Ethan Furmand65b9032021-02-08 17:32:38 -0800148 # check for member
149 if self.name in ownerclass._member_map_:
150 import warnings
151 warnings.warn(
152 "accessing one member from another is not supported, "
Ethan Furman44e580f2021-03-03 09:54:30 -0800153 " and will be disabled in 3.12",
Ethan Furmand65b9032021-02-08 17:32:38 -0800154 DeprecationWarning,
155 stacklevel=2,
156 )
157 return ownerclass._member_map_[self.name]
Ethan Furman7aaeb2a2021-01-25 14:26:19 -0800158 raise AttributeError(
Ethan Furmand65b9032021-02-08 17:32:38 -0800159 '%s: no instance attribute %r' % (ownerclass.__name__, self.name)
Ethan Furman7aaeb2a2021-01-25 14:26:19 -0800160 )
Ethan Furmanc314e602021-01-12 23:47:57 -0800161 else:
162 return self.fget(instance)
163
164 def __set__(self, instance, value):
165 if self.fset is None:
Ethan Furman7aaeb2a2021-01-25 14:26:19 -0800166 raise AttributeError(
Ethan Furmand65b9032021-02-08 17:32:38 -0800167 "%s: cannot set instance attribute %r" % (self.clsname, self.name)
Ethan Furman7aaeb2a2021-01-25 14:26:19 -0800168 )
Ethan Furmanc314e602021-01-12 23:47:57 -0800169 else:
170 return self.fset(instance, value)
171
172 def __delete__(self, instance):
173 if self.fdel is None:
Ethan Furman7aaeb2a2021-01-25 14:26:19 -0800174 raise AttributeError(
Ethan Furmand65b9032021-02-08 17:32:38 -0800175 "%s: cannot delete instance attribute %r" % (self.clsname, self.name)
Ethan Furman7aaeb2a2021-01-25 14:26:19 -0800176 )
Ethan Furmanc314e602021-01-12 23:47:57 -0800177 else:
178 return self.fdel(instance)
179
180 def __set_name__(self, ownerclass, name):
181 self.name = name
182 self.clsname = ownerclass.__name__
183
184
185class _proto_member:
186 """
187 intermediate step for enum members between class execution and final creation
188 """
189
190 def __init__(self, value):
191 self.value = value
192
193 def __set_name__(self, enum_class, member_name):
194 """
195 convert each quasi-member into an instance of the new enum class
196 """
197 # first step: remove ourself from enum_class
198 delattr(enum_class, member_name)
199 # second step: create member based on enum_class
200 value = self.value
201 if not isinstance(value, tuple):
202 args = (value, )
203 else:
204 args = value
205 if enum_class._member_type_ is tuple: # special case for tuple enums
206 args = (args, ) # wrap it one more time
207 if not enum_class._use_args_:
208 enum_member = enum_class._new_member_(enum_class)
209 if not hasattr(enum_member, '_value_'):
210 enum_member._value_ = value
211 else:
212 enum_member = enum_class._new_member_(enum_class, *args)
213 if not hasattr(enum_member, '_value_'):
214 if enum_class._member_type_ is object:
215 enum_member._value_ = value
216 else:
Ethan Furman7aaeb2a2021-01-25 14:26:19 -0800217 try:
218 enum_member._value_ = enum_class._member_type_(*args)
219 except Exception as exc:
220 raise TypeError(
221 '_value_ not set in __new__, unable to create it'
222 ) from None
Ethan Furmanc314e602021-01-12 23:47:57 -0800223 value = enum_member._value_
224 enum_member._name_ = member_name
225 enum_member.__objclass__ = enum_class
226 enum_member.__init__(*args)
Ethan Furman7aaeb2a2021-01-25 14:26:19 -0800227 enum_member._sort_order_ = len(enum_class._member_names_)
Ethan Furmanc314e602021-01-12 23:47:57 -0800228 # If another member with the same value was already defined, the
229 # new member becomes an alias to the existing one.
230 for name, canonical_member in enum_class._member_map_.items():
231 if canonical_member._value_ == enum_member._value_:
232 enum_member = canonical_member
233 break
234 else:
Ethan Furman7aaeb2a2021-01-25 14:26:19 -0800235 # this could still be an alias if the value is multi-bit and the
236 # class is a flag class
237 if (
238 Flag is None
239 or not issubclass(enum_class, Flag)
240 ):
241 # no other instances found, record this member in _member_names_
242 enum_class._member_names_.append(member_name)
243 elif (
244 Flag is not None
245 and issubclass(enum_class, Flag)
246 and _is_single_bit(value)
247 ):
248 # no other instances found, record this member in _member_names_
249 enum_class._member_names_.append(member_name)
Ethan Furmanc314e602021-01-12 23:47:57 -0800250 # get redirect in place before adding to _member_map_
251 # but check for other instances in parent classes first
252 need_override = False
253 descriptor = None
254 for base in enum_class.__mro__[1:]:
255 descriptor = base.__dict__.get(member_name)
256 if descriptor is not None:
257 if isinstance(descriptor, (property, DynamicClassAttribute)):
258 break
259 else:
260 need_override = True
261 # keep looking for an enum.property
262 if descriptor and not need_override:
263 # previous enum.property found, no further action needed
264 pass
265 else:
266 redirect = property()
267 redirect.__set_name__(enum_class, member_name)
268 if descriptor and need_override:
269 # previous enum.property found, but some other inherited attribute
270 # is in the way; copy fget, fset, fdel to this one
271 redirect.fget = descriptor.fget
272 redirect.fset = descriptor.fset
273 redirect.fdel = descriptor.fdel
274 setattr(enum_class, member_name, redirect)
275 # now add to _member_map_ (even aliases)
276 enum_class._member_map_[member_name] = enum_member
277 try:
278 # This may fail if value is not hashable. We can't add the value
279 # to the map, and by-value lookups for this value will be
280 # linear.
Ethan Furman7aaeb2a2021-01-25 14:26:19 -0800281 enum_class._value2member_map_.setdefault(value, enum_member)
Ethan Furmanc314e602021-01-12 23:47:57 -0800282 except TypeError:
283 pass
284
Ethan Furman101e0742013-09-15 12:34:36 -0700285
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700286class _EnumDict(dict):
Ethan Furman6d3dfee2020-12-08 12:26:56 -0800287 """
288 Track enum member order and ensure member names are not reused.
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700289
Ethan Furmanb7751062021-03-30 21:17:26 -0700290 EnumType will use the names found in self._member_names as the
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700291 enumeration member names.
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700292 """
293 def __init__(self):
294 super().__init__()
295 self._member_names = []
Ethan Furmanc16595e2016-09-10 23:36:59 -0700296 self._last_values = []
Ethan Furmana4b1bb42018-01-22 07:56:37 -0800297 self._ignore = []
Ethan Onstottd9a43e22020-04-28 13:20:55 -0400298 self._auto_called = False
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700299
300 def __setitem__(self, key, value):
Ethan Furman6d3dfee2020-12-08 12:26:56 -0800301 """
302 Changes anything not dundered or not a descriptor.
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700303
304 If an enum member name is used twice, an error is raised; duplicate
305 values are not checked for.
306
307 Single underscore (sunder) names are reserved.
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700308 """
Ethan Furman7cf0aad2020-12-09 17:12:11 -0800309 if _is_private(self._cls_name, key):
310 # do nothing, name will be a normal attribute
311 pass
312 elif _is_sunder(key):
Ethan Furmanee47e5c2016-08-31 00:12:15 -0700313 if key not in (
Ethan Furman0dca5eb2021-04-15 06:49:54 -0700314 '_order_',
Ethan Furmana4b1bb42018-01-22 07:56:37 -0800315 '_generate_next_value_', '_missing_', '_ignore_',
Ethan Furman7aaeb2a2021-01-25 14:26:19 -0800316 '_iter_member_', '_iter_member_by_value_', '_iter_member_by_def_',
Ethan Furmanee47e5c2016-08-31 00:12:15 -0700317 ):
Ethan Furman6d3dfee2020-12-08 12:26:56 -0800318 raise ValueError(
319 '_sunder_ names, such as %r, are reserved for future Enum use'
320 % (key, )
321 )
Ethan Furmanc16595e2016-09-10 23:36:59 -0700322 if key == '_generate_next_value_':
Ethan Onstottd9a43e22020-04-28 13:20:55 -0400323 # check if members already defined as auto()
324 if self._auto_called:
325 raise TypeError("_generate_next_value_ must be defined before members")
Ethan Furmanb7751062021-03-30 21:17:26 -0700326 _gnv = value.__func__ if isinstance(value, staticmethod) else value
327 setattr(self, '_generate_next_value', _gnv)
Ethan Furmana4b1bb42018-01-22 07:56:37 -0800328 elif key == '_ignore_':
329 if isinstance(value, str):
330 value = value.replace(',',' ').split()
331 else:
332 value = list(value)
333 self._ignore = value
334 already = set(value) & set(self._member_names)
335 if already:
Ethan Furman6d3dfee2020-12-08 12:26:56 -0800336 raise ValueError(
337 '_ignore_ cannot specify already set names: %r'
338 % (already, )
339 )
Ethan Furman101e0742013-09-15 12:34:36 -0700340 elif _is_dunder(key):
Ethan Furmane8e61272016-08-20 07:19:31 -0700341 if key == '__order__':
342 key = '_order_'
Ethan Furman101e0742013-09-15 12:34:36 -0700343 elif key in self._member_names:
344 # descriptor overwriting an enum?
Ethan Furmana6582872020-12-10 13:07:00 -0800345 raise TypeError('%r already defined as: %r' % (key, self[key]))
Ethan Furmana4b1bb42018-01-22 07:56:37 -0800346 elif key in self._ignore:
347 pass
Ethan Furman101e0742013-09-15 12:34:36 -0700348 elif not _is_descriptor(value):
349 if key in self:
350 # enum overwriting a descriptor?
Ethan Furmanee47e5c2016-08-31 00:12:15 -0700351 raise TypeError('%r already defined as: %r' % (key, self[key]))
Ethan Furmanc16595e2016-09-10 23:36:59 -0700352 if isinstance(value, auto):
Ethan Furman3515dcc2016-09-18 13:15:41 -0700353 if value.value == _auto_null:
Ethan Furman6d3dfee2020-12-08 12:26:56 -0800354 value.value = self._generate_next_value(
Ethan Furman7aaeb2a2021-01-25 14:26:19 -0800355 key, 1, len(self._member_names), self._last_values[:],
Ethan Furman6d3dfee2020-12-08 12:26:56 -0800356 )
Ethan Furmanfc23a942020-09-16 12:37:54 -0700357 self._auto_called = True
Ethan Furman3515dcc2016-09-18 13:15:41 -0700358 value = value.value
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700359 self._member_names.append(key)
Ethan Furmanc16595e2016-09-10 23:36:59 -0700360 self._last_values.append(value)
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700361 super().__setitem__(key, value)
362
Ethan Furmana6582872020-12-10 13:07:00 -0800363 def update(self, members, **more_members):
364 try:
365 for name in members.keys():
366 self[name] = members[name]
367 except AttributeError:
368 for name, value in members:
369 self[name] = value
370 for name, value in more_members.items():
371 self[name] = value
372
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700373
Ethan Furmanb7751062021-03-30 21:17:26 -0700374class EnumType(type):
Ethan Furman6d3dfee2020-12-08 12:26:56 -0800375 """
376 Metaclass for Enum
377 """
Ethan Furman7aaeb2a2021-01-25 14:26:19 -0800378
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700379 @classmethod
Ethan Furman6ec0ade2020-12-24 10:05:02 -0800380 def __prepare__(metacls, cls, bases, **kwds):
Ethan Furman3064dbf2020-09-16 07:11:57 -0700381 # check that previous enum members do not exist
382 metacls._check_for_existing_members(cls, bases)
Ethan Furmanee47e5c2016-08-31 00:12:15 -0700383 # create the namespace dict
384 enum_dict = _EnumDict()
Ethan Furman7cf0aad2020-12-09 17:12:11 -0800385 enum_dict._cls_name = cls
Ethan Furmanee47e5c2016-08-31 00:12:15 -0700386 # inherit previous flags and _generate_next_value_ function
Ethan Furman3064dbf2020-09-16 07:11:57 -0700387 member_type, first_enum = metacls._get_mixins_(cls, bases)
Ethan Furmanee47e5c2016-08-31 00:12:15 -0700388 if first_enum is not None:
Ethan Furman6d3dfee2020-12-08 12:26:56 -0800389 enum_dict['_generate_next_value_'] = getattr(
390 first_enum, '_generate_next_value_', None,
391 )
Ethan Furmanee47e5c2016-08-31 00:12:15 -0700392 return enum_dict
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700393
Ethan Furman7aaeb2a2021-01-25 14:26:19 -0800394 def __new__(metacls, cls, bases, classdict, boundary=None, **kwds):
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700395 # an Enum class is final once enumeration items have been defined; it
396 # cannot be mixed with other types (int, float, etc.) if it has an
397 # inherited __new__ unless a new __new__ is defined (or the resulting
398 # class will fail).
Ethan Furmana4b1bb42018-01-22 07:56:37 -0800399 #
400 # remove any keys listed in _ignore_
401 classdict.setdefault('_ignore_', []).append('_ignore_')
402 ignore = classdict['_ignore_']
403 for key in ignore:
404 classdict.pop(key, None)
Ethan Furmanc314e602021-01-12 23:47:57 -0800405 #
406 # grab member names
407 member_names = classdict._member_names
408 #
409 # check for illegal enum names (any others?)
410 invalid_names = set(member_names) & {'mro', ''}
411 if invalid_names:
412 raise ValueError('Invalid enum member name: {0}'.format(
413 ','.join(invalid_names)))
414 #
415 # adjust the sunders
416 _order_ = classdict.pop('_order_', None)
417 # convert to normal dict
418 classdict = dict(classdict.items())
419 #
420 # data type of member and the controlling Enum class
Ethan Furman3064dbf2020-09-16 07:11:57 -0700421 member_type, first_enum = metacls._get_mixins_(cls, bases)
Ethan Furmanc2667362020-12-07 00:17:31 -0800422 __new__, save_new, use_args = metacls._find_new_(
423 classdict, member_type, first_enum,
424 )
Ethan Furmanc314e602021-01-12 23:47:57 -0800425 classdict['_new_member_'] = __new__
426 classdict['_use_args_'] = use_args
427 #
428 # convert future enum members into temporary _proto_members
Ethan Furman7aaeb2a2021-01-25 14:26:19 -0800429 # and record integer values in case this will be a Flag
430 flag_mask = 0
Ethan Furmanc314e602021-01-12 23:47:57 -0800431 for name in member_names:
Ethan Furman7aaeb2a2021-01-25 14:26:19 -0800432 value = classdict[name]
433 if isinstance(value, int):
434 flag_mask |= value
435 classdict[name] = _proto_member(value)
Ethan Furmanc314e602021-01-12 23:47:57 -0800436 #
Ethan Furman7aaeb2a2021-01-25 14:26:19 -0800437 # house-keeping structures
Ethan Furmanc314e602021-01-12 23:47:57 -0800438 classdict['_member_names_'] = []
439 classdict['_member_map_'] = {}
440 classdict['_value2member_map_'] = {}
441 classdict['_member_type_'] = member_type
442 #
Ethan Furman7aaeb2a2021-01-25 14:26:19 -0800443 # Flag structures (will be removed if final class is not a Flag
444 classdict['_boundary_'] = (
445 boundary
446 or getattr(first_enum, '_boundary_', None)
447 )
448 classdict['_flag_mask_'] = flag_mask
449 classdict['_all_bits_'] = 2 ** ((flag_mask).bit_length()) - 1
450 classdict['_inverted_'] = None
451 #
Ethan Furman2da95042014-03-03 12:42:52 -0800452 # If a custom type is mixed into the Enum, and it does not know how
453 # to pickle itself, pickle.dumps will succeed but pickle.loads will
454 # fail. Rather than have the error show up later and possibly far
455 # from the source, sabotage the pickle protocol for this class so
456 # that pickle.dumps also fails.
457 #
458 # However, if the new class implements its own __reduce_ex__, do not
459 # sabotage -- it's on them to make sure it works correctly. We use
460 # __reduce_ex__ instead of any of the others as it is preferred by
461 # pickle over __reduce__, and it handles all pickle protocols.
462 if '__reduce_ex__' not in classdict:
Ethan Furmandc870522014-02-18 12:37:12 -0800463 if member_type is not object:
464 methods = ('__getnewargs_ex__', '__getnewargs__',
465 '__reduce_ex__', '__reduce__')
Ethan Furman2da95042014-03-03 12:42:52 -0800466 if not any(m in member_type.__dict__ for m in methods):
Ethan Furmanc314e602021-01-12 23:47:57 -0800467 _make_class_unpicklable(classdict)
468 #
469 # create a default docstring if one has not been provided
470 if '__doc__' not in classdict:
471 classdict['__doc__'] = 'An enumeration.'
472 try:
473 exc = None
474 enum_class = super().__new__(metacls, cls, bases, classdict, **kwds)
475 except RuntimeError as e:
476 # any exceptions raised by member.__new__ will get converted to a
477 # RuntimeError, so get that original exception back and raise it instead
478 exc = e.__cause__ or e
479 if exc is not None:
480 raise exc
481 #
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700482 # double check that repr and friends are not the mixin's or various
483 # things break (such as pickle)
Ethan Furman22415ad2020-09-15 16:28:25 -0700484 # however, if the method is defined in the Enum itself, don't replace
485 # it
Ethan Furmandc870522014-02-18 12:37:12 -0800486 for name in ('__repr__', '__str__', '__format__', '__reduce_ex__'):
Ethan Furman22415ad2020-09-15 16:28:25 -0700487 if name in classdict:
488 continue
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700489 class_method = getattr(enum_class, name)
490 obj_method = getattr(member_type, name, None)
491 enum_method = getattr(first_enum, name, None)
492 if obj_method is not None and obj_method is class_method:
493 setattr(enum_class, name, enum_method)
Ethan Furmanc314e602021-01-12 23:47:57 -0800494 #
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700495 # replace any other __new__ with our own (as long as Enum is not None,
496 # anyway) -- again, this is to support pickle
497 if Enum is not None:
498 # if the user defined their own __new__, save it before it gets
499 # clobbered in case they subclass later
500 if save_new:
501 enum_class.__new_member__ = __new__
502 enum_class.__new__ = Enum.__new__
Ethan Furmanc314e602021-01-12 23:47:57 -0800503 #
Ethan Furmane8e61272016-08-20 07:19:31 -0700504 # py3 support for definition order (helps keep py2/py3 code in sync)
Ethan Furman7aaeb2a2021-01-25 14:26:19 -0800505 #
506 # _order_ checking is spread out into three/four steps
507 # - if enum_class is a Flag:
508 # - remove any non-single-bit flags from _order_
509 # - remove any aliases from _order_
510 # - check that _order_ and _member_names_ match
511 #
512 # step 1: ensure we have a list
Ethan Furmane8e61272016-08-20 07:19:31 -0700513 if _order_ is not None:
514 if isinstance(_order_, str):
515 _order_ = _order_.replace(',', ' ').split()
Ethan Furman7aaeb2a2021-01-25 14:26:19 -0800516 #
517 # remove Flag structures if final class is not a Flag
518 if (
519 Flag is None and cls != 'Flag'
520 or Flag is not None and not issubclass(enum_class, Flag)
521 ):
522 delattr(enum_class, '_boundary_')
523 delattr(enum_class, '_flag_mask_')
524 delattr(enum_class, '_all_bits_')
525 delattr(enum_class, '_inverted_')
526 elif Flag is not None and issubclass(enum_class, Flag):
527 # ensure _all_bits_ is correct and there are no missing flags
528 single_bit_total = 0
529 multi_bit_total = 0
530 for flag in enum_class._member_map_.values():
531 flag_value = flag._value_
532 if _is_single_bit(flag_value):
533 single_bit_total |= flag_value
534 else:
535 # multi-bit flags are considered aliases
536 multi_bit_total |= flag_value
537 if enum_class._boundary_ is not KEEP:
538 missed = list(_iter_bits_lsb(multi_bit_total & ~single_bit_total))
539 if missed:
540 raise TypeError(
541 'invalid Flag %r -- missing values: %s'
542 % (cls, ', '.join((str(i) for i in missed)))
543 )
544 enum_class._flag_mask_ = single_bit_total
545 #
546 # set correct __iter__
547 member_list = [m._value_ for m in enum_class]
548 if member_list != sorted(member_list):
549 enum_class._iter_member_ = enum_class._iter_member_by_def_
550 if _order_:
551 # _order_ step 2: remove any items from _order_ that are not single-bit
552 _order_ = [
553 o
554 for o in _order_
555 if o not in enum_class._member_map_ or _is_single_bit(enum_class[o]._value_)
556 ]
557 #
558 if _order_:
559 # _order_ step 3: remove aliases from _order_
560 _order_ = [
561 o
562 for o in _order_
563 if (
564 o not in enum_class._member_map_
565 or
566 (o in enum_class._member_map_ and o in enum_class._member_names_)
567 )]
568 # _order_ step 4: verify that _order_ and _member_names_ match
Ethan Furmane8e61272016-08-20 07:19:31 -0700569 if _order_ != enum_class._member_names_:
Ethan Furman7aaeb2a2021-01-25 14:26:19 -0800570 raise TypeError(
571 'member order does not match _order_:\n%r\n%r'
572 % (enum_class._member_names_, _order_)
573 )
Ethan Furmanc314e602021-01-12 23:47:57 -0800574 #
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700575 return enum_class
576
Ethan Furman5de67b12016-04-13 23:52:09 -0700577 def __bool__(self):
578 """
579 classes/types should always be True.
580 """
581 return True
582
Ethan Furman7aaeb2a2021-01-25 14:26:19 -0800583 def __call__(cls, value, names=None, *, module=None, qualname=None, type=None, start=1, boundary=None):
Ethan Furman6d3dfee2020-12-08 12:26:56 -0800584 """
585 Either returns an existing member, or creates a new enum class.
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700586
587 This method is used both when an enum class is given a value to match
588 to an enumeration member (i.e. Color(3)) and for the functional API
Ethan Furman23bb6f42016-11-21 09:22:05 -0800589 (i.e. Color = Enum('Color', names='RED GREEN BLUE')).
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700590
Ethan Furman2da95042014-03-03 12:42:52 -0800591 When used for the functional API:
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700592
Ethan Furman2da95042014-03-03 12:42:52 -0800593 `value` will be the name of the new class.
594
595 `names` should be either a string of white-space/comma delimited names
Ethan Furmand9925a12014-09-16 20:35:55 -0700596 (values will start at `start`), or an iterator/mapping of name, value pairs.
Ethan Furman2da95042014-03-03 12:42:52 -0800597
598 `module` should be set to the module this class is being created in;
599 if it is not set, an attempt to find that module will be made, but if
600 it fails the class will not be picklable.
601
602 `qualname` should be set to the actual location this class can be found
603 at in its module; by default it is set to the global scope. If this is
604 not correct, unpickling will fail in some circumstances.
605
606 `type`, if set, will be mixed in as the first base class.
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700607 """
608 if names is None: # simple value lookup
609 return cls.__new__(cls, value)
610 # otherwise, functional API: we're creating a new Enum type
Ethan Furman6d3dfee2020-12-08 12:26:56 -0800611 return cls._create_(
612 value,
613 names,
614 module=module,
615 qualname=qualname,
616 type=type,
617 start=start,
Ethan Furman7aaeb2a2021-01-25 14:26:19 -0800618 boundary=boundary,
Ethan Furman6d3dfee2020-12-08 12:26:56 -0800619 )
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700620
621 def __contains__(cls, member):
Rahul Jha94306522018-09-10 23:51:04 +0530622 if not isinstance(member, Enum):
623 raise TypeError(
624 "unsupported operand type(s) for 'in': '%s' and '%s'" % (
625 type(member).__qualname__, cls.__class__.__qualname__))
Ethan Furman0081f232014-09-16 17:31:23 -0700626 return isinstance(member, cls) and member._name_ in cls._member_map_
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700627
Ethan Furman64a99722013-09-22 16:18:19 -0700628 def __delattr__(cls, attr):
629 # nicer error message when someone tries to delete an attribute
630 # (see issue19025).
631 if attr in cls._member_map_:
Ethan Furman6d3dfee2020-12-08 12:26:56 -0800632 raise AttributeError("%s: cannot delete Enum member %r." % (cls.__name__, attr))
Ethan Furman64a99722013-09-22 16:18:19 -0700633 super().__delattr__(attr)
634
Ethan Furman388a3922013-08-12 06:51:41 -0700635 def __dir__(self):
Ethan Furman6d3dfee2020-12-08 12:26:56 -0800636 return (
637 ['__class__', '__doc__', '__members__', '__module__']
638 + self._member_names_
639 )
Ethan Furman388a3922013-08-12 06:51:41 -0700640
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700641 def __getattr__(cls, name):
Ethan Furman6d3dfee2020-12-08 12:26:56 -0800642 """
643 Return the enum member matching `name`
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700644
645 We use __getattr__ instead of descriptors or inserting into the enum
646 class' __dict__ in order to support `name` and `value` being both
647 properties for enum members (which live in the class' __dict__) and
648 enum members themselves.
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700649 """
650 if _is_dunder(name):
651 raise AttributeError(name)
652 try:
Ethan Furman520ad572013-07-19 19:47:21 -0700653 return cls._member_map_[name]
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700654 except KeyError:
655 raise AttributeError(name) from None
656
657 def __getitem__(cls, name):
Ethan Furman520ad572013-07-19 19:47:21 -0700658 return cls._member_map_[name]
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700659
660 def __iter__(cls):
Ethan Furman6d3dfee2020-12-08 12:26:56 -0800661 """
662 Returns members in definition order.
663 """
Ethan Furman520ad572013-07-19 19:47:21 -0700664 return (cls._member_map_[name] for name in cls._member_names_)
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700665
666 def __len__(cls):
Ethan Furman520ad572013-07-19 19:47:21 -0700667 return len(cls._member_names_)
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700668
Ethan Furmanc314e602021-01-12 23:47:57 -0800669 @_bltin_property
Ethan Furman2131a4a2013-09-14 18:11:24 -0700670 def __members__(cls):
Ethan Furman6d3dfee2020-12-08 12:26:56 -0800671 """
672 Returns a mapping of member name->value.
Ethan Furman2131a4a2013-09-14 18:11:24 -0700673
674 This mapping lists all enum members, including aliases. Note that this
675 is a read-only view of the internal mapping.
Ethan Furman2131a4a2013-09-14 18:11:24 -0700676 """
677 return MappingProxyType(cls._member_map_)
678
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700679 def __repr__(cls):
680 return "<enum %r>" % cls.__name__
681
Ethan Furman2131a4a2013-09-14 18:11:24 -0700682 def __reversed__(cls):
Ethan Furman6d3dfee2020-12-08 12:26:56 -0800683 """
684 Returns members in reverse definition order.
685 """
Ethan Furman2131a4a2013-09-14 18:11:24 -0700686 return (cls._member_map_[name] for name in reversed(cls._member_names_))
687
Ethan Furmanf203f2d2013-09-06 07:16:48 -0700688 def __setattr__(cls, name, value):
Ethan Furman6d3dfee2020-12-08 12:26:56 -0800689 """
690 Block attempts to reassign Enum members.
Ethan Furmanf203f2d2013-09-06 07:16:48 -0700691
692 A simple assignment to the class namespace only changes one of the
693 several possible ways to get an Enum member from the Enum class,
694 resulting in an inconsistent Enumeration.
Ethan Furmanf203f2d2013-09-06 07:16:48 -0700695 """
696 member_map = cls.__dict__.get('_member_map_', {})
697 if name in member_map:
698 raise AttributeError('Cannot reassign members.')
699 super().__setattr__(name, value)
700
Ethan Furman7aaeb2a2021-01-25 14:26:19 -0800701 def _create_(cls, class_name, names, *, module=None, qualname=None, type=None, start=1, boundary=None):
Ethan Furman6d3dfee2020-12-08 12:26:56 -0800702 """
703 Convenience method to create a new Enum class.
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700704
705 `names` can be:
706
707 * A string containing member names, separated either with spaces or
Ethan Furmand9925a12014-09-16 20:35:55 -0700708 commas. Values are incremented by 1 from `start`.
709 * An iterable of member names. Values are incremented by 1 from `start`.
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700710 * An iterable of (member name, value) pairs.
Ethan Furmand9925a12014-09-16 20:35:55 -0700711 * A mapping of member name -> value pairs.
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700712 """
713 metacls = cls.__class__
714 bases = (cls, ) if type is None else (type, cls)
Ethan Furman3064dbf2020-09-16 07:11:57 -0700715 _, first_enum = cls._get_mixins_(cls, bases)
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700716 classdict = metacls.__prepare__(class_name, bases)
717
718 # special processing needed for names?
719 if isinstance(names, str):
720 names = names.replace(',', ' ').split()
Dong-hee Nadcc8ce42017-06-22 01:52:32 +0900721 if isinstance(names, (tuple, list)) and names and isinstance(names[0], str):
Ethan Furmanee47e5c2016-08-31 00:12:15 -0700722 original_names, names = names, []
Ethan Furmanc16595e2016-09-10 23:36:59 -0700723 last_values = []
Ethan Furmanee47e5c2016-08-31 00:12:15 -0700724 for count, name in enumerate(original_names):
Ethan Furmanc16595e2016-09-10 23:36:59 -0700725 value = first_enum._generate_next_value_(name, start, count, last_values[:])
726 last_values.append(value)
727 names.append((name, value))
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700728
729 # Here, names is either an iterable of (name, value) or a mapping.
730 for item in names:
731 if isinstance(item, str):
732 member_name, member_value = item, names[item]
733 else:
734 member_name, member_value = item
735 classdict[member_name] = member_value
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700736
737 # TODO: replace the frame hack if a blessed way to know the calling
738 # module is ever developed
739 if module is None:
740 try:
741 module = sys._getframe(2).f_globals['__name__']
Pablo Galindo293dd232019-11-19 21:34:03 +0000742 except (AttributeError, ValueError, KeyError):
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700743 pass
744 if module is None:
Ethan Furmanc314e602021-01-12 23:47:57 -0800745 _make_class_unpicklable(classdict)
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700746 else:
Ethan Furmanc314e602021-01-12 23:47:57 -0800747 classdict['__module__'] = module
Ethan Furmanca1b7942014-02-08 11:36:27 -0800748 if qualname is not None:
Ethan Furmanc314e602021-01-12 23:47:57 -0800749 classdict['__qualname__'] = qualname
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700750
Ethan Furman7aaeb2a2021-01-25 14:26:19 -0800751 return metacls.__new__(metacls, class_name, bases, classdict, boundary=boundary)
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700752
Ethan Furman7aaeb2a2021-01-25 14:26:19 -0800753 def _convert_(cls, name, module, filter, source=None, boundary=None):
orlnub1230fb9fad2018-09-12 20:28:53 +0300754 """
755 Create a new Enum subclass that replaces a collection of global constants
756 """
757 # convert all constants from source (or module) that pass filter() to
758 # a new Enum called name, and export the enum and its members back to
759 # module;
760 # also, replace the __reduce_ex__ method so unpickling works in
761 # previous Python versions
Ethan Furmanb7751062021-03-30 21:17:26 -0700762 module_globals = sys.modules[module].__dict__
orlnub1230fb9fad2018-09-12 20:28:53 +0300763 if source:
Ethan Furmanb7751062021-03-30 21:17:26 -0700764 source = source.__dict__
orlnub1230fb9fad2018-09-12 20:28:53 +0300765 else:
766 source = module_globals
767 # _value2member_map_ is populated in the same order every time
768 # for a consistent reverse mapping of number to name when there
769 # are multiple names for the same number.
770 members = [
771 (name, value)
772 for name, value in source.items()
773 if filter(name)]
774 try:
775 # sort by value
776 members.sort(key=lambda t: (t[1], t[0]))
777 except TypeError:
778 # unless some values aren't comparable, in which case sort by name
779 members.sort(key=lambda t: t[0])
Ethan Furman7aaeb2a2021-01-25 14:26:19 -0800780 cls = cls(name, members, module=module, boundary=boundary or KEEP)
orlnub1230fb9fad2018-09-12 20:28:53 +0300781 cls.__reduce_ex__ = _reduce_ex_by_name
Ethan Furmanb7751062021-03-30 21:17:26 -0700782 global_enum(cls)
orlnub1230fb9fad2018-09-12 20:28:53 +0300783 module_globals[name] = cls
784 return cls
785
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700786 @staticmethod
Ethan Furman3064dbf2020-09-16 07:11:57 -0700787 def _check_for_existing_members(class_name, bases):
788 for chain in bases:
789 for base in chain.__mro__:
790 if issubclass(base, Enum) and base._member_names_:
Ethan Furman6d3dfee2020-12-08 12:26:56 -0800791 raise TypeError(
792 "%s: cannot extend enumeration %r"
793 % (class_name, base.__name__)
794 )
Ethan Furman3064dbf2020-09-16 07:11:57 -0700795
796 @staticmethod
797 def _get_mixins_(class_name, bases):
Ethan Furman6d3dfee2020-12-08 12:26:56 -0800798 """
799 Returns the type for creating enum members, and the first inherited
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700800 enum class.
801
802 bases: the tuple of bases that was given to __new__
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700803 """
804 if not bases:
805 return object, Enum
806
Ethan Furman5bdab642018-09-21 19:03:09 -0700807 def _find_data_type(bases):
Ethan Furmanbff01f32020-09-15 15:56:26 -0700808 data_types = []
Ethan Furman5bdab642018-09-21 19:03:09 -0700809 for chain in bases:
Ethan Furmanbff01f32020-09-15 15:56:26 -0700810 candidate = None
Ethan Furman5bdab642018-09-21 19:03:09 -0700811 for base in chain.__mro__:
812 if base is object:
813 continue
Ethan Furmanc2667362020-12-07 00:17:31 -0800814 elif issubclass(base, Enum):
815 if base._member_type_ is not object:
816 data_types.append(base._member_type_)
817 break
Ethan Furman5bdab642018-09-21 19:03:09 -0700818 elif '__new__' in base.__dict__:
Ethan Furmancd453852018-10-05 23:29:36 -0700819 if issubclass(base, Enum):
Ethan Furman5bdab642018-09-21 19:03:09 -0700820 continue
Ethan Furmanbff01f32020-09-15 15:56:26 -0700821 data_types.append(candidate or base)
822 break
Ethan Furmanc2667362020-12-07 00:17:31 -0800823 else:
Ethan Furmanbff01f32020-09-15 15:56:26 -0700824 candidate = base
825 if len(data_types) > 1:
Ethan Furman3064dbf2020-09-16 07:11:57 -0700826 raise TypeError('%r: too many data types: %r' % (class_name, data_types))
Ethan Furmanbff01f32020-09-15 15:56:26 -0700827 elif data_types:
828 return data_types[0]
829 else:
830 return None
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700831
Ethan Furman5bdab642018-09-21 19:03:09 -0700832 # ensure final parent class is an Enum derivative, find any concrete
833 # data type, and check that Enum has no members
834 first_enum = bases[-1]
835 if not issubclass(first_enum, Enum):
836 raise TypeError("new enumerations should be created as "
837 "`EnumName([mixin_type, ...] [data_type,] enum_type)`")
838 member_type = _find_data_type(bases) or object
839 if first_enum._member_names_:
840 raise TypeError("Cannot extend enumerations")
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700841 return member_type, first_enum
842
843 @staticmethod
844 def _find_new_(classdict, member_type, first_enum):
Ethan Furman6d3dfee2020-12-08 12:26:56 -0800845 """
846 Returns the __new__ to be used for creating the enum members.
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700847
848 classdict: the class dictionary given to __new__
849 member_type: the data type whose __new__ will be used by default
850 first_enum: enumeration to check for an overriding __new__
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700851 """
852 # now find the correct __new__, checking to see of one was defined
853 # by the user; also check earlier enum classes in case a __new__ was
854 # saved as __new_member__
855 __new__ = classdict.get('__new__', None)
856
857 # should __new__ be saved as __new_member__ later?
858 save_new = __new__ is not None
859
860 if __new__ is None:
861 # check all possibles for __new_member__ before falling back to
862 # __new__
863 for method in ('__new_member__', '__new__'):
864 for possible in (member_type, first_enum):
865 target = getattr(possible, method, None)
866 if target not in {
867 None,
868 None.__new__,
869 object.__new__,
870 Enum.__new__,
871 }:
872 __new__ = target
873 break
874 if __new__ is not None:
875 break
876 else:
877 __new__ = object.__new__
878
879 # if a non-object.__new__ is used then whatever value/tuple was
880 # assigned to the enum member name will be passed to __new__ and to the
881 # new enum member's __init__
882 if __new__ is object.__new__:
883 use_args = False
884 else:
885 use_args = True
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700886 return __new__, save_new, use_args
Ethan Furmanb7751062021-03-30 21:17:26 -0700887EnumMeta = EnumType
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700888
889
Ethan Furmanb7751062021-03-30 21:17:26 -0700890class Enum(metaclass=EnumType):
Ethan Furman6d3dfee2020-12-08 12:26:56 -0800891 """
892 Generic enumeration.
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700893
894 Derive from this class to define new enumerations.
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700895 """
Ethan Furman7aaeb2a2021-01-25 14:26:19 -0800896
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700897 def __new__(cls, value):
898 # all enum instances are actually created during class construction
899 # without calling this method; this method is called by the metaclass'
900 # __call__ (i.e. Color(3) ), and by pickle
901 if type(value) is cls:
Ethan Furman23bb6f42016-11-21 09:22:05 -0800902 # For lookups like Color(Color.RED)
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700903 return value
904 # by-value search for a matching enum member
905 # see if it's in the reverse mapping (for hashable values)
Ethan Furman2aa27322013-07-19 19:35:56 -0700906 try:
Andrew Svetlov34ae04f2018-12-26 20:45:33 +0200907 return cls._value2member_map_[value]
908 except KeyError:
909 # Not found, no need to do long O(n) search
910 pass
Ethan Furman2aa27322013-07-19 19:35:56 -0700911 except TypeError:
912 # not there, now do long search -- O(n) behavior
Ethan Furman520ad572013-07-19 19:47:21 -0700913 for member in cls._member_map_.values():
Ethan Furman0081f232014-09-16 17:31:23 -0700914 if member._value_ == value:
Ethan Furman2aa27322013-07-19 19:35:56 -0700915 return member
Ethan Furmanee47e5c2016-08-31 00:12:15 -0700916 # still not found -- try _missing_ hook
Ethan Furman019f0a02018-09-12 11:43:34 -0700917 try:
918 exc = None
919 result = cls._missing_(value)
920 except Exception as e:
921 exc = e
922 result = None
Ethan Furman8c14f5a2021-04-12 08:51:20 -0700923 try:
924 if isinstance(result, cls):
925 return result
926 elif (
927 Flag is not None and issubclass(cls, Flag)
928 and cls._boundary_ is EJECT and isinstance(result, int)
929 ):
930 return result
931 else:
932 ve_exc = ValueError("%r is not a valid %s" % (value, cls.__qualname__))
933 if result is None and exc is None:
934 raise ve_exc
935 elif exc is None:
936 exc = TypeError(
937 'error in %s._missing_: returned %r instead of None or a valid member'
938 % (cls.__name__, result)
939 )
940 if not isinstance(exc, ValueError):
941 exc.__context__ = ve_exc
942 raise exc
943 finally:
944 # ensure all variables that could hold an exception are destroyed
945 exc = None
946 ve_exc = None
Ethan Furmanee47e5c2016-08-31 00:12:15 -0700947
Ethan Furmanc16595e2016-09-10 23:36:59 -0700948 def _generate_next_value_(name, start, count, last_values):
Ethan Furman6d3dfee2020-12-08 12:26:56 -0800949 """
950 Generate the next value when not given.
951
952 name: the name of the member
953 start: the initial start value or None
954 count: the number of existing members
955 last_value: the last value assigned or None
956 """
Ethan Furmanc16595e2016-09-10 23:36:59 -0700957 for last_value in reversed(last_values):
958 try:
959 return last_value + 1
960 except TypeError:
961 pass
962 else:
Ethan Furmanee47e5c2016-08-31 00:12:15 -0700963 return start
Ethan Furmanc16595e2016-09-10 23:36:59 -0700964
Ethan Furmanee47e5c2016-08-31 00:12:15 -0700965 @classmethod
966 def _missing_(cls, value):
Ethan Furmanc95ad7a2020-09-16 10:26:50 -0700967 return None
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700968
969 def __repr__(self):
Ethan Furmanb7751062021-03-30 21:17:26 -0700970 return "%s.%s" % ( self.__class__.__name__, self._name_)
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700971
972 def __str__(self):
Ethan Furmanb7751062021-03-30 21:17:26 -0700973 return "%s" % (self._name_, )
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700974
Ethan Furman388a3922013-08-12 06:51:41 -0700975 def __dir__(self):
Ethan Furman6d3dfee2020-12-08 12:26:56 -0800976 """
977 Returns all members and all public methods
978 """
Ethan Furman0ae550b2014-10-14 08:58:32 -0700979 added_behavior = [
980 m
981 for cls in self.__class__.mro()
982 for m in cls.__dict__
Ethan Furman354ecf12015-03-11 08:43:12 -0700983 if m[0] != '_' and m not in self._member_map_
Angelin BOOZ68526fe2020-09-21 15:11:06 +0200984 ] + [m for m in self.__dict__ if m[0] != '_']
Ethan Furmanec5f8eb2014-10-21 13:40:35 -0700985 return (['__class__', '__doc__', '__module__'] + added_behavior)
Ethan Furman388a3922013-08-12 06:51:41 -0700986
Ethan Furmanec15a822013-08-31 19:17:41 -0700987 def __format__(self, format_spec):
Ethan Furman6d3dfee2020-12-08 12:26:56 -0800988 """
989 Returns format using actual value type unless __str__ has been overridden.
990 """
Ethan Furmanec15a822013-08-31 19:17:41 -0700991 # mixed-in Enums should use the mixed-in type's __format__, otherwise
992 # we can get strange results with the Enum name showing up instead of
993 # the value
994
thatneat2f19e822019-07-04 11:28:37 -0700995 # pure Enum branch, or branch with __str__ explicitly overridden
Ethan Furman37440ee2020-12-08 11:14:10 -0800996 str_overridden = type(self).__str__ not in (Enum.__str__, Flag.__str__)
thatneat2f19e822019-07-04 11:28:37 -0700997 if self._member_type_ is object or str_overridden:
Ethan Furmanec15a822013-08-31 19:17:41 -0700998 cls = str
999 val = str(self)
1000 # mix-in branch
1001 else:
1002 cls = self._member_type_
Ethan Furman0081f232014-09-16 17:31:23 -07001003 val = self._value_
Ethan Furmanec15a822013-08-31 19:17:41 -07001004 return cls.__format__(val, format_spec)
1005
Ethan Furman6b3d64a2013-06-14 16:55:46 -07001006 def __hash__(self):
Ethan Furman520ad572013-07-19 19:47:21 -07001007 return hash(self._name_)
Ethan Furman6b3d64a2013-06-14 16:55:46 -07001008
Ethan Furmanca1b7942014-02-08 11:36:27 -08001009 def __reduce_ex__(self, proto):
Ethan Furmandc870522014-02-18 12:37:12 -08001010 return self.__class__, (self._value_, )
Ethan Furmanca1b7942014-02-08 11:36:27 -08001011
Ethan Furmanc314e602021-01-12 23:47:57 -08001012 # enum.property is used to provide access to the `name` and
1013 # `value` attributes of enum members while keeping some measure of
Ethan Furman6b3d64a2013-06-14 16:55:46 -07001014 # protection from modification, while still allowing for an enumeration
1015 # to have members named `name` and `value`. This works because enumeration
Ethan Furmanc314e602021-01-12 23:47:57 -08001016 # members are not set directly on the enum class; they are kept in a
1017 # separate structure, _member_map_, which is where enum.property looks for
1018 # them
Ethan Furman6b3d64a2013-06-14 16:55:46 -07001019
Ethan Furmanc314e602021-01-12 23:47:57 -08001020 @property
Ethan Furman6b3d64a2013-06-14 16:55:46 -07001021 def name(self):
Ethan Furmanc850f342013-09-15 16:59:35 -07001022 """The name of the Enum member."""
Ethan Furman520ad572013-07-19 19:47:21 -07001023 return self._name_
Ethan Furman6b3d64a2013-06-14 16:55:46 -07001024
Ethan Furmanc314e602021-01-12 23:47:57 -08001025 @property
Ethan Furman6b3d64a2013-06-14 16:55:46 -07001026 def value(self):
Ethan Furmanc850f342013-09-15 16:59:35 -07001027 """The value of the Enum member."""
Ethan Furman520ad572013-07-19 19:47:21 -07001028 return self._value_
Ethan Furman6b3d64a2013-06-14 16:55:46 -07001029
1030
1031class IntEnum(int, Enum):
Ethan Furman0063ff42020-09-21 17:23:13 -07001032 """
1033 Enum where members are also (and must be) ints
1034 """
1035
1036
1037class StrEnum(str, Enum):
1038 """
1039 Enum where members are also (and must be) strings
1040 """
1041
1042 def __new__(cls, *values):
1043 if len(values) > 3:
1044 raise TypeError('too many arguments for str(): %r' % (values, ))
1045 if len(values) == 1:
1046 # it must be a string
1047 if not isinstance(values[0], str):
1048 raise TypeError('%r is not a string' % (values[0], ))
Ethan Furman7aaeb2a2021-01-25 14:26:19 -08001049 if len(values) >= 2:
Ethan Furman0063ff42020-09-21 17:23:13 -07001050 # check that encoding argument is a string
1051 if not isinstance(values[1], str):
1052 raise TypeError('encoding must be a string, not %r' % (values[1], ))
Ethan Furman7aaeb2a2021-01-25 14:26:19 -08001053 if len(values) == 3:
1054 # check that errors argument is a string
1055 if not isinstance(values[2], str):
1056 raise TypeError('errors must be a string, not %r' % (values[2]))
Ethan Furman0063ff42020-09-21 17:23:13 -07001057 value = str(*values)
1058 member = str.__new__(cls, value)
1059 member._value_ = value
1060 return member
Ethan Furmanf24bb352013-07-18 17:05:39 -07001061
Ethan Furmand986d162020-09-22 13:00:07 -07001062 __str__ = str.__str__
1063
Ethan Furmanefb13be2020-12-10 12:20:06 -08001064 def _generate_next_value_(name, start, count, last_values):
1065 """
1066 Return the lower-cased version of the member name.
1067 """
1068 return name.lower()
1069
Ethan Furmanf24bb352013-07-18 17:05:39 -07001070
Ethan Furman24e837f2015-03-18 17:27:57 -07001071def _reduce_ex_by_name(self, proto):
1072 return self.name
1073
Ethan Furman7aaeb2a2021-01-25 14:26:19 -08001074class FlagBoundary(StrEnum):
1075 """
1076 control how out of range values are handled
1077 "strict" -> error is raised [default for Flag]
1078 "conform" -> extra bits are discarded
1079 "eject" -> lose flag status [default for IntFlag]
1080 "keep" -> keep flag status and all bits
1081 """
1082 STRICT = auto()
1083 CONFORM = auto()
1084 EJECT = auto()
1085 KEEP = auto()
1086STRICT, CONFORM, EJECT, KEEP = FlagBoundary
1087
1088
1089class Flag(Enum, boundary=STRICT):
Ethan Furman6d3dfee2020-12-08 12:26:56 -08001090 """
1091 Support for flags
1092 """
Ethan Furmanc16595e2016-09-10 23:36:59 -07001093
1094 def _generate_next_value_(name, start, count, last_values):
Ethan Furmanee47e5c2016-08-31 00:12:15 -07001095 """
1096 Generate the next value when not given.
1097
1098 name: the name of the member
HongWeipengbb16fb22019-09-21 13:22:54 +08001099 start: the initial start value or None
Ethan Furmanee47e5c2016-08-31 00:12:15 -07001100 count: the number of existing members
1101 last_value: the last value assigned or None
1102 """
1103 if not count:
1104 return start if start is not None else 1
Ethan Furman7aaeb2a2021-01-25 14:26:19 -08001105 last_value = max(last_values)
1106 try:
1107 high_bit = _high_bit(last_value)
1108 except Exception:
1109 raise TypeError('Invalid Flag value: %r' % last_value) from None
Ethan Furmanee47e5c2016-08-31 00:12:15 -07001110 return 2 ** (high_bit+1)
1111
1112 @classmethod
Ethan Furman7aaeb2a2021-01-25 14:26:19 -08001113 def _iter_member_by_value_(cls, value):
1114 """
1115 Extract all members from the value in definition (i.e. increasing value) order.
1116 """
1117 for val in _iter_bits_lsb(value & cls._flag_mask_):
1118 yield cls._value2member_map_.get(val)
1119
1120 _iter_member_ = _iter_member_by_value_
1121
1122 @classmethod
1123 def _iter_member_by_def_(cls, value):
1124 """
1125 Extract all members from the value in definition order.
1126 """
1127 yield from sorted(
1128 cls._iter_member_by_value_(value),
1129 key=lambda m: m._sort_order_,
1130 )
1131
1132 @classmethod
Ethan Furmanee47e5c2016-08-31 00:12:15 -07001133 def _missing_(cls, value):
Ethan Furman6d3dfee2020-12-08 12:26:56 -08001134 """
Ethan Furman0dca5eb2021-04-15 06:49:54 -07001135 Create a composite member containing all canonical members present in `value`.
1136
1137 If non-member values are present, result depends on `_boundary_` setting.
Ethan Furman3515dcc2016-09-18 13:15:41 -07001138 """
Ethan Furman7aaeb2a2021-01-25 14:26:19 -08001139 if not isinstance(value, int):
1140 raise ValueError(
1141 "%r is not a valid %s" % (value, cls.__qualname__)
1142 )
1143 # check boundaries
1144 # - value must be in range (e.g. -16 <-> +15, i.e. ~15 <-> 15)
1145 # - value must not include any skipped flags (e.g. if bit 2 is not
1146 # defined, then 0d10 is invalid)
1147 flag_mask = cls._flag_mask_
1148 all_bits = cls._all_bits_
1149 neg_value = None
1150 if (
1151 not ~all_bits <= value <= all_bits
1152 or value & (all_bits ^ flag_mask)
1153 ):
1154 if cls._boundary_ is STRICT:
1155 max_bits = max(value.bit_length(), flag_mask.bit_length())
1156 raise ValueError(
1157 "%s: invalid value: %r\n given %s\n allowed %s" % (
1158 cls.__name__, value, bin(value, max_bits), bin(flag_mask, max_bits),
1159 ))
1160 elif cls._boundary_ is CONFORM:
1161 value = value & flag_mask
1162 elif cls._boundary_ is EJECT:
1163 return value
1164 elif cls._boundary_ is KEEP:
1165 if value < 0:
1166 value = (
1167 max(all_bits+1, 2**(value.bit_length()))
1168 + value
1169 )
1170 else:
1171 raise ValueError(
1172 'unknown flag boundary: %r' % (cls._boundary_, )
1173 )
1174 if value < 0:
1175 neg_value = value
1176 value = all_bits + 1 + value
1177 # get members and unknown
1178 unknown = value & ~flag_mask
1179 member_value = value & flag_mask
1180 if unknown and cls._boundary_ is not KEEP:
1181 raise ValueError(
1182 '%s(%r) --> unknown values %r [%s]'
1183 % (cls.__name__, value, unknown, bin(unknown))
1184 )
1185 # normal Flag?
1186 __new__ = getattr(cls, '__new_member__', None)
1187 if cls._member_type_ is object and not __new__:
Ethan Furman3515dcc2016-09-18 13:15:41 -07001188 # construct a singleton enum pseudo-member
Ethan Furmanee47e5c2016-08-31 00:12:15 -07001189 pseudo_member = object.__new__(cls)
Ethan Furman7aaeb2a2021-01-25 14:26:19 -08001190 else:
1191 pseudo_member = (__new__ or cls._member_type_.__new__)(cls, value)
1192 if not hasattr(pseudo_member, 'value'):
Ethan Furmanee47e5c2016-08-31 00:12:15 -07001193 pseudo_member._value_ = value
Ethan Furman7aaeb2a2021-01-25 14:26:19 -08001194 if member_value:
1195 pseudo_member._name_ = '|'.join([
1196 m._name_ for m in cls._iter_member_(member_value)
1197 ])
1198 if unknown:
1199 pseudo_member._name_ += '|0x%x' % unknown
1200 else:
1201 pseudo_member._name_ = None
1202 # use setdefault in case another thread already created a composite
1203 # with this value, but only if all members are known
1204 # note: zero is a special case -- add it
1205 if not unknown:
Ethan Furman28cf6632017-01-24 12:12:06 -08001206 pseudo_member = cls._value2member_map_.setdefault(value, pseudo_member)
Ethan Furman7aaeb2a2021-01-25 14:26:19 -08001207 if neg_value is not None:
1208 cls._value2member_map_[neg_value] = pseudo_member
Ethan Furmanee47e5c2016-08-31 00:12:15 -07001209 return pseudo_member
1210
Ethan Furmanee47e5c2016-08-31 00:12:15 -07001211 def __contains__(self, other):
Ethan Furman6d3dfee2020-12-08 12:26:56 -08001212 """
1213 Returns True if self has at least the same flags set as other.
1214 """
Ethan Furmanee47e5c2016-08-31 00:12:15 -07001215 if not isinstance(other, self.__class__):
Rahul Jha94306522018-09-10 23:51:04 +05301216 raise TypeError(
1217 "unsupported operand type(s) for 'in': '%s' and '%s'" % (
1218 type(other).__qualname__, self.__class__.__qualname__))
Ethan Furman7aaeb2a2021-01-25 14:26:19 -08001219 if other._value_ == 0 or self._value_ == 0:
1220 return False
Ethan Furmanee47e5c2016-08-31 00:12:15 -07001221 return other._value_ & self._value_ == other._value_
1222
Ethan Furman7219e272020-09-16 13:01:00 -07001223 def __iter__(self):
Ethan Furman6d3dfee2020-12-08 12:26:56 -08001224 """
Ethan Furman7aaeb2a2021-01-25 14:26:19 -08001225 Returns flags in definition order.
Ethan Furman6d3dfee2020-12-08 12:26:56 -08001226 """
Ethan Furman7aaeb2a2021-01-25 14:26:19 -08001227 yield from self._iter_member_(self._value_)
1228
1229 def __len__(self):
1230 return self._value_.bit_count()
Ethan Furman7219e272020-09-16 13:01:00 -07001231
Ethan Furmanee47e5c2016-08-31 00:12:15 -07001232 def __repr__(self):
Ethan Furmanb7751062021-03-30 21:17:26 -07001233 cls_name = self.__class__.__name__
1234 if self._name_ is None:
1235 return "0x%x" % (self._value_, )
1236 if _is_single_bit(self._value_):
1237 return '%s.%s' % (cls_name, self._name_)
1238 if self._boundary_ is not FlagBoundary.KEEP:
1239 return '%s.' % cls_name + ('|%s.' % cls_name).join(self.name.split('|'))
Ethan Furman7aaeb2a2021-01-25 14:26:19 -08001240 else:
Ethan Furmanb7751062021-03-30 21:17:26 -07001241 name = []
1242 for n in self._name_.split('|'):
1243 if n.startswith('0'):
1244 name.append(n)
1245 else:
1246 name.append('%s.%s' % (cls_name, n))
1247 return '|'.join(name)
Ethan Furmanee47e5c2016-08-31 00:12:15 -07001248
1249 def __str__(self):
1250 cls = self.__class__
Ethan Furmanb7751062021-03-30 21:17:26 -07001251 if self._name_ is None:
1252 return '%s(%x)' % (cls.__name__, self._value_)
Ethan Furmanee47e5c2016-08-31 00:12:15 -07001253 else:
Ethan Furmanb7751062021-03-30 21:17:26 -07001254 return self._name_
Ethan Furmanee47e5c2016-08-31 00:12:15 -07001255
Ethan Furman25d94bb2016-09-02 16:32:32 -07001256 def __bool__(self):
1257 return bool(self._value_)
1258
Ethan Furmanee47e5c2016-08-31 00:12:15 -07001259 def __or__(self, other):
1260 if not isinstance(other, self.__class__):
1261 return NotImplemented
1262 return self.__class__(self._value_ | other._value_)
1263
1264 def __and__(self, other):
1265 if not isinstance(other, self.__class__):
1266 return NotImplemented
1267 return self.__class__(self._value_ & other._value_)
1268
1269 def __xor__(self, other):
1270 if not isinstance(other, self.__class__):
1271 return NotImplemented
1272 return self.__class__(self._value_ ^ other._value_)
1273
1274 def __invert__(self):
Ethan Furman7aaeb2a2021-01-25 14:26:19 -08001275 if self._inverted_ is None:
1276 if self._boundary_ is KEEP:
1277 # use all bits
1278 self._inverted_ = self.__class__(~self._value_)
1279 else:
1280 # calculate flags not in this member
1281 self._inverted_ = self.__class__(self._flag_mask_ ^ self._value_)
1282 self._inverted_._inverted_ = self
1283 return self._inverted_
Ethan Furmanee47e5c2016-08-31 00:12:15 -07001284
1285
Ethan Furman7aaeb2a2021-01-25 14:26:19 -08001286class IntFlag(int, Flag, boundary=EJECT):
Ethan Furman6d3dfee2020-12-08 12:26:56 -08001287 """
1288 Support for integer-based Flags
1289 """
Ethan Furmanee47e5c2016-08-31 00:12:15 -07001290
Ethan Furmanee47e5c2016-08-31 00:12:15 -07001291 def __or__(self, other):
Ethan Furman7aaeb2a2021-01-25 14:26:19 -08001292 if isinstance(other, self.__class__):
1293 other = other._value_
1294 elif isinstance(other, int):
1295 other = other
1296 else:
Ethan Furmanee47e5c2016-08-31 00:12:15 -07001297 return NotImplemented
Ethan Furman7aaeb2a2021-01-25 14:26:19 -08001298 value = self._value_
1299 return self.__class__(value | other)
Ethan Furmanee47e5c2016-08-31 00:12:15 -07001300
1301 def __and__(self, other):
Ethan Furman7aaeb2a2021-01-25 14:26:19 -08001302 if isinstance(other, self.__class__):
1303 other = other._value_
1304 elif isinstance(other, int):
1305 other = other
1306 else:
Ethan Furmanee47e5c2016-08-31 00:12:15 -07001307 return NotImplemented
Ethan Furman7aaeb2a2021-01-25 14:26:19 -08001308 value = self._value_
1309 return self.__class__(value & other)
Ethan Furmanee47e5c2016-08-31 00:12:15 -07001310
1311 def __xor__(self, other):
Ethan Furman7aaeb2a2021-01-25 14:26:19 -08001312 if isinstance(other, self.__class__):
1313 other = other._value_
1314 elif isinstance(other, int):
1315 other = other
1316 else:
Ethan Furmanee47e5c2016-08-31 00:12:15 -07001317 return NotImplemented
Ethan Furman7aaeb2a2021-01-25 14:26:19 -08001318 value = self._value_
1319 return self.__class__(value ^ other)
Ethan Furmanee47e5c2016-08-31 00:12:15 -07001320
1321 __ror__ = __or__
1322 __rand__ = __and__
1323 __rxor__ = __xor__
Ethan Furman7aaeb2a2021-01-25 14:26:19 -08001324 __invert__ = Flag.__invert__
Ethan Furmanee47e5c2016-08-31 00:12:15 -07001325
1326def _high_bit(value):
Ethan Furman6d3dfee2020-12-08 12:26:56 -08001327 """
1328 returns index of highest bit, or -1 if value is zero or negative
1329 """
Ethan Furman3515dcc2016-09-18 13:15:41 -07001330 return value.bit_length() - 1
Ethan Furmanee47e5c2016-08-31 00:12:15 -07001331
Ethan Furmanf24bb352013-07-18 17:05:39 -07001332def unique(enumeration):
Ethan Furman6d3dfee2020-12-08 12:26:56 -08001333 """
1334 Class decorator for enumerations ensuring unique member values.
1335 """
Ethan Furmanf24bb352013-07-18 17:05:39 -07001336 duplicates = []
1337 for name, member in enumeration.__members__.items():
1338 if name != member.name:
1339 duplicates.append((name, member.name))
1340 if duplicates:
1341 alias_details = ', '.join(
1342 ["%s -> %s" % (alias, name) for (alias, name) in duplicates])
1343 raise ValueError('duplicate values found in %r: %s' %
1344 (enumeration, alias_details))
1345 return enumeration
Ethan Furman3515dcc2016-09-18 13:15:41 -07001346
Ethan Furman7aaeb2a2021-01-25 14:26:19 -08001347def _power_of_two(value):
1348 if value < 1:
1349 return False
1350 return value == 2 ** _high_bit(value)
Ethan Furmanb7751062021-03-30 21:17:26 -07001351
1352def global_enum_repr(self):
1353 return '%s.%s' % (self.__class__.__module__, self._name_)
1354
1355def global_flag_repr(self):
1356 module = self.__class__.__module__
1357 cls_name = self.__class__.__name__
1358 if self._name_ is None:
1359 return "%x" % (module, cls_name, self._value_)
1360 if _is_single_bit(self):
1361 return '%s.%s' % (module, self._name_)
1362 if self._boundary_ is not FlagBoundary.KEEP:
1363 return module + module.join(self.name.split('|'))
1364 else:
1365 name = []
1366 for n in self._name_.split('|'):
1367 if n.startswith('0'):
1368 name.append(n)
1369 else:
1370 name.append('%s.%s' % (module, n))
1371 return '|'.join(name)
1372
1373
1374def global_enum(cls):
1375 """
1376 decorator that makes the repr() of an enum member reference its module
1377 instead of its class; also exports all members to the enum's module's
1378 global namespace
1379 """
1380 if issubclass(cls, Flag):
1381 cls.__repr__ = global_flag_repr
1382 else:
1383 cls.__repr__ = global_enum_repr
1384 sys.modules[cls.__module__].__dict__.update(cls.__members__)
1385 return cls