blob: bcf411c6b6550cbbbeeb437caee783c502c23c63 [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 Furmana02cb472021-04-21 10:20:44 -0700394 def __new__(metacls, cls, bases, classdict, *, boundary=None, _simple=False, **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_
Ethan Furmana02cb472021-04-21 10:20:44 -0700401 if _simple:
402 return super().__new__(metacls, cls, bases, classdict, **kwds)
Ethan Furmana4b1bb42018-01-22 07:56:37 -0800403 classdict.setdefault('_ignore_', []).append('_ignore_')
404 ignore = classdict['_ignore_']
405 for key in ignore:
406 classdict.pop(key, None)
Ethan Furmanc314e602021-01-12 23:47:57 -0800407 #
408 # grab member names
409 member_names = classdict._member_names
410 #
411 # check for illegal enum names (any others?)
412 invalid_names = set(member_names) & {'mro', ''}
413 if invalid_names:
414 raise ValueError('Invalid enum member name: {0}'.format(
415 ','.join(invalid_names)))
416 #
417 # adjust the sunders
418 _order_ = classdict.pop('_order_', None)
419 # convert to normal dict
420 classdict = dict(classdict.items())
421 #
422 # data type of member and the controlling Enum class
Ethan Furman3064dbf2020-09-16 07:11:57 -0700423 member_type, first_enum = metacls._get_mixins_(cls, bases)
Ethan Furmanc2667362020-12-07 00:17:31 -0800424 __new__, save_new, use_args = metacls._find_new_(
425 classdict, member_type, first_enum,
426 )
Ethan Furmanc314e602021-01-12 23:47:57 -0800427 classdict['_new_member_'] = __new__
428 classdict['_use_args_'] = use_args
429 #
430 # convert future enum members into temporary _proto_members
Ethan Furman7aaeb2a2021-01-25 14:26:19 -0800431 # and record integer values in case this will be a Flag
432 flag_mask = 0
Ethan Furmanc314e602021-01-12 23:47:57 -0800433 for name in member_names:
Ethan Furman7aaeb2a2021-01-25 14:26:19 -0800434 value = classdict[name]
435 if isinstance(value, int):
436 flag_mask |= value
437 classdict[name] = _proto_member(value)
Ethan Furmanc314e602021-01-12 23:47:57 -0800438 #
Ethan Furman7aaeb2a2021-01-25 14:26:19 -0800439 # house-keeping structures
Ethan Furmanc314e602021-01-12 23:47:57 -0800440 classdict['_member_names_'] = []
441 classdict['_member_map_'] = {}
442 classdict['_value2member_map_'] = {}
443 classdict['_member_type_'] = member_type
444 #
Ethan Furman7aaeb2a2021-01-25 14:26:19 -0800445 # Flag structures (will be removed if final class is not a Flag
446 classdict['_boundary_'] = (
447 boundary
448 or getattr(first_enum, '_boundary_', None)
449 )
450 classdict['_flag_mask_'] = flag_mask
451 classdict['_all_bits_'] = 2 ** ((flag_mask).bit_length()) - 1
452 classdict['_inverted_'] = None
453 #
Ethan Furman2da95042014-03-03 12:42:52 -0800454 # If a custom type is mixed into the Enum, and it does not know how
455 # to pickle itself, pickle.dumps will succeed but pickle.loads will
456 # fail. Rather than have the error show up later and possibly far
457 # from the source, sabotage the pickle protocol for this class so
458 # that pickle.dumps also fails.
459 #
460 # However, if the new class implements its own __reduce_ex__, do not
461 # sabotage -- it's on them to make sure it works correctly. We use
462 # __reduce_ex__ instead of any of the others as it is preferred by
463 # pickle over __reduce__, and it handles all pickle protocols.
464 if '__reduce_ex__' not in classdict:
Ethan Furmandc870522014-02-18 12:37:12 -0800465 if member_type is not object:
466 methods = ('__getnewargs_ex__', '__getnewargs__',
467 '__reduce_ex__', '__reduce__')
Ethan Furman2da95042014-03-03 12:42:52 -0800468 if not any(m in member_type.__dict__ for m in methods):
Ethan Furmanc314e602021-01-12 23:47:57 -0800469 _make_class_unpicklable(classdict)
470 #
471 # create a default docstring if one has not been provided
472 if '__doc__' not in classdict:
473 classdict['__doc__'] = 'An enumeration.'
474 try:
475 exc = None
476 enum_class = super().__new__(metacls, cls, bases, classdict, **kwds)
477 except RuntimeError as e:
478 # any exceptions raised by member.__new__ will get converted to a
479 # RuntimeError, so get that original exception back and raise it instead
480 exc = e.__cause__ or e
481 if exc is not None:
482 raise exc
483 #
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700484 # double check that repr and friends are not the mixin's or various
485 # things break (such as pickle)
Ethan Furman22415ad2020-09-15 16:28:25 -0700486 # however, if the method is defined in the Enum itself, don't replace
487 # it
Ethan Furmandc870522014-02-18 12:37:12 -0800488 for name in ('__repr__', '__str__', '__format__', '__reduce_ex__'):
Ethan Furman22415ad2020-09-15 16:28:25 -0700489 if name in classdict:
490 continue
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700491 class_method = getattr(enum_class, name)
492 obj_method = getattr(member_type, name, None)
493 enum_method = getattr(first_enum, name, None)
494 if obj_method is not None and obj_method is class_method:
495 setattr(enum_class, name, enum_method)
Ethan Furmanc314e602021-01-12 23:47:57 -0800496 #
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700497 # replace any other __new__ with our own (as long as Enum is not None,
498 # anyway) -- again, this is to support pickle
499 if Enum is not None:
500 # if the user defined their own __new__, save it before it gets
501 # clobbered in case they subclass later
502 if save_new:
503 enum_class.__new_member__ = __new__
504 enum_class.__new__ = Enum.__new__
Ethan Furmanc314e602021-01-12 23:47:57 -0800505 #
Ethan Furmane8e61272016-08-20 07:19:31 -0700506 # py3 support for definition order (helps keep py2/py3 code in sync)
Ethan Furman7aaeb2a2021-01-25 14:26:19 -0800507 #
508 # _order_ checking is spread out into three/four steps
509 # - if enum_class is a Flag:
510 # - remove any non-single-bit flags from _order_
511 # - remove any aliases from _order_
512 # - check that _order_ and _member_names_ match
513 #
514 # step 1: ensure we have a list
Ethan Furmane8e61272016-08-20 07:19:31 -0700515 if _order_ is not None:
516 if isinstance(_order_, str):
517 _order_ = _order_.replace(',', ' ').split()
Ethan Furman7aaeb2a2021-01-25 14:26:19 -0800518 #
519 # remove Flag structures if final class is not a Flag
520 if (
521 Flag is None and cls != 'Flag'
522 or Flag is not None and not issubclass(enum_class, Flag)
523 ):
524 delattr(enum_class, '_boundary_')
525 delattr(enum_class, '_flag_mask_')
526 delattr(enum_class, '_all_bits_')
527 delattr(enum_class, '_inverted_')
528 elif Flag is not None and issubclass(enum_class, Flag):
529 # ensure _all_bits_ is correct and there are no missing flags
530 single_bit_total = 0
531 multi_bit_total = 0
532 for flag in enum_class._member_map_.values():
533 flag_value = flag._value_
534 if _is_single_bit(flag_value):
535 single_bit_total |= flag_value
536 else:
537 # multi-bit flags are considered aliases
538 multi_bit_total |= flag_value
539 if enum_class._boundary_ is not KEEP:
540 missed = list(_iter_bits_lsb(multi_bit_total & ~single_bit_total))
541 if missed:
542 raise TypeError(
543 'invalid Flag %r -- missing values: %s'
544 % (cls, ', '.join((str(i) for i in missed)))
545 )
546 enum_class._flag_mask_ = single_bit_total
547 #
548 # set correct __iter__
549 member_list = [m._value_ for m in enum_class]
550 if member_list != sorted(member_list):
551 enum_class._iter_member_ = enum_class._iter_member_by_def_
552 if _order_:
553 # _order_ step 2: remove any items from _order_ that are not single-bit
554 _order_ = [
555 o
556 for o in _order_
557 if o not in enum_class._member_map_ or _is_single_bit(enum_class[o]._value_)
558 ]
559 #
560 if _order_:
561 # _order_ step 3: remove aliases from _order_
562 _order_ = [
563 o
564 for o in _order_
565 if (
566 o not in enum_class._member_map_
567 or
568 (o in enum_class._member_map_ and o in enum_class._member_names_)
569 )]
570 # _order_ step 4: verify that _order_ and _member_names_ match
Ethan Furmane8e61272016-08-20 07:19:31 -0700571 if _order_ != enum_class._member_names_:
Ethan Furman7aaeb2a2021-01-25 14:26:19 -0800572 raise TypeError(
573 'member order does not match _order_:\n%r\n%r'
574 % (enum_class._member_names_, _order_)
575 )
Ethan Furmanc314e602021-01-12 23:47:57 -0800576 #
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700577 return enum_class
578
Ethan Furman5de67b12016-04-13 23:52:09 -0700579 def __bool__(self):
580 """
581 classes/types should always be True.
582 """
583 return True
584
Ethan Furman7aaeb2a2021-01-25 14:26:19 -0800585 def __call__(cls, value, names=None, *, module=None, qualname=None, type=None, start=1, boundary=None):
Ethan Furman6d3dfee2020-12-08 12:26:56 -0800586 """
587 Either returns an existing member, or creates a new enum class.
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700588
589 This method is used both when an enum class is given a value to match
590 to an enumeration member (i.e. Color(3)) and for the functional API
Ethan Furman23bb6f42016-11-21 09:22:05 -0800591 (i.e. Color = Enum('Color', names='RED GREEN BLUE')).
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700592
Ethan Furman2da95042014-03-03 12:42:52 -0800593 When used for the functional API:
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700594
Ethan Furman2da95042014-03-03 12:42:52 -0800595 `value` will be the name of the new class.
596
597 `names` should be either a string of white-space/comma delimited names
Ethan Furmand9925a12014-09-16 20:35:55 -0700598 (values will start at `start`), or an iterator/mapping of name, value pairs.
Ethan Furman2da95042014-03-03 12:42:52 -0800599
600 `module` should be set to the module this class is being created in;
601 if it is not set, an attempt to find that module will be made, but if
602 it fails the class will not be picklable.
603
604 `qualname` should be set to the actual location this class can be found
605 at in its module; by default it is set to the global scope. If this is
606 not correct, unpickling will fail in some circumstances.
607
608 `type`, if set, will be mixed in as the first base class.
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700609 """
610 if names is None: # simple value lookup
611 return cls.__new__(cls, value)
612 # otherwise, functional API: we're creating a new Enum type
Ethan Furman6d3dfee2020-12-08 12:26:56 -0800613 return cls._create_(
614 value,
615 names,
616 module=module,
617 qualname=qualname,
618 type=type,
619 start=start,
Ethan Furman7aaeb2a2021-01-25 14:26:19 -0800620 boundary=boundary,
Ethan Furman6d3dfee2020-12-08 12:26:56 -0800621 )
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700622
623 def __contains__(cls, member):
Rahul Jha94306522018-09-10 23:51:04 +0530624 if not isinstance(member, Enum):
625 raise TypeError(
626 "unsupported operand type(s) for 'in': '%s' and '%s'" % (
627 type(member).__qualname__, cls.__class__.__qualname__))
Ethan Furman0081f232014-09-16 17:31:23 -0700628 return isinstance(member, cls) and member._name_ in cls._member_map_
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700629
Ethan Furman64a99722013-09-22 16:18:19 -0700630 def __delattr__(cls, attr):
631 # nicer error message when someone tries to delete an attribute
632 # (see issue19025).
633 if attr in cls._member_map_:
Ethan Furman6d3dfee2020-12-08 12:26:56 -0800634 raise AttributeError("%s: cannot delete Enum member %r." % (cls.__name__, attr))
Ethan Furman64a99722013-09-22 16:18:19 -0700635 super().__delattr__(attr)
636
Ethan Furman388a3922013-08-12 06:51:41 -0700637 def __dir__(self):
Ethan Furman6d3dfee2020-12-08 12:26:56 -0800638 return (
639 ['__class__', '__doc__', '__members__', '__module__']
640 + self._member_names_
641 )
Ethan Furman388a3922013-08-12 06:51:41 -0700642
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700643 def __getattr__(cls, name):
Ethan Furman6d3dfee2020-12-08 12:26:56 -0800644 """
645 Return the enum member matching `name`
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700646
647 We use __getattr__ instead of descriptors or inserting into the enum
648 class' __dict__ in order to support `name` and `value` being both
649 properties for enum members (which live in the class' __dict__) and
650 enum members themselves.
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700651 """
652 if _is_dunder(name):
653 raise AttributeError(name)
654 try:
Ethan Furman520ad572013-07-19 19:47:21 -0700655 return cls._member_map_[name]
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700656 except KeyError:
657 raise AttributeError(name) from None
658
659 def __getitem__(cls, name):
Ethan Furman520ad572013-07-19 19:47:21 -0700660 return cls._member_map_[name]
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700661
662 def __iter__(cls):
Ethan Furman6d3dfee2020-12-08 12:26:56 -0800663 """
664 Returns members in definition order.
665 """
Ethan Furman520ad572013-07-19 19:47:21 -0700666 return (cls._member_map_[name] for name in cls._member_names_)
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700667
668 def __len__(cls):
Ethan Furman520ad572013-07-19 19:47:21 -0700669 return len(cls._member_names_)
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700670
Ethan Furmanc314e602021-01-12 23:47:57 -0800671 @_bltin_property
Ethan Furman2131a4a2013-09-14 18:11:24 -0700672 def __members__(cls):
Ethan Furman6d3dfee2020-12-08 12:26:56 -0800673 """
674 Returns a mapping of member name->value.
Ethan Furman2131a4a2013-09-14 18:11:24 -0700675
676 This mapping lists all enum members, including aliases. Note that this
677 is a read-only view of the internal mapping.
Ethan Furman2131a4a2013-09-14 18:11:24 -0700678 """
679 return MappingProxyType(cls._member_map_)
680
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700681 def __repr__(cls):
682 return "<enum %r>" % cls.__name__
683
Ethan Furman2131a4a2013-09-14 18:11:24 -0700684 def __reversed__(cls):
Ethan Furman6d3dfee2020-12-08 12:26:56 -0800685 """
686 Returns members in reverse definition order.
687 """
Ethan Furman2131a4a2013-09-14 18:11:24 -0700688 return (cls._member_map_[name] for name in reversed(cls._member_names_))
689
Ethan Furmanf203f2d2013-09-06 07:16:48 -0700690 def __setattr__(cls, name, value):
Ethan Furman6d3dfee2020-12-08 12:26:56 -0800691 """
692 Block attempts to reassign Enum members.
Ethan Furmanf203f2d2013-09-06 07:16:48 -0700693
694 A simple assignment to the class namespace only changes one of the
695 several possible ways to get an Enum member from the Enum class,
696 resulting in an inconsistent Enumeration.
Ethan Furmanf203f2d2013-09-06 07:16:48 -0700697 """
698 member_map = cls.__dict__.get('_member_map_', {})
699 if name in member_map:
Ethan Furmana02cb472021-04-21 10:20:44 -0700700 raise AttributeError('Cannot reassign member %r.' % (name, ))
Ethan Furmanf203f2d2013-09-06 07:16:48 -0700701 super().__setattr__(name, value)
702
Ethan Furman7aaeb2a2021-01-25 14:26:19 -0800703 def _create_(cls, class_name, names, *, module=None, qualname=None, type=None, start=1, boundary=None):
Ethan Furman6d3dfee2020-12-08 12:26:56 -0800704 """
705 Convenience method to create a new Enum class.
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700706
707 `names` can be:
708
709 * A string containing member names, separated either with spaces or
Ethan Furmand9925a12014-09-16 20:35:55 -0700710 commas. Values are incremented by 1 from `start`.
711 * An iterable of member names. Values are incremented by 1 from `start`.
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700712 * An iterable of (member name, value) pairs.
Ethan Furmand9925a12014-09-16 20:35:55 -0700713 * A mapping of member name -> value pairs.
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700714 """
715 metacls = cls.__class__
716 bases = (cls, ) if type is None else (type, cls)
Ethan Furman3064dbf2020-09-16 07:11:57 -0700717 _, first_enum = cls._get_mixins_(cls, bases)
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700718 classdict = metacls.__prepare__(class_name, bases)
719
720 # special processing needed for names?
721 if isinstance(names, str):
722 names = names.replace(',', ' ').split()
Dong-hee Nadcc8ce42017-06-22 01:52:32 +0900723 if isinstance(names, (tuple, list)) and names and isinstance(names[0], str):
Ethan Furmanee47e5c2016-08-31 00:12:15 -0700724 original_names, names = names, []
Ethan Furmanc16595e2016-09-10 23:36:59 -0700725 last_values = []
Ethan Furmanee47e5c2016-08-31 00:12:15 -0700726 for count, name in enumerate(original_names):
Ethan Furmanc16595e2016-09-10 23:36:59 -0700727 value = first_enum._generate_next_value_(name, start, count, last_values[:])
728 last_values.append(value)
729 names.append((name, value))
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700730
731 # Here, names is either an iterable of (name, value) or a mapping.
732 for item in names:
733 if isinstance(item, str):
734 member_name, member_value = item, names[item]
735 else:
736 member_name, member_value = item
737 classdict[member_name] = member_value
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700738
739 # TODO: replace the frame hack if a blessed way to know the calling
740 # module is ever developed
741 if module is None:
742 try:
743 module = sys._getframe(2).f_globals['__name__']
Pablo Galindo293dd232019-11-19 21:34:03 +0000744 except (AttributeError, ValueError, KeyError):
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700745 pass
746 if module is None:
Ethan Furmanc314e602021-01-12 23:47:57 -0800747 _make_class_unpicklable(classdict)
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700748 else:
Ethan Furmanc314e602021-01-12 23:47:57 -0800749 classdict['__module__'] = module
Ethan Furmanca1b7942014-02-08 11:36:27 -0800750 if qualname is not None:
Ethan Furmanc314e602021-01-12 23:47:57 -0800751 classdict['__qualname__'] = qualname
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700752
Ethan Furman7aaeb2a2021-01-25 14:26:19 -0800753 return metacls.__new__(metacls, class_name, bases, classdict, boundary=boundary)
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700754
Ethan Furmana02cb472021-04-21 10:20:44 -0700755 def _convert_(cls, name, module, filter, source=None, *, boundary=None):
756
orlnub1230fb9fad2018-09-12 20:28:53 +0300757 """
758 Create a new Enum subclass that replaces a collection of global constants
759 """
760 # convert all constants from source (or module) that pass filter() to
761 # a new Enum called name, and export the enum and its members back to
762 # module;
763 # also, replace the __reduce_ex__ method so unpickling works in
764 # previous Python versions
Ethan Furmanb7751062021-03-30 21:17:26 -0700765 module_globals = sys.modules[module].__dict__
orlnub1230fb9fad2018-09-12 20:28:53 +0300766 if source:
Ethan Furmanb7751062021-03-30 21:17:26 -0700767 source = source.__dict__
orlnub1230fb9fad2018-09-12 20:28:53 +0300768 else:
769 source = module_globals
770 # _value2member_map_ is populated in the same order every time
771 # for a consistent reverse mapping of number to name when there
772 # are multiple names for the same number.
773 members = [
774 (name, value)
775 for name, value in source.items()
776 if filter(name)]
777 try:
778 # sort by value
779 members.sort(key=lambda t: (t[1], t[0]))
780 except TypeError:
781 # unless some values aren't comparable, in which case sort by name
782 members.sort(key=lambda t: t[0])
Ethan Furmana02cb472021-04-21 10:20:44 -0700783 body = {t[0]: t[1] for t in members}
784 body['__module__'] = module
785 tmp_cls = type(name, (object, ), body)
786 cls = _simple_enum(etype=cls, boundary=boundary or KEEP)(tmp_cls)
orlnub1230fb9fad2018-09-12 20:28:53 +0300787 cls.__reduce_ex__ = _reduce_ex_by_name
Ethan Furmanb7751062021-03-30 21:17:26 -0700788 global_enum(cls)
orlnub1230fb9fad2018-09-12 20:28:53 +0300789 module_globals[name] = cls
790 return cls
791
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700792 @staticmethod
Ethan Furman3064dbf2020-09-16 07:11:57 -0700793 def _check_for_existing_members(class_name, bases):
794 for chain in bases:
795 for base in chain.__mro__:
796 if issubclass(base, Enum) and base._member_names_:
Ethan Furman6d3dfee2020-12-08 12:26:56 -0800797 raise TypeError(
798 "%s: cannot extend enumeration %r"
799 % (class_name, base.__name__)
800 )
Ethan Furman3064dbf2020-09-16 07:11:57 -0700801
802 @staticmethod
803 def _get_mixins_(class_name, bases):
Ethan Furman6d3dfee2020-12-08 12:26:56 -0800804 """
805 Returns the type for creating enum members, and the first inherited
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700806 enum class.
807
808 bases: the tuple of bases that was given to __new__
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700809 """
810 if not bases:
811 return object, Enum
812
Ethan Furman5bdab642018-09-21 19:03:09 -0700813 def _find_data_type(bases):
Ethan Furmanbff01f32020-09-15 15:56:26 -0700814 data_types = []
Ethan Furman5bdab642018-09-21 19:03:09 -0700815 for chain in bases:
Ethan Furmanbff01f32020-09-15 15:56:26 -0700816 candidate = None
Ethan Furman5bdab642018-09-21 19:03:09 -0700817 for base in chain.__mro__:
818 if base is object:
819 continue
Ethan Furmanc2667362020-12-07 00:17:31 -0800820 elif issubclass(base, Enum):
821 if base._member_type_ is not object:
822 data_types.append(base._member_type_)
823 break
Ethan Furman5bdab642018-09-21 19:03:09 -0700824 elif '__new__' in base.__dict__:
Ethan Furmancd453852018-10-05 23:29:36 -0700825 if issubclass(base, Enum):
Ethan Furman5bdab642018-09-21 19:03:09 -0700826 continue
Ethan Furmanbff01f32020-09-15 15:56:26 -0700827 data_types.append(candidate or base)
828 break
Ethan Furmanc2667362020-12-07 00:17:31 -0800829 else:
Ethan Furmanbff01f32020-09-15 15:56:26 -0700830 candidate = base
831 if len(data_types) > 1:
Ethan Furman3064dbf2020-09-16 07:11:57 -0700832 raise TypeError('%r: too many data types: %r' % (class_name, data_types))
Ethan Furmanbff01f32020-09-15 15:56:26 -0700833 elif data_types:
834 return data_types[0]
835 else:
836 return None
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700837
Ethan Furman5bdab642018-09-21 19:03:09 -0700838 # ensure final parent class is an Enum derivative, find any concrete
839 # data type, and check that Enum has no members
840 first_enum = bases[-1]
841 if not issubclass(first_enum, Enum):
842 raise TypeError("new enumerations should be created as "
843 "`EnumName([mixin_type, ...] [data_type,] enum_type)`")
844 member_type = _find_data_type(bases) or object
845 if first_enum._member_names_:
846 raise TypeError("Cannot extend enumerations")
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700847 return member_type, first_enum
848
849 @staticmethod
850 def _find_new_(classdict, member_type, first_enum):
Ethan Furman6d3dfee2020-12-08 12:26:56 -0800851 """
852 Returns the __new__ to be used for creating the enum members.
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700853
854 classdict: the class dictionary given to __new__
855 member_type: the data type whose __new__ will be used by default
856 first_enum: enumeration to check for an overriding __new__
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700857 """
858 # now find the correct __new__, checking to see of one was defined
859 # by the user; also check earlier enum classes in case a __new__ was
860 # saved as __new_member__
861 __new__ = classdict.get('__new__', None)
862
863 # should __new__ be saved as __new_member__ later?
Ethan Furmana02cb472021-04-21 10:20:44 -0700864 save_new = first_enum is not None and __new__ is not None
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700865
866 if __new__ is None:
867 # check all possibles for __new_member__ before falling back to
868 # __new__
869 for method in ('__new_member__', '__new__'):
870 for possible in (member_type, first_enum):
871 target = getattr(possible, method, None)
872 if target not in {
873 None,
874 None.__new__,
875 object.__new__,
876 Enum.__new__,
877 }:
878 __new__ = target
879 break
880 if __new__ is not None:
881 break
882 else:
883 __new__ = object.__new__
884
885 # if a non-object.__new__ is used then whatever value/tuple was
886 # assigned to the enum member name will be passed to __new__ and to the
887 # new enum member's __init__
Ethan Furmana02cb472021-04-21 10:20:44 -0700888 if first_enum is None or __new__ in (Enum.__new__, object.__new__):
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700889 use_args = False
890 else:
891 use_args = True
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700892 return __new__, save_new, use_args
Ethan Furmanb7751062021-03-30 21:17:26 -0700893EnumMeta = EnumType
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700894
895
Ethan Furmanb7751062021-03-30 21:17:26 -0700896class Enum(metaclass=EnumType):
Ethan Furman6d3dfee2020-12-08 12:26:56 -0800897 """
898 Generic enumeration.
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700899
900 Derive from this class to define new enumerations.
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700901 """
Ethan Furman7aaeb2a2021-01-25 14:26:19 -0800902
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700903 def __new__(cls, value):
904 # all enum instances are actually created during class construction
905 # without calling this method; this method is called by the metaclass'
906 # __call__ (i.e. Color(3) ), and by pickle
907 if type(value) is cls:
Ethan Furman23bb6f42016-11-21 09:22:05 -0800908 # For lookups like Color(Color.RED)
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700909 return value
910 # by-value search for a matching enum member
911 # see if it's in the reverse mapping (for hashable values)
Ethan Furman2aa27322013-07-19 19:35:56 -0700912 try:
Andrew Svetlov34ae04f2018-12-26 20:45:33 +0200913 return cls._value2member_map_[value]
914 except KeyError:
915 # Not found, no need to do long O(n) search
916 pass
Ethan Furman2aa27322013-07-19 19:35:56 -0700917 except TypeError:
918 # not there, now do long search -- O(n) behavior
Ethan Furman520ad572013-07-19 19:47:21 -0700919 for member in cls._member_map_.values():
Ethan Furman0081f232014-09-16 17:31:23 -0700920 if member._value_ == value:
Ethan Furman2aa27322013-07-19 19:35:56 -0700921 return member
Ethan Furmanee47e5c2016-08-31 00:12:15 -0700922 # still not found -- try _missing_ hook
Ethan Furman019f0a02018-09-12 11:43:34 -0700923 try:
924 exc = None
925 result = cls._missing_(value)
926 except Exception as e:
927 exc = e
928 result = None
Ethan Furman8c14f5a2021-04-12 08:51:20 -0700929 try:
930 if isinstance(result, cls):
931 return result
932 elif (
933 Flag is not None and issubclass(cls, Flag)
934 and cls._boundary_ is EJECT and isinstance(result, int)
935 ):
936 return result
937 else:
938 ve_exc = ValueError("%r is not a valid %s" % (value, cls.__qualname__))
939 if result is None and exc is None:
940 raise ve_exc
941 elif exc is None:
942 exc = TypeError(
943 'error in %s._missing_: returned %r instead of None or a valid member'
944 % (cls.__name__, result)
945 )
946 if not isinstance(exc, ValueError):
947 exc.__context__ = ve_exc
948 raise exc
949 finally:
950 # ensure all variables that could hold an exception are destroyed
951 exc = None
952 ve_exc = None
Ethan Furmanee47e5c2016-08-31 00:12:15 -0700953
Ethan Furmanc16595e2016-09-10 23:36:59 -0700954 def _generate_next_value_(name, start, count, last_values):
Ethan Furman6d3dfee2020-12-08 12:26:56 -0800955 """
956 Generate the next value when not given.
957
958 name: the name of the member
959 start: the initial start value or None
960 count: the number of existing members
961 last_value: the last value assigned or None
962 """
Ethan Furmanc16595e2016-09-10 23:36:59 -0700963 for last_value in reversed(last_values):
964 try:
965 return last_value + 1
966 except TypeError:
967 pass
968 else:
Ethan Furmanee47e5c2016-08-31 00:12:15 -0700969 return start
Ethan Furmanc16595e2016-09-10 23:36:59 -0700970
Ethan Furmanee47e5c2016-08-31 00:12:15 -0700971 @classmethod
972 def _missing_(cls, value):
Ethan Furmanc95ad7a2020-09-16 10:26:50 -0700973 return None
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700974
975 def __repr__(self):
Ethan Furmanb7751062021-03-30 21:17:26 -0700976 return "%s.%s" % ( self.__class__.__name__, self._name_)
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700977
978 def __str__(self):
Ethan Furmanb7751062021-03-30 21:17:26 -0700979 return "%s" % (self._name_, )
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700980
Ethan Furman388a3922013-08-12 06:51:41 -0700981 def __dir__(self):
Ethan Furman6d3dfee2020-12-08 12:26:56 -0800982 """
983 Returns all members and all public methods
984 """
Ethan Furman0ae550b2014-10-14 08:58:32 -0700985 added_behavior = [
986 m
987 for cls in self.__class__.mro()
988 for m in cls.__dict__
Ethan Furman354ecf12015-03-11 08:43:12 -0700989 if m[0] != '_' and m not in self._member_map_
Angelin BOOZ68526fe2020-09-21 15:11:06 +0200990 ] + [m for m in self.__dict__ if m[0] != '_']
Ethan Furmanec5f8eb2014-10-21 13:40:35 -0700991 return (['__class__', '__doc__', '__module__'] + added_behavior)
Ethan Furman388a3922013-08-12 06:51:41 -0700992
Ethan Furmanec15a822013-08-31 19:17:41 -0700993 def __format__(self, format_spec):
Ethan Furman6d3dfee2020-12-08 12:26:56 -0800994 """
995 Returns format using actual value type unless __str__ has been overridden.
996 """
Ethan Furmanec15a822013-08-31 19:17:41 -0700997 # mixed-in Enums should use the mixed-in type's __format__, otherwise
998 # we can get strange results with the Enum name showing up instead of
999 # the value
1000
thatneat2f19e822019-07-04 11:28:37 -07001001 # pure Enum branch, or branch with __str__ explicitly overridden
Ethan Furman37440ee2020-12-08 11:14:10 -08001002 str_overridden = type(self).__str__ not in (Enum.__str__, Flag.__str__)
thatneat2f19e822019-07-04 11:28:37 -07001003 if self._member_type_ is object or str_overridden:
Ethan Furmanec15a822013-08-31 19:17:41 -07001004 cls = str
1005 val = str(self)
1006 # mix-in branch
1007 else:
Ethan Furman5987b8c2021-04-26 22:42:57 -07001008 import warnings
1009 warnings.warn(
1010 "in 3.12 format() will use the enum member, not the enum member's value;\n"
1011 "use a format specifier, such as :d for an IntEnum member, to maintain"
1012 "the current display",
1013 DeprecationWarning,
1014 stacklevel=2,
1015 )
Ethan Furmanec15a822013-08-31 19:17:41 -07001016 cls = self._member_type_
Ethan Furman0081f232014-09-16 17:31:23 -07001017 val = self._value_
Ethan Furmanec15a822013-08-31 19:17:41 -07001018 return cls.__format__(val, format_spec)
1019
Ethan Furman6b3d64a2013-06-14 16:55:46 -07001020 def __hash__(self):
Ethan Furman520ad572013-07-19 19:47:21 -07001021 return hash(self._name_)
Ethan Furman6b3d64a2013-06-14 16:55:46 -07001022
Ethan Furmanca1b7942014-02-08 11:36:27 -08001023 def __reduce_ex__(self, proto):
Ethan Furmandc870522014-02-18 12:37:12 -08001024 return self.__class__, (self._value_, )
Ethan Furmanca1b7942014-02-08 11:36:27 -08001025
Ethan Furmanc314e602021-01-12 23:47:57 -08001026 # enum.property is used to provide access to the `name` and
1027 # `value` attributes of enum members while keeping some measure of
Ethan Furman6b3d64a2013-06-14 16:55:46 -07001028 # protection from modification, while still allowing for an enumeration
1029 # to have members named `name` and `value`. This works because enumeration
Ethan Furmanc314e602021-01-12 23:47:57 -08001030 # members are not set directly on the enum class; they are kept in a
1031 # separate structure, _member_map_, which is where enum.property looks for
1032 # them
Ethan Furman6b3d64a2013-06-14 16:55:46 -07001033
Ethan Furmanc314e602021-01-12 23:47:57 -08001034 @property
Ethan Furman6b3d64a2013-06-14 16:55:46 -07001035 def name(self):
Ethan Furmanc850f342013-09-15 16:59:35 -07001036 """The name of the Enum member."""
Ethan Furman520ad572013-07-19 19:47:21 -07001037 return self._name_
Ethan Furman6b3d64a2013-06-14 16:55:46 -07001038
Ethan Furmanc314e602021-01-12 23:47:57 -08001039 @property
Ethan Furman6b3d64a2013-06-14 16:55:46 -07001040 def value(self):
Ethan Furmanc850f342013-09-15 16:59:35 -07001041 """The value of the Enum member."""
Ethan Furman520ad572013-07-19 19:47:21 -07001042 return self._value_
Ethan Furman6b3d64a2013-06-14 16:55:46 -07001043
1044
1045class IntEnum(int, Enum):
Ethan Furman0063ff42020-09-21 17:23:13 -07001046 """
1047 Enum where members are also (and must be) ints
1048 """
1049
1050
1051class StrEnum(str, Enum):
1052 """
1053 Enum where members are also (and must be) strings
1054 """
1055
1056 def __new__(cls, *values):
1057 if len(values) > 3:
1058 raise TypeError('too many arguments for str(): %r' % (values, ))
1059 if len(values) == 1:
1060 # it must be a string
1061 if not isinstance(values[0], str):
1062 raise TypeError('%r is not a string' % (values[0], ))
Ethan Furman7aaeb2a2021-01-25 14:26:19 -08001063 if len(values) >= 2:
Ethan Furman0063ff42020-09-21 17:23:13 -07001064 # check that encoding argument is a string
1065 if not isinstance(values[1], str):
1066 raise TypeError('encoding must be a string, not %r' % (values[1], ))
Ethan Furman7aaeb2a2021-01-25 14:26:19 -08001067 if len(values) == 3:
1068 # check that errors argument is a string
1069 if not isinstance(values[2], str):
1070 raise TypeError('errors must be a string, not %r' % (values[2]))
Ethan Furman0063ff42020-09-21 17:23:13 -07001071 value = str(*values)
1072 member = str.__new__(cls, value)
1073 member._value_ = value
1074 return member
Ethan Furmanf24bb352013-07-18 17:05:39 -07001075
Ethan Furmand986d162020-09-22 13:00:07 -07001076 __str__ = str.__str__
1077
Ethan Furmanefb13be2020-12-10 12:20:06 -08001078 def _generate_next_value_(name, start, count, last_values):
1079 """
1080 Return the lower-cased version of the member name.
1081 """
1082 return name.lower()
1083
Ethan Furmanf24bb352013-07-18 17:05:39 -07001084
Ethan Furman24e837f2015-03-18 17:27:57 -07001085def _reduce_ex_by_name(self, proto):
1086 return self.name
1087
Ethan Furman7aaeb2a2021-01-25 14:26:19 -08001088class FlagBoundary(StrEnum):
1089 """
1090 control how out of range values are handled
1091 "strict" -> error is raised [default for Flag]
1092 "conform" -> extra bits are discarded
1093 "eject" -> lose flag status [default for IntFlag]
1094 "keep" -> keep flag status and all bits
1095 """
1096 STRICT = auto()
1097 CONFORM = auto()
1098 EJECT = auto()
1099 KEEP = auto()
1100STRICT, CONFORM, EJECT, KEEP = FlagBoundary
1101
1102
1103class Flag(Enum, boundary=STRICT):
Ethan Furman6d3dfee2020-12-08 12:26:56 -08001104 """
1105 Support for flags
1106 """
Ethan Furmanc16595e2016-09-10 23:36:59 -07001107
1108 def _generate_next_value_(name, start, count, last_values):
Ethan Furmanee47e5c2016-08-31 00:12:15 -07001109 """
1110 Generate the next value when not given.
1111
1112 name: the name of the member
HongWeipengbb16fb22019-09-21 13:22:54 +08001113 start: the initial start value or None
Ethan Furmanee47e5c2016-08-31 00:12:15 -07001114 count: the number of existing members
1115 last_value: the last value assigned or None
1116 """
1117 if not count:
1118 return start if start is not None else 1
Ethan Furman7aaeb2a2021-01-25 14:26:19 -08001119 last_value = max(last_values)
1120 try:
1121 high_bit = _high_bit(last_value)
1122 except Exception:
1123 raise TypeError('Invalid Flag value: %r' % last_value) from None
Ethan Furmanee47e5c2016-08-31 00:12:15 -07001124 return 2 ** (high_bit+1)
1125
1126 @classmethod
Ethan Furman7aaeb2a2021-01-25 14:26:19 -08001127 def _iter_member_by_value_(cls, value):
1128 """
1129 Extract all members from the value in definition (i.e. increasing value) order.
1130 """
1131 for val in _iter_bits_lsb(value & cls._flag_mask_):
1132 yield cls._value2member_map_.get(val)
1133
1134 _iter_member_ = _iter_member_by_value_
1135
1136 @classmethod
1137 def _iter_member_by_def_(cls, value):
1138 """
1139 Extract all members from the value in definition order.
1140 """
1141 yield from sorted(
1142 cls._iter_member_by_value_(value),
1143 key=lambda m: m._sort_order_,
1144 )
1145
1146 @classmethod
Ethan Furmanee47e5c2016-08-31 00:12:15 -07001147 def _missing_(cls, value):
Ethan Furman6d3dfee2020-12-08 12:26:56 -08001148 """
Ethan Furman0dca5eb2021-04-15 06:49:54 -07001149 Create a composite member containing all canonical members present in `value`.
1150
1151 If non-member values are present, result depends on `_boundary_` setting.
Ethan Furman3515dcc2016-09-18 13:15:41 -07001152 """
Ethan Furman7aaeb2a2021-01-25 14:26:19 -08001153 if not isinstance(value, int):
1154 raise ValueError(
1155 "%r is not a valid %s" % (value, cls.__qualname__)
1156 )
1157 # check boundaries
1158 # - value must be in range (e.g. -16 <-> +15, i.e. ~15 <-> 15)
1159 # - value must not include any skipped flags (e.g. if bit 2 is not
1160 # defined, then 0d10 is invalid)
1161 flag_mask = cls._flag_mask_
1162 all_bits = cls._all_bits_
1163 neg_value = None
1164 if (
1165 not ~all_bits <= value <= all_bits
1166 or value & (all_bits ^ flag_mask)
1167 ):
1168 if cls._boundary_ is STRICT:
1169 max_bits = max(value.bit_length(), flag_mask.bit_length())
1170 raise ValueError(
1171 "%s: invalid value: %r\n given %s\n allowed %s" % (
1172 cls.__name__, value, bin(value, max_bits), bin(flag_mask, max_bits),
1173 ))
1174 elif cls._boundary_ is CONFORM:
1175 value = value & flag_mask
1176 elif cls._boundary_ is EJECT:
1177 return value
1178 elif cls._boundary_ is KEEP:
1179 if value < 0:
1180 value = (
1181 max(all_bits+1, 2**(value.bit_length()))
1182 + value
1183 )
1184 else:
1185 raise ValueError(
1186 'unknown flag boundary: %r' % (cls._boundary_, )
1187 )
1188 if value < 0:
1189 neg_value = value
1190 value = all_bits + 1 + value
1191 # get members and unknown
1192 unknown = value & ~flag_mask
1193 member_value = value & flag_mask
1194 if unknown and cls._boundary_ is not KEEP:
1195 raise ValueError(
1196 '%s(%r) --> unknown values %r [%s]'
1197 % (cls.__name__, value, unknown, bin(unknown))
1198 )
1199 # normal Flag?
1200 __new__ = getattr(cls, '__new_member__', None)
1201 if cls._member_type_ is object and not __new__:
Ethan Furman3515dcc2016-09-18 13:15:41 -07001202 # construct a singleton enum pseudo-member
Ethan Furmanee47e5c2016-08-31 00:12:15 -07001203 pseudo_member = object.__new__(cls)
Ethan Furman7aaeb2a2021-01-25 14:26:19 -08001204 else:
1205 pseudo_member = (__new__ or cls._member_type_.__new__)(cls, value)
Ethan Furmana02cb472021-04-21 10:20:44 -07001206 if not hasattr(pseudo_member, '_value_'):
Ethan Furmanee47e5c2016-08-31 00:12:15 -07001207 pseudo_member._value_ = value
Ethan Furman7aaeb2a2021-01-25 14:26:19 -08001208 if member_value:
1209 pseudo_member._name_ = '|'.join([
1210 m._name_ for m in cls._iter_member_(member_value)
1211 ])
1212 if unknown:
1213 pseudo_member._name_ += '|0x%x' % unknown
1214 else:
1215 pseudo_member._name_ = None
1216 # use setdefault in case another thread already created a composite
1217 # with this value, but only if all members are known
1218 # note: zero is a special case -- add it
1219 if not unknown:
Ethan Furman28cf6632017-01-24 12:12:06 -08001220 pseudo_member = cls._value2member_map_.setdefault(value, pseudo_member)
Ethan Furman7aaeb2a2021-01-25 14:26:19 -08001221 if neg_value is not None:
1222 cls._value2member_map_[neg_value] = pseudo_member
Ethan Furmanee47e5c2016-08-31 00:12:15 -07001223 return pseudo_member
1224
Ethan Furmanee47e5c2016-08-31 00:12:15 -07001225 def __contains__(self, other):
Ethan Furman6d3dfee2020-12-08 12:26:56 -08001226 """
1227 Returns True if self has at least the same flags set as other.
1228 """
Ethan Furmanee47e5c2016-08-31 00:12:15 -07001229 if not isinstance(other, self.__class__):
Rahul Jha94306522018-09-10 23:51:04 +05301230 raise TypeError(
1231 "unsupported operand type(s) for 'in': '%s' and '%s'" % (
1232 type(other).__qualname__, self.__class__.__qualname__))
Ethan Furman7aaeb2a2021-01-25 14:26:19 -08001233 if other._value_ == 0 or self._value_ == 0:
1234 return False
Ethan Furmanee47e5c2016-08-31 00:12:15 -07001235 return other._value_ & self._value_ == other._value_
1236
Ethan Furman7219e272020-09-16 13:01:00 -07001237 def __iter__(self):
Ethan Furman6d3dfee2020-12-08 12:26:56 -08001238 """
Ethan Furman7aaeb2a2021-01-25 14:26:19 -08001239 Returns flags in definition order.
Ethan Furman6d3dfee2020-12-08 12:26:56 -08001240 """
Ethan Furman7aaeb2a2021-01-25 14:26:19 -08001241 yield from self._iter_member_(self._value_)
1242
1243 def __len__(self):
1244 return self._value_.bit_count()
Ethan Furman7219e272020-09-16 13:01:00 -07001245
Ethan Furmanee47e5c2016-08-31 00:12:15 -07001246 def __repr__(self):
Ethan Furmanb7751062021-03-30 21:17:26 -07001247 cls_name = self.__class__.__name__
1248 if self._name_ is None:
1249 return "0x%x" % (self._value_, )
1250 if _is_single_bit(self._value_):
1251 return '%s.%s' % (cls_name, self._name_)
1252 if self._boundary_ is not FlagBoundary.KEEP:
1253 return '%s.' % cls_name + ('|%s.' % cls_name).join(self.name.split('|'))
Ethan Furman7aaeb2a2021-01-25 14:26:19 -08001254 else:
Ethan Furmanb7751062021-03-30 21:17:26 -07001255 name = []
1256 for n in self._name_.split('|'):
1257 if n.startswith('0'):
1258 name.append(n)
1259 else:
1260 name.append('%s.%s' % (cls_name, n))
1261 return '|'.join(name)
Ethan Furmanee47e5c2016-08-31 00:12:15 -07001262
1263 def __str__(self):
1264 cls = self.__class__
Ethan Furmanb7751062021-03-30 21:17:26 -07001265 if self._name_ is None:
1266 return '%s(%x)' % (cls.__name__, self._value_)
Ethan Furmanee47e5c2016-08-31 00:12:15 -07001267 else:
Ethan Furmanb7751062021-03-30 21:17:26 -07001268 return self._name_
Ethan Furmanee47e5c2016-08-31 00:12:15 -07001269
Ethan Furman25d94bb2016-09-02 16:32:32 -07001270 def __bool__(self):
1271 return bool(self._value_)
1272
Ethan Furmanee47e5c2016-08-31 00:12:15 -07001273 def __or__(self, other):
1274 if not isinstance(other, self.__class__):
1275 return NotImplemented
1276 return self.__class__(self._value_ | other._value_)
1277
1278 def __and__(self, other):
1279 if not isinstance(other, self.__class__):
1280 return NotImplemented
1281 return self.__class__(self._value_ & other._value_)
1282
1283 def __xor__(self, other):
1284 if not isinstance(other, self.__class__):
1285 return NotImplemented
1286 return self.__class__(self._value_ ^ other._value_)
1287
1288 def __invert__(self):
Ethan Furman7aaeb2a2021-01-25 14:26:19 -08001289 if self._inverted_ is None:
1290 if self._boundary_ is KEEP:
1291 # use all bits
1292 self._inverted_ = self.__class__(~self._value_)
1293 else:
1294 # calculate flags not in this member
1295 self._inverted_ = self.__class__(self._flag_mask_ ^ self._value_)
1296 self._inverted_._inverted_ = self
1297 return self._inverted_
Ethan Furmanee47e5c2016-08-31 00:12:15 -07001298
1299
Ethan Furman7aaeb2a2021-01-25 14:26:19 -08001300class IntFlag(int, Flag, boundary=EJECT):
Ethan Furman6d3dfee2020-12-08 12:26:56 -08001301 """
1302 Support for integer-based Flags
1303 """
Ethan Furmanee47e5c2016-08-31 00:12:15 -07001304
Ethan Furmanee47e5c2016-08-31 00:12:15 -07001305 def __or__(self, other):
Ethan Furman7aaeb2a2021-01-25 14:26:19 -08001306 if isinstance(other, self.__class__):
1307 other = other._value_
1308 elif isinstance(other, int):
1309 other = other
1310 else:
Ethan Furmanee47e5c2016-08-31 00:12:15 -07001311 return NotImplemented
Ethan Furman7aaeb2a2021-01-25 14:26:19 -08001312 value = self._value_
1313 return self.__class__(value | other)
Ethan Furmanee47e5c2016-08-31 00:12:15 -07001314
1315 def __and__(self, other):
Ethan Furman7aaeb2a2021-01-25 14:26:19 -08001316 if isinstance(other, self.__class__):
1317 other = other._value_
1318 elif isinstance(other, int):
1319 other = other
1320 else:
Ethan Furmanee47e5c2016-08-31 00:12:15 -07001321 return NotImplemented
Ethan Furman7aaeb2a2021-01-25 14:26:19 -08001322 value = self._value_
1323 return self.__class__(value & other)
Ethan Furmanee47e5c2016-08-31 00:12:15 -07001324
1325 def __xor__(self, other):
Ethan Furman7aaeb2a2021-01-25 14:26:19 -08001326 if isinstance(other, self.__class__):
1327 other = other._value_
1328 elif isinstance(other, int):
1329 other = other
1330 else:
Ethan Furmanee47e5c2016-08-31 00:12:15 -07001331 return NotImplemented
Ethan Furman7aaeb2a2021-01-25 14:26:19 -08001332 value = self._value_
1333 return self.__class__(value ^ other)
Ethan Furmanee47e5c2016-08-31 00:12:15 -07001334
1335 __ror__ = __or__
1336 __rand__ = __and__
1337 __rxor__ = __xor__
Ethan Furman7aaeb2a2021-01-25 14:26:19 -08001338 __invert__ = Flag.__invert__
Ethan Furmanee47e5c2016-08-31 00:12:15 -07001339
1340def _high_bit(value):
Ethan Furman6d3dfee2020-12-08 12:26:56 -08001341 """
1342 returns index of highest bit, or -1 if value is zero or negative
1343 """
Ethan Furman3515dcc2016-09-18 13:15:41 -07001344 return value.bit_length() - 1
Ethan Furmanee47e5c2016-08-31 00:12:15 -07001345
Ethan Furmanf24bb352013-07-18 17:05:39 -07001346def unique(enumeration):
Ethan Furman6d3dfee2020-12-08 12:26:56 -08001347 """
1348 Class decorator for enumerations ensuring unique member values.
1349 """
Ethan Furmanf24bb352013-07-18 17:05:39 -07001350 duplicates = []
1351 for name, member in enumeration.__members__.items():
1352 if name != member.name:
1353 duplicates.append((name, member.name))
1354 if duplicates:
1355 alias_details = ', '.join(
1356 ["%s -> %s" % (alias, name) for (alias, name) in duplicates])
1357 raise ValueError('duplicate values found in %r: %s' %
1358 (enumeration, alias_details))
1359 return enumeration
Ethan Furman3515dcc2016-09-18 13:15:41 -07001360
Ethan Furman7aaeb2a2021-01-25 14:26:19 -08001361def _power_of_two(value):
1362 if value < 1:
1363 return False
1364 return value == 2 ** _high_bit(value)
Ethan Furmanb7751062021-03-30 21:17:26 -07001365
1366def global_enum_repr(self):
1367 return '%s.%s' % (self.__class__.__module__, self._name_)
1368
1369def global_flag_repr(self):
1370 module = self.__class__.__module__
1371 cls_name = self.__class__.__name__
1372 if self._name_ is None:
1373 return "%x" % (module, cls_name, self._value_)
1374 if _is_single_bit(self):
1375 return '%s.%s' % (module, self._name_)
1376 if self._boundary_ is not FlagBoundary.KEEP:
1377 return module + module.join(self.name.split('|'))
1378 else:
1379 name = []
1380 for n in self._name_.split('|'):
1381 if n.startswith('0'):
1382 name.append(n)
1383 else:
1384 name.append('%s.%s' % (module, n))
1385 return '|'.join(name)
1386
1387
1388def global_enum(cls):
1389 """
1390 decorator that makes the repr() of an enum member reference its module
1391 instead of its class; also exports all members to the enum's module's
1392 global namespace
1393 """
1394 if issubclass(cls, Flag):
1395 cls.__repr__ = global_flag_repr
1396 else:
1397 cls.__repr__ = global_enum_repr
1398 sys.modules[cls.__module__].__dict__.update(cls.__members__)
1399 return cls
Ethan Furmana02cb472021-04-21 10:20:44 -07001400
1401def _simple_enum(etype=Enum, *, boundary=None, use_args=None):
1402 """
1403 Class decorator that converts a normal class into an :class:`Enum`. No
1404 safety checks are done, and some advanced behavior (such as
1405 :func:`__init_subclass__`) is not available. Enum creation can be faster
1406 using :func:`simple_enum`.
1407
1408 >>> from enum import Enum, _simple_enum
1409 >>> @_simple_enum(Enum)
1410 ... class Color:
1411 ... RED = auto()
1412 ... GREEN = auto()
1413 ... BLUE = auto()
1414 >>> Color
1415 <enum 'Color'>
1416 """
1417 def convert_class(cls):
1418 nonlocal use_args
1419 cls_name = cls.__name__
1420 if use_args is None:
1421 use_args = etype._use_args_
1422 __new__ = cls.__dict__.get('__new__')
1423 if __new__ is not None:
1424 new_member = __new__.__func__
1425 else:
1426 new_member = etype._member_type_.__new__
1427 attrs = {}
1428 body = {}
1429 if __new__ is not None:
1430 body['__new_member__'] = new_member
1431 body['_new_member_'] = new_member
1432 body['_use_args_'] = use_args
1433 body['_generate_next_value_'] = gnv = etype._generate_next_value_
1434 body['_member_names_'] = member_names = []
1435 body['_member_map_'] = member_map = {}
1436 body['_value2member_map_'] = value2member_map = {}
1437 body['_member_type_'] = member_type = etype._member_type_
1438 if issubclass(etype, Flag):
1439 body['_boundary_'] = boundary or etype._boundary_
1440 body['_flag_mask_'] = None
1441 body['_all_bits_'] = None
1442 body['_inverted_'] = None
1443 for name, obj in cls.__dict__.items():
1444 if name in ('__dict__', '__weakref__'):
1445 continue
1446 if _is_dunder(name) or _is_private(cls_name, name) or _is_sunder(name) or _is_descriptor(obj):
1447 body[name] = obj
1448 else:
1449 attrs[name] = obj
1450 if cls.__dict__.get('__doc__') is None:
1451 body['__doc__'] = 'An enumeration.'
1452 #
1453 # double check that repr and friends are not the mixin's or various
1454 # things break (such as pickle)
1455 # however, if the method is defined in the Enum itself, don't replace
1456 # it
1457 enum_class = type(cls_name, (etype, ), body, boundary=boundary, _simple=True)
1458 for name in ('__repr__', '__str__', '__format__', '__reduce_ex__'):
1459 if name in body:
1460 continue
1461 class_method = getattr(enum_class, name)
1462 obj_method = getattr(member_type, name, None)
1463 enum_method = getattr(etype, name, None)
1464 if obj_method is not None and obj_method is class_method:
1465 setattr(enum_class, name, enum_method)
1466 gnv_last_values = []
1467 if issubclass(enum_class, Flag):
1468 # Flag / IntFlag
1469 single_bits = multi_bits = 0
1470 for name, value in attrs.items():
1471 if isinstance(value, auto) and auto.value is _auto_null:
1472 value = gnv(name, 1, len(member_names), gnv_last_values)
1473 if value in value2member_map:
1474 # an alias to an existing member
1475 redirect = property()
1476 redirect.__set_name__(enum_class, name)
1477 setattr(enum_class, name, redirect)
1478 member_map[name] = value2member_map[value]
1479 else:
1480 # create the member
1481 if use_args:
1482 if not isinstance(value, tuple):
1483 value = (value, )
1484 member = new_member(enum_class, *value)
1485 value = value[0]
1486 else:
1487 member = new_member(enum_class)
1488 if __new__ is None:
1489 member._value_ = value
1490 member._name_ = name
1491 member.__objclass__ = enum_class
1492 member.__init__(value)
1493 redirect = property()
1494 redirect.__set_name__(enum_class, name)
1495 setattr(enum_class, name, redirect)
1496 member_map[name] = member
1497 member._sort_order_ = len(member_names)
1498 value2member_map[value] = member
1499 if _is_single_bit(value):
1500 # not a multi-bit alias, record in _member_names_ and _flag_mask_
1501 member_names.append(name)
1502 single_bits |= value
1503 else:
1504 multi_bits |= value
1505 gnv_last_values.append(value)
1506 enum_class._flag_mask_ = single_bits
1507 enum_class._all_bits_ = 2 ** ((single_bits|multi_bits).bit_length()) - 1
1508 # set correct __iter__
1509 member_list = [m._value_ for m in enum_class]
1510 if member_list != sorted(member_list):
1511 enum_class._iter_member_ = enum_class._iter_member_by_def_
1512 else:
1513 # Enum / IntEnum / StrEnum
1514 for name, value in attrs.items():
1515 if isinstance(value, auto):
1516 if value.value is _auto_null:
1517 value.value = gnv(name, 1, len(member_names), gnv_last_values)
1518 value = value.value
1519 if value in value2member_map:
1520 # an alias to an existing member
1521 redirect = property()
1522 redirect.__set_name__(enum_class, name)
1523 setattr(enum_class, name, redirect)
1524 member_map[name] = value2member_map[value]
1525 else:
1526 # create the member
1527 if use_args:
1528 if not isinstance(value, tuple):
1529 value = (value, )
1530 member = new_member(enum_class, *value)
1531 value = value[0]
1532 else:
1533 member = new_member(enum_class)
1534 if __new__ is None:
1535 member._value_ = value
1536 member._name_ = name
1537 member.__objclass__ = enum_class
1538 member.__init__(value)
1539 member._sort_order_ = len(member_names)
1540 redirect = property()
1541 redirect.__set_name__(enum_class, name)
1542 setattr(enum_class, name, redirect)
1543 member_map[name] = member
1544 value2member_map[value] = member
1545 member_names.append(name)
1546 gnv_last_values.append(value)
1547 if '__new__' in body:
1548 enum_class.__new_member__ = enum_class.__new__
1549 enum_class.__new__ = Enum.__new__
1550 return enum_class
1551 return convert_class
1552
1553def _test_simple_enum(checked_enum, simple_enum):
1554 """
1555 A function that can be used to test an enum created with :func:`_simple_enum`
1556 against the version created by subclassing :class:`Enum`::
1557
1558 >>> from enum import Enum, _simple_enum, _test_simple_enum
1559 >>> @_simple_enum(Enum)
1560 ... class Color:
1561 ... RED = auto()
1562 ... GREEN = auto()
1563 ... BLUE = auto()
1564 >>> class CheckedColor(Enum):
1565 ... RED = auto()
1566 ... GREEN = auto()
1567 ... BLUE = auto()
1568 >>> _test_simple_enum(CheckedColor, Color)
1569
1570 If differences are found, a :exc:`TypeError` is raised.
1571 """
1572 failed = []
1573 if checked_enum.__dict__ != simple_enum.__dict__:
1574 checked_dict = checked_enum.__dict__
1575 checked_keys = list(checked_dict.keys())
1576 simple_dict = simple_enum.__dict__
1577 simple_keys = list(simple_dict.keys())
1578 member_names = set(
1579 list(checked_enum._member_map_.keys())
1580 + list(simple_enum._member_map_.keys())
1581 )
1582 for key in set(checked_keys + simple_keys):
1583 if key in ('__module__', '_member_map_', '_value2member_map_'):
1584 # keys known to be different
1585 continue
1586 elif key in member_names:
1587 # members are checked below
1588 continue
1589 elif key not in simple_keys:
1590 failed.append("missing key: %r" % (key, ))
1591 elif key not in checked_keys:
1592 failed.append("extra key: %r" % (key, ))
1593 else:
1594 checked_value = checked_dict[key]
1595 simple_value = simple_dict[key]
1596 if callable(checked_value):
1597 continue
1598 if key == '__doc__':
1599 # remove all spaces/tabs
1600 compressed_checked_value = checked_value.replace(' ','').replace('\t','')
1601 compressed_simple_value = simple_value.replace(' ','').replace('\t','')
1602 if compressed_checked_value != compressed_simple_value:
1603 failed.append("%r:\n %s\n %s" % (
1604 key,
1605 "checked -> %r" % (checked_value, ),
1606 "simple -> %r" % (simple_value, ),
1607 ))
1608 elif checked_value != simple_value:
1609 failed.append("%r:\n %s\n %s" % (
1610 key,
1611 "checked -> %r" % (checked_value, ),
1612 "simple -> %r" % (simple_value, ),
1613 ))
1614 failed.sort()
1615 for name in member_names:
1616 failed_member = []
1617 if name not in simple_keys:
1618 failed.append('missing member from simple enum: %r' % name)
1619 elif name not in checked_keys:
1620 failed.append('extra member in simple enum: %r' % name)
1621 else:
1622 checked_member_dict = checked_enum[name].__dict__
1623 checked_member_keys = list(checked_member_dict.keys())
1624 simple_member_dict = simple_enum[name].__dict__
1625 simple_member_keys = list(simple_member_dict.keys())
1626 for key in set(checked_member_keys + simple_member_keys):
Ethan Furman6c681e12021-04-23 19:08:22 -07001627 if key in ('__module__', '__objclass__', '_inverted_'):
1628 # keys known to be different or absent
Ethan Furmana02cb472021-04-21 10:20:44 -07001629 continue
1630 elif key not in simple_member_keys:
1631 failed_member.append("missing key %r not in the simple enum member %r" % (key, name))
1632 elif key not in checked_member_keys:
1633 failed_member.append("extra key %r in simple enum member %r" % (key, name))
1634 else:
1635 checked_value = checked_member_dict[key]
1636 simple_value = simple_member_dict[key]
1637 if checked_value != simple_value:
1638 failed_member.append("%r:\n %s\n %s" % (
1639 key,
1640 "checked member -> %r" % (checked_value, ),
1641 "simple member -> %r" % (simple_value, ),
1642 ))
1643 if failed_member:
1644 failed.append('%r member mismatch:\n %s' % (
1645 name, '\n '.join(failed_member),
1646 ))
1647 for method in (
1648 '__str__', '__repr__', '__reduce_ex__', '__format__',
1649 '__getnewargs_ex__', '__getnewargs__', '__reduce_ex__', '__reduce__'
1650 ):
1651 if method in simple_keys and method in checked_keys:
1652 # cannot compare functions, and it exists in both, so we're good
1653 continue
1654 elif method not in simple_keys and method not in checked_keys:
1655 # method is inherited -- check it out
1656 checked_method = getattr(checked_enum, method, None)
1657 simple_method = getattr(simple_enum, method, None)
1658 if hasattr(checked_method, '__func__'):
1659 checked_method = checked_method.__func__
1660 simple_method = simple_method.__func__
1661 if checked_method != simple_method:
1662 failed.append("%r: %-30s %s" % (
1663 method,
1664 "checked -> %r" % (checked_method, ),
1665 "simple -> %r" % (simple_method, ),
1666 ))
1667 else:
1668 # if the method existed in only one of the enums, it will have been caught
1669 # in the first checks above
1670 pass
1671 if failed:
1672 raise TypeError('enum mismatch:\n %s' % '\n '.join(failed))
1673
1674def _old_convert_(etype, name, module, filter, source=None, *, boundary=None):
1675 """
1676 Create a new Enum subclass that replaces a collection of global constants
1677 """
1678 # convert all constants from source (or module) that pass filter() to
1679 # a new Enum called name, and export the enum and its members back to
1680 # module;
1681 # also, replace the __reduce_ex__ method so unpickling works in
1682 # previous Python versions
1683 module_globals = sys.modules[module].__dict__
1684 if source:
1685 source = source.__dict__
1686 else:
1687 source = module_globals
1688 # _value2member_map_ is populated in the same order every time
1689 # for a consistent reverse mapping of number to name when there
1690 # are multiple names for the same number.
1691 members = [
1692 (name, value)
1693 for name, value in source.items()
1694 if filter(name)]
1695 try:
1696 # sort by value
1697 members.sort(key=lambda t: (t[1], t[0]))
1698 except TypeError:
1699 # unless some values aren't comparable, in which case sort by name
1700 members.sort(key=lambda t: t[0])
1701 cls = etype(name, members, module=module, boundary=boundary or KEEP)
1702 cls.__reduce_ex__ = _reduce_ex_by_name
1703 cls.__repr__ = global_enum_repr
1704 return cls