blob: 82be1fbaf70b0a94015fd1c34b5fc5329e6a750d [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:
1008 cls = self._member_type_
Ethan Furman0081f232014-09-16 17:31:23 -07001009 val = self._value_
Ethan Furmanec15a822013-08-31 19:17:41 -07001010 return cls.__format__(val, format_spec)
1011
Ethan Furman6b3d64a2013-06-14 16:55:46 -07001012 def __hash__(self):
Ethan Furman520ad572013-07-19 19:47:21 -07001013 return hash(self._name_)
Ethan Furman6b3d64a2013-06-14 16:55:46 -07001014
Ethan Furmanca1b7942014-02-08 11:36:27 -08001015 def __reduce_ex__(self, proto):
Ethan Furmandc870522014-02-18 12:37:12 -08001016 return self.__class__, (self._value_, )
Ethan Furmanca1b7942014-02-08 11:36:27 -08001017
Ethan Furmanc314e602021-01-12 23:47:57 -08001018 # enum.property is used to provide access to the `name` and
1019 # `value` attributes of enum members while keeping some measure of
Ethan Furman6b3d64a2013-06-14 16:55:46 -07001020 # protection from modification, while still allowing for an enumeration
1021 # to have members named `name` and `value`. This works because enumeration
Ethan Furmanc314e602021-01-12 23:47:57 -08001022 # members are not set directly on the enum class; they are kept in a
1023 # separate structure, _member_map_, which is where enum.property looks for
1024 # them
Ethan Furman6b3d64a2013-06-14 16:55:46 -07001025
Ethan Furmanc314e602021-01-12 23:47:57 -08001026 @property
Ethan Furman6b3d64a2013-06-14 16:55:46 -07001027 def name(self):
Ethan Furmanc850f342013-09-15 16:59:35 -07001028 """The name of the Enum member."""
Ethan Furman520ad572013-07-19 19:47:21 -07001029 return self._name_
Ethan Furman6b3d64a2013-06-14 16:55:46 -07001030
Ethan Furmanc314e602021-01-12 23:47:57 -08001031 @property
Ethan Furman6b3d64a2013-06-14 16:55:46 -07001032 def value(self):
Ethan Furmanc850f342013-09-15 16:59:35 -07001033 """The value of the Enum member."""
Ethan Furman520ad572013-07-19 19:47:21 -07001034 return self._value_
Ethan Furman6b3d64a2013-06-14 16:55:46 -07001035
1036
1037class IntEnum(int, Enum):
Ethan Furman0063ff42020-09-21 17:23:13 -07001038 """
1039 Enum where members are also (and must be) ints
1040 """
1041
1042
1043class StrEnum(str, Enum):
1044 """
1045 Enum where members are also (and must be) strings
1046 """
1047
1048 def __new__(cls, *values):
1049 if len(values) > 3:
1050 raise TypeError('too many arguments for str(): %r' % (values, ))
1051 if len(values) == 1:
1052 # it must be a string
1053 if not isinstance(values[0], str):
1054 raise TypeError('%r is not a string' % (values[0], ))
Ethan Furman7aaeb2a2021-01-25 14:26:19 -08001055 if len(values) >= 2:
Ethan Furman0063ff42020-09-21 17:23:13 -07001056 # check that encoding argument is a string
1057 if not isinstance(values[1], str):
1058 raise TypeError('encoding must be a string, not %r' % (values[1], ))
Ethan Furman7aaeb2a2021-01-25 14:26:19 -08001059 if len(values) == 3:
1060 # check that errors argument is a string
1061 if not isinstance(values[2], str):
1062 raise TypeError('errors must be a string, not %r' % (values[2]))
Ethan Furman0063ff42020-09-21 17:23:13 -07001063 value = str(*values)
1064 member = str.__new__(cls, value)
1065 member._value_ = value
1066 return member
Ethan Furmanf24bb352013-07-18 17:05:39 -07001067
Ethan Furmand986d162020-09-22 13:00:07 -07001068 __str__ = str.__str__
1069
Ethan Furmanefb13be2020-12-10 12:20:06 -08001070 def _generate_next_value_(name, start, count, last_values):
1071 """
1072 Return the lower-cased version of the member name.
1073 """
1074 return name.lower()
1075
Ethan Furmanf24bb352013-07-18 17:05:39 -07001076
Ethan Furman24e837f2015-03-18 17:27:57 -07001077def _reduce_ex_by_name(self, proto):
1078 return self.name
1079
Ethan Furman7aaeb2a2021-01-25 14:26:19 -08001080class FlagBoundary(StrEnum):
1081 """
1082 control how out of range values are handled
1083 "strict" -> error is raised [default for Flag]
1084 "conform" -> extra bits are discarded
1085 "eject" -> lose flag status [default for IntFlag]
1086 "keep" -> keep flag status and all bits
1087 """
1088 STRICT = auto()
1089 CONFORM = auto()
1090 EJECT = auto()
1091 KEEP = auto()
1092STRICT, CONFORM, EJECT, KEEP = FlagBoundary
1093
1094
1095class Flag(Enum, boundary=STRICT):
Ethan Furman6d3dfee2020-12-08 12:26:56 -08001096 """
1097 Support for flags
1098 """
Ethan Furmanc16595e2016-09-10 23:36:59 -07001099
1100 def _generate_next_value_(name, start, count, last_values):
Ethan Furmanee47e5c2016-08-31 00:12:15 -07001101 """
1102 Generate the next value when not given.
1103
1104 name: the name of the member
HongWeipengbb16fb22019-09-21 13:22:54 +08001105 start: the initial start value or None
Ethan Furmanee47e5c2016-08-31 00:12:15 -07001106 count: the number of existing members
1107 last_value: the last value assigned or None
1108 """
1109 if not count:
1110 return start if start is not None else 1
Ethan Furman7aaeb2a2021-01-25 14:26:19 -08001111 last_value = max(last_values)
1112 try:
1113 high_bit = _high_bit(last_value)
1114 except Exception:
1115 raise TypeError('Invalid Flag value: %r' % last_value) from None
Ethan Furmanee47e5c2016-08-31 00:12:15 -07001116 return 2 ** (high_bit+1)
1117
1118 @classmethod
Ethan Furman7aaeb2a2021-01-25 14:26:19 -08001119 def _iter_member_by_value_(cls, value):
1120 """
1121 Extract all members from the value in definition (i.e. increasing value) order.
1122 """
1123 for val in _iter_bits_lsb(value & cls._flag_mask_):
1124 yield cls._value2member_map_.get(val)
1125
1126 _iter_member_ = _iter_member_by_value_
1127
1128 @classmethod
1129 def _iter_member_by_def_(cls, value):
1130 """
1131 Extract all members from the value in definition order.
1132 """
1133 yield from sorted(
1134 cls._iter_member_by_value_(value),
1135 key=lambda m: m._sort_order_,
1136 )
1137
1138 @classmethod
Ethan Furmanee47e5c2016-08-31 00:12:15 -07001139 def _missing_(cls, value):
Ethan Furman6d3dfee2020-12-08 12:26:56 -08001140 """
Ethan Furman0dca5eb2021-04-15 06:49:54 -07001141 Create a composite member containing all canonical members present in `value`.
1142
1143 If non-member values are present, result depends on `_boundary_` setting.
Ethan Furman3515dcc2016-09-18 13:15:41 -07001144 """
Ethan Furman7aaeb2a2021-01-25 14:26:19 -08001145 if not isinstance(value, int):
1146 raise ValueError(
1147 "%r is not a valid %s" % (value, cls.__qualname__)
1148 )
1149 # check boundaries
1150 # - value must be in range (e.g. -16 <-> +15, i.e. ~15 <-> 15)
1151 # - value must not include any skipped flags (e.g. if bit 2 is not
1152 # defined, then 0d10 is invalid)
1153 flag_mask = cls._flag_mask_
1154 all_bits = cls._all_bits_
1155 neg_value = None
1156 if (
1157 not ~all_bits <= value <= all_bits
1158 or value & (all_bits ^ flag_mask)
1159 ):
1160 if cls._boundary_ is STRICT:
1161 max_bits = max(value.bit_length(), flag_mask.bit_length())
1162 raise ValueError(
1163 "%s: invalid value: %r\n given %s\n allowed %s" % (
1164 cls.__name__, value, bin(value, max_bits), bin(flag_mask, max_bits),
1165 ))
1166 elif cls._boundary_ is CONFORM:
1167 value = value & flag_mask
1168 elif cls._boundary_ is EJECT:
1169 return value
1170 elif cls._boundary_ is KEEP:
1171 if value < 0:
1172 value = (
1173 max(all_bits+1, 2**(value.bit_length()))
1174 + value
1175 )
1176 else:
1177 raise ValueError(
1178 'unknown flag boundary: %r' % (cls._boundary_, )
1179 )
1180 if value < 0:
1181 neg_value = value
1182 value = all_bits + 1 + value
1183 # get members and unknown
1184 unknown = value & ~flag_mask
1185 member_value = value & flag_mask
1186 if unknown and cls._boundary_ is not KEEP:
1187 raise ValueError(
1188 '%s(%r) --> unknown values %r [%s]'
1189 % (cls.__name__, value, unknown, bin(unknown))
1190 )
1191 # normal Flag?
1192 __new__ = getattr(cls, '__new_member__', None)
1193 if cls._member_type_ is object and not __new__:
Ethan Furman3515dcc2016-09-18 13:15:41 -07001194 # construct a singleton enum pseudo-member
Ethan Furmanee47e5c2016-08-31 00:12:15 -07001195 pseudo_member = object.__new__(cls)
Ethan Furman7aaeb2a2021-01-25 14:26:19 -08001196 else:
1197 pseudo_member = (__new__ or cls._member_type_.__new__)(cls, value)
Ethan Furmana02cb472021-04-21 10:20:44 -07001198 if not hasattr(pseudo_member, '_value_'):
Ethan Furmanee47e5c2016-08-31 00:12:15 -07001199 pseudo_member._value_ = value
Ethan Furman7aaeb2a2021-01-25 14:26:19 -08001200 if member_value:
1201 pseudo_member._name_ = '|'.join([
1202 m._name_ for m in cls._iter_member_(member_value)
1203 ])
1204 if unknown:
1205 pseudo_member._name_ += '|0x%x' % unknown
1206 else:
1207 pseudo_member._name_ = None
1208 # use setdefault in case another thread already created a composite
1209 # with this value, but only if all members are known
1210 # note: zero is a special case -- add it
1211 if not unknown:
Ethan Furman28cf6632017-01-24 12:12:06 -08001212 pseudo_member = cls._value2member_map_.setdefault(value, pseudo_member)
Ethan Furman7aaeb2a2021-01-25 14:26:19 -08001213 if neg_value is not None:
1214 cls._value2member_map_[neg_value] = pseudo_member
Ethan Furmanee47e5c2016-08-31 00:12:15 -07001215 return pseudo_member
1216
Ethan Furmanee47e5c2016-08-31 00:12:15 -07001217 def __contains__(self, other):
Ethan Furman6d3dfee2020-12-08 12:26:56 -08001218 """
1219 Returns True if self has at least the same flags set as other.
1220 """
Ethan Furmanee47e5c2016-08-31 00:12:15 -07001221 if not isinstance(other, self.__class__):
Rahul Jha94306522018-09-10 23:51:04 +05301222 raise TypeError(
1223 "unsupported operand type(s) for 'in': '%s' and '%s'" % (
1224 type(other).__qualname__, self.__class__.__qualname__))
Ethan Furman7aaeb2a2021-01-25 14:26:19 -08001225 if other._value_ == 0 or self._value_ == 0:
1226 return False
Ethan Furmanee47e5c2016-08-31 00:12:15 -07001227 return other._value_ & self._value_ == other._value_
1228
Ethan Furman7219e272020-09-16 13:01:00 -07001229 def __iter__(self):
Ethan Furman6d3dfee2020-12-08 12:26:56 -08001230 """
Ethan Furman7aaeb2a2021-01-25 14:26:19 -08001231 Returns flags in definition order.
Ethan Furman6d3dfee2020-12-08 12:26:56 -08001232 """
Ethan Furman7aaeb2a2021-01-25 14:26:19 -08001233 yield from self._iter_member_(self._value_)
1234
1235 def __len__(self):
1236 return self._value_.bit_count()
Ethan Furman7219e272020-09-16 13:01:00 -07001237
Ethan Furmanee47e5c2016-08-31 00:12:15 -07001238 def __repr__(self):
Ethan Furmanb7751062021-03-30 21:17:26 -07001239 cls_name = self.__class__.__name__
1240 if self._name_ is None:
1241 return "0x%x" % (self._value_, )
1242 if _is_single_bit(self._value_):
1243 return '%s.%s' % (cls_name, self._name_)
1244 if self._boundary_ is not FlagBoundary.KEEP:
1245 return '%s.' % cls_name + ('|%s.' % cls_name).join(self.name.split('|'))
Ethan Furman7aaeb2a2021-01-25 14:26:19 -08001246 else:
Ethan Furmanb7751062021-03-30 21:17:26 -07001247 name = []
1248 for n in self._name_.split('|'):
1249 if n.startswith('0'):
1250 name.append(n)
1251 else:
1252 name.append('%s.%s' % (cls_name, n))
1253 return '|'.join(name)
Ethan Furmanee47e5c2016-08-31 00:12:15 -07001254
1255 def __str__(self):
1256 cls = self.__class__
Ethan Furmanb7751062021-03-30 21:17:26 -07001257 if self._name_ is None:
1258 return '%s(%x)' % (cls.__name__, self._value_)
Ethan Furmanee47e5c2016-08-31 00:12:15 -07001259 else:
Ethan Furmanb7751062021-03-30 21:17:26 -07001260 return self._name_
Ethan Furmanee47e5c2016-08-31 00:12:15 -07001261
Ethan Furman25d94bb2016-09-02 16:32:32 -07001262 def __bool__(self):
1263 return bool(self._value_)
1264
Ethan Furmanee47e5c2016-08-31 00:12:15 -07001265 def __or__(self, other):
1266 if not isinstance(other, self.__class__):
1267 return NotImplemented
1268 return self.__class__(self._value_ | other._value_)
1269
1270 def __and__(self, other):
1271 if not isinstance(other, self.__class__):
1272 return NotImplemented
1273 return self.__class__(self._value_ & other._value_)
1274
1275 def __xor__(self, other):
1276 if not isinstance(other, self.__class__):
1277 return NotImplemented
1278 return self.__class__(self._value_ ^ other._value_)
1279
1280 def __invert__(self):
Ethan Furman7aaeb2a2021-01-25 14:26:19 -08001281 if self._inverted_ is None:
1282 if self._boundary_ is KEEP:
1283 # use all bits
1284 self._inverted_ = self.__class__(~self._value_)
1285 else:
1286 # calculate flags not in this member
1287 self._inverted_ = self.__class__(self._flag_mask_ ^ self._value_)
1288 self._inverted_._inverted_ = self
1289 return self._inverted_
Ethan Furmanee47e5c2016-08-31 00:12:15 -07001290
1291
Ethan Furman7aaeb2a2021-01-25 14:26:19 -08001292class IntFlag(int, Flag, boundary=EJECT):
Ethan Furman6d3dfee2020-12-08 12:26:56 -08001293 """
1294 Support for integer-based Flags
1295 """
Ethan Furmanee47e5c2016-08-31 00:12:15 -07001296
Ethan Furmanee47e5c2016-08-31 00:12:15 -07001297 def __or__(self, other):
Ethan Furman7aaeb2a2021-01-25 14:26:19 -08001298 if isinstance(other, self.__class__):
1299 other = other._value_
1300 elif isinstance(other, int):
1301 other = other
1302 else:
Ethan Furmanee47e5c2016-08-31 00:12:15 -07001303 return NotImplemented
Ethan Furman7aaeb2a2021-01-25 14:26:19 -08001304 value = self._value_
1305 return self.__class__(value | other)
Ethan Furmanee47e5c2016-08-31 00:12:15 -07001306
1307 def __and__(self, other):
Ethan Furman7aaeb2a2021-01-25 14:26:19 -08001308 if isinstance(other, self.__class__):
1309 other = other._value_
1310 elif isinstance(other, int):
1311 other = other
1312 else:
Ethan Furmanee47e5c2016-08-31 00:12:15 -07001313 return NotImplemented
Ethan Furman7aaeb2a2021-01-25 14:26:19 -08001314 value = self._value_
1315 return self.__class__(value & other)
Ethan Furmanee47e5c2016-08-31 00:12:15 -07001316
1317 def __xor__(self, other):
Ethan Furman7aaeb2a2021-01-25 14:26:19 -08001318 if isinstance(other, self.__class__):
1319 other = other._value_
1320 elif isinstance(other, int):
1321 other = other
1322 else:
Ethan Furmanee47e5c2016-08-31 00:12:15 -07001323 return NotImplemented
Ethan Furman7aaeb2a2021-01-25 14:26:19 -08001324 value = self._value_
1325 return self.__class__(value ^ other)
Ethan Furmanee47e5c2016-08-31 00:12:15 -07001326
1327 __ror__ = __or__
1328 __rand__ = __and__
1329 __rxor__ = __xor__
Ethan Furman7aaeb2a2021-01-25 14:26:19 -08001330 __invert__ = Flag.__invert__
Ethan Furmanee47e5c2016-08-31 00:12:15 -07001331
1332def _high_bit(value):
Ethan Furman6d3dfee2020-12-08 12:26:56 -08001333 """
1334 returns index of highest bit, or -1 if value is zero or negative
1335 """
Ethan Furman3515dcc2016-09-18 13:15:41 -07001336 return value.bit_length() - 1
Ethan Furmanee47e5c2016-08-31 00:12:15 -07001337
Ethan Furmanf24bb352013-07-18 17:05:39 -07001338def unique(enumeration):
Ethan Furman6d3dfee2020-12-08 12:26:56 -08001339 """
1340 Class decorator for enumerations ensuring unique member values.
1341 """
Ethan Furmanf24bb352013-07-18 17:05:39 -07001342 duplicates = []
1343 for name, member in enumeration.__members__.items():
1344 if name != member.name:
1345 duplicates.append((name, member.name))
1346 if duplicates:
1347 alias_details = ', '.join(
1348 ["%s -> %s" % (alias, name) for (alias, name) in duplicates])
1349 raise ValueError('duplicate values found in %r: %s' %
1350 (enumeration, alias_details))
1351 return enumeration
Ethan Furman3515dcc2016-09-18 13:15:41 -07001352
Ethan Furman7aaeb2a2021-01-25 14:26:19 -08001353def _power_of_two(value):
1354 if value < 1:
1355 return False
1356 return value == 2 ** _high_bit(value)
Ethan Furmanb7751062021-03-30 21:17:26 -07001357
1358def global_enum_repr(self):
1359 return '%s.%s' % (self.__class__.__module__, self._name_)
1360
1361def global_flag_repr(self):
1362 module = self.__class__.__module__
1363 cls_name = self.__class__.__name__
1364 if self._name_ is None:
1365 return "%x" % (module, cls_name, self._value_)
1366 if _is_single_bit(self):
1367 return '%s.%s' % (module, self._name_)
1368 if self._boundary_ is not FlagBoundary.KEEP:
1369 return module + module.join(self.name.split('|'))
1370 else:
1371 name = []
1372 for n in self._name_.split('|'):
1373 if n.startswith('0'):
1374 name.append(n)
1375 else:
1376 name.append('%s.%s' % (module, n))
1377 return '|'.join(name)
1378
1379
1380def global_enum(cls):
1381 """
1382 decorator that makes the repr() of an enum member reference its module
1383 instead of its class; also exports all members to the enum's module's
1384 global namespace
1385 """
1386 if issubclass(cls, Flag):
1387 cls.__repr__ = global_flag_repr
1388 else:
1389 cls.__repr__ = global_enum_repr
1390 sys.modules[cls.__module__].__dict__.update(cls.__members__)
1391 return cls
Ethan Furmana02cb472021-04-21 10:20:44 -07001392
1393def _simple_enum(etype=Enum, *, boundary=None, use_args=None):
1394 """
1395 Class decorator that converts a normal class into an :class:`Enum`. No
1396 safety checks are done, and some advanced behavior (such as
1397 :func:`__init_subclass__`) is not available. Enum creation can be faster
1398 using :func:`simple_enum`.
1399
1400 >>> from enum import Enum, _simple_enum
1401 >>> @_simple_enum(Enum)
1402 ... class Color:
1403 ... RED = auto()
1404 ... GREEN = auto()
1405 ... BLUE = auto()
1406 >>> Color
1407 <enum 'Color'>
1408 """
1409 def convert_class(cls):
1410 nonlocal use_args
1411 cls_name = cls.__name__
1412 if use_args is None:
1413 use_args = etype._use_args_
1414 __new__ = cls.__dict__.get('__new__')
1415 if __new__ is not None:
1416 new_member = __new__.__func__
1417 else:
1418 new_member = etype._member_type_.__new__
1419 attrs = {}
1420 body = {}
1421 if __new__ is not None:
1422 body['__new_member__'] = new_member
1423 body['_new_member_'] = new_member
1424 body['_use_args_'] = use_args
1425 body['_generate_next_value_'] = gnv = etype._generate_next_value_
1426 body['_member_names_'] = member_names = []
1427 body['_member_map_'] = member_map = {}
1428 body['_value2member_map_'] = value2member_map = {}
1429 body['_member_type_'] = member_type = etype._member_type_
1430 if issubclass(etype, Flag):
1431 body['_boundary_'] = boundary or etype._boundary_
1432 body['_flag_mask_'] = None
1433 body['_all_bits_'] = None
1434 body['_inverted_'] = None
1435 for name, obj in cls.__dict__.items():
1436 if name in ('__dict__', '__weakref__'):
1437 continue
1438 if _is_dunder(name) or _is_private(cls_name, name) or _is_sunder(name) or _is_descriptor(obj):
1439 body[name] = obj
1440 else:
1441 attrs[name] = obj
1442 if cls.__dict__.get('__doc__') is None:
1443 body['__doc__'] = 'An enumeration.'
1444 #
1445 # double check that repr and friends are not the mixin's or various
1446 # things break (such as pickle)
1447 # however, if the method is defined in the Enum itself, don't replace
1448 # it
1449 enum_class = type(cls_name, (etype, ), body, boundary=boundary, _simple=True)
1450 for name in ('__repr__', '__str__', '__format__', '__reduce_ex__'):
1451 if name in body:
1452 continue
1453 class_method = getattr(enum_class, name)
1454 obj_method = getattr(member_type, name, None)
1455 enum_method = getattr(etype, name, None)
1456 if obj_method is not None and obj_method is class_method:
1457 setattr(enum_class, name, enum_method)
1458 gnv_last_values = []
1459 if issubclass(enum_class, Flag):
1460 # Flag / IntFlag
1461 single_bits = multi_bits = 0
1462 for name, value in attrs.items():
1463 if isinstance(value, auto) and auto.value is _auto_null:
1464 value = gnv(name, 1, len(member_names), gnv_last_values)
1465 if value in value2member_map:
1466 # an alias to an existing member
1467 redirect = property()
1468 redirect.__set_name__(enum_class, name)
1469 setattr(enum_class, name, redirect)
1470 member_map[name] = value2member_map[value]
1471 else:
1472 # create the member
1473 if use_args:
1474 if not isinstance(value, tuple):
1475 value = (value, )
1476 member = new_member(enum_class, *value)
1477 value = value[0]
1478 else:
1479 member = new_member(enum_class)
1480 if __new__ is None:
1481 member._value_ = value
1482 member._name_ = name
1483 member.__objclass__ = enum_class
1484 member.__init__(value)
1485 redirect = property()
1486 redirect.__set_name__(enum_class, name)
1487 setattr(enum_class, name, redirect)
1488 member_map[name] = member
1489 member._sort_order_ = len(member_names)
1490 value2member_map[value] = member
1491 if _is_single_bit(value):
1492 # not a multi-bit alias, record in _member_names_ and _flag_mask_
1493 member_names.append(name)
1494 single_bits |= value
1495 else:
1496 multi_bits |= value
1497 gnv_last_values.append(value)
1498 enum_class._flag_mask_ = single_bits
1499 enum_class._all_bits_ = 2 ** ((single_bits|multi_bits).bit_length()) - 1
1500 # set correct __iter__
1501 member_list = [m._value_ for m in enum_class]
1502 if member_list != sorted(member_list):
1503 enum_class._iter_member_ = enum_class._iter_member_by_def_
1504 else:
1505 # Enum / IntEnum / StrEnum
1506 for name, value in attrs.items():
1507 if isinstance(value, auto):
1508 if value.value is _auto_null:
1509 value.value = gnv(name, 1, len(member_names), gnv_last_values)
1510 value = value.value
1511 if value in value2member_map:
1512 # an alias to an existing member
1513 redirect = property()
1514 redirect.__set_name__(enum_class, name)
1515 setattr(enum_class, name, redirect)
1516 member_map[name] = value2member_map[value]
1517 else:
1518 # create the member
1519 if use_args:
1520 if not isinstance(value, tuple):
1521 value = (value, )
1522 member = new_member(enum_class, *value)
1523 value = value[0]
1524 else:
1525 member = new_member(enum_class)
1526 if __new__ is None:
1527 member._value_ = value
1528 member._name_ = name
1529 member.__objclass__ = enum_class
1530 member.__init__(value)
1531 member._sort_order_ = len(member_names)
1532 redirect = property()
1533 redirect.__set_name__(enum_class, name)
1534 setattr(enum_class, name, redirect)
1535 member_map[name] = member
1536 value2member_map[value] = member
1537 member_names.append(name)
1538 gnv_last_values.append(value)
1539 if '__new__' in body:
1540 enum_class.__new_member__ = enum_class.__new__
1541 enum_class.__new__ = Enum.__new__
1542 return enum_class
1543 return convert_class
1544
1545def _test_simple_enum(checked_enum, simple_enum):
1546 """
1547 A function that can be used to test an enum created with :func:`_simple_enum`
1548 against the version created by subclassing :class:`Enum`::
1549
1550 >>> from enum import Enum, _simple_enum, _test_simple_enum
1551 >>> @_simple_enum(Enum)
1552 ... class Color:
1553 ... RED = auto()
1554 ... GREEN = auto()
1555 ... BLUE = auto()
1556 >>> class CheckedColor(Enum):
1557 ... RED = auto()
1558 ... GREEN = auto()
1559 ... BLUE = auto()
1560 >>> _test_simple_enum(CheckedColor, Color)
1561
1562 If differences are found, a :exc:`TypeError` is raised.
1563 """
1564 failed = []
1565 if checked_enum.__dict__ != simple_enum.__dict__:
1566 checked_dict = checked_enum.__dict__
1567 checked_keys = list(checked_dict.keys())
1568 simple_dict = simple_enum.__dict__
1569 simple_keys = list(simple_dict.keys())
1570 member_names = set(
1571 list(checked_enum._member_map_.keys())
1572 + list(simple_enum._member_map_.keys())
1573 )
1574 for key in set(checked_keys + simple_keys):
1575 if key in ('__module__', '_member_map_', '_value2member_map_'):
1576 # keys known to be different
1577 continue
1578 elif key in member_names:
1579 # members are checked below
1580 continue
1581 elif key not in simple_keys:
1582 failed.append("missing key: %r" % (key, ))
1583 elif key not in checked_keys:
1584 failed.append("extra key: %r" % (key, ))
1585 else:
1586 checked_value = checked_dict[key]
1587 simple_value = simple_dict[key]
1588 if callable(checked_value):
1589 continue
1590 if key == '__doc__':
1591 # remove all spaces/tabs
1592 compressed_checked_value = checked_value.replace(' ','').replace('\t','')
1593 compressed_simple_value = simple_value.replace(' ','').replace('\t','')
1594 if compressed_checked_value != compressed_simple_value:
1595 failed.append("%r:\n %s\n %s" % (
1596 key,
1597 "checked -> %r" % (checked_value, ),
1598 "simple -> %r" % (simple_value, ),
1599 ))
1600 elif checked_value != simple_value:
1601 failed.append("%r:\n %s\n %s" % (
1602 key,
1603 "checked -> %r" % (checked_value, ),
1604 "simple -> %r" % (simple_value, ),
1605 ))
1606 failed.sort()
1607 for name in member_names:
1608 failed_member = []
1609 if name not in simple_keys:
1610 failed.append('missing member from simple enum: %r' % name)
1611 elif name not in checked_keys:
1612 failed.append('extra member in simple enum: %r' % name)
1613 else:
1614 checked_member_dict = checked_enum[name].__dict__
1615 checked_member_keys = list(checked_member_dict.keys())
1616 simple_member_dict = simple_enum[name].__dict__
1617 simple_member_keys = list(simple_member_dict.keys())
1618 for key in set(checked_member_keys + simple_member_keys):
1619 if key in ('__module__', '__objclass__'):
1620 # keys known to be different
1621 continue
1622 elif key not in simple_member_keys:
1623 failed_member.append("missing key %r not in the simple enum member %r" % (key, name))
1624 elif key not in checked_member_keys:
1625 failed_member.append("extra key %r in simple enum member %r" % (key, name))
1626 else:
1627 checked_value = checked_member_dict[key]
1628 simple_value = simple_member_dict[key]
1629 if checked_value != simple_value:
1630 failed_member.append("%r:\n %s\n %s" % (
1631 key,
1632 "checked member -> %r" % (checked_value, ),
1633 "simple member -> %r" % (simple_value, ),
1634 ))
1635 if failed_member:
1636 failed.append('%r member mismatch:\n %s' % (
1637 name, '\n '.join(failed_member),
1638 ))
1639 for method in (
1640 '__str__', '__repr__', '__reduce_ex__', '__format__',
1641 '__getnewargs_ex__', '__getnewargs__', '__reduce_ex__', '__reduce__'
1642 ):
1643 if method in simple_keys and method in checked_keys:
1644 # cannot compare functions, and it exists in both, so we're good
1645 continue
1646 elif method not in simple_keys and method not in checked_keys:
1647 # method is inherited -- check it out
1648 checked_method = getattr(checked_enum, method, None)
1649 simple_method = getattr(simple_enum, method, None)
1650 if hasattr(checked_method, '__func__'):
1651 checked_method = checked_method.__func__
1652 simple_method = simple_method.__func__
1653 if checked_method != simple_method:
1654 failed.append("%r: %-30s %s" % (
1655 method,
1656 "checked -> %r" % (checked_method, ),
1657 "simple -> %r" % (simple_method, ),
1658 ))
1659 else:
1660 # if the method existed in only one of the enums, it will have been caught
1661 # in the first checks above
1662 pass
1663 if failed:
1664 raise TypeError('enum mismatch:\n %s' % '\n '.join(failed))
1665
1666def _old_convert_(etype, name, module, filter, source=None, *, boundary=None):
1667 """
1668 Create a new Enum subclass that replaces a collection of global constants
1669 """
1670 # convert all constants from source (or module) that pass filter() to
1671 # a new Enum called name, and export the enum and its members back to
1672 # module;
1673 # also, replace the __reduce_ex__ method so unpickling works in
1674 # previous Python versions
1675 module_globals = sys.modules[module].__dict__
1676 if source:
1677 source = source.__dict__
1678 else:
1679 source = module_globals
1680 # _value2member_map_ is populated in the same order every time
1681 # for a consistent reverse mapping of number to name when there
1682 # are multiple names for the same number.
1683 members = [
1684 (name, value)
1685 for name, value in source.items()
1686 if filter(name)]
1687 try:
1688 # sort by value
1689 members.sort(key=lambda t: (t[1], t[0]))
1690 except TypeError:
1691 # unless some values aren't comparable, in which case sort by name
1692 members.sort(key=lambda t: t[0])
1693 cls = etype(name, members, module=module, boundary=boundary or KEEP)
1694 cls.__reduce_ex__ = _reduce_ex_by_name
1695 cls.__repr__ = global_enum_repr
1696 return cls