blob: 5263e510d59361880a22919d7460eccbb898f7c5 [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 Furman74964862021-06-10 07:24:20 -07009 'auto', 'unique', 'property', 'verify',
Ethan Furman7aaeb2a2021-01-25 14:26:19 -080010 'FlagBoundary', 'STRICT', 'CONFORM', 'EJECT', 'KEEP',
Ethan Furmanb7751062021-03-30 21:17:26 -070011 'global_flag_repr', 'global_enum_repr', 'global_enum',
Ethan Furman74964862021-06-10 07:24:20 -070012 'EnumCheck', 'CONTINUOUS', 'NAMED_FLAGS', 'UNIQUE',
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):
Ethan Furman74964862021-06-10 07:24:20 -070092 # num must be an integer
93 if isinstance(num, Enum):
94 num = num.value
Ethan Furman7aaeb2a2021-01-25 14:26:19 -080095 while num:
96 b = num & (~num + 1)
97 yield b
98 num ^= b
99
100def bin(num, max_bits=None):
101 """
102 Like built-in bin(), except negative values are represented in
103 twos-compliment, and the leading bit always indicates sign
104 (0=positive, 1=negative).
105
106 >>> bin(10)
107 '0b0 1010'
108 >>> bin(~10) # ~10 is -11
109 '0b1 0101'
110 """
111
112 ceiling = 2 ** (num).bit_length()
113 if num >= 0:
114 s = _bltin_bin(num + ceiling).replace('1', '0', 1)
115 else:
116 s = _bltin_bin(~num ^ (ceiling - 1) + ceiling)
117 sign = s[:3]
118 digits = s[3:]
119 if max_bits is not None:
120 if len(digits) < max_bits:
121 digits = (sign[-1] * max_bits + digits)[-max_bits:]
122 return "%s %s" % (sign, digits)
123
124
Ethan Furman3515dcc2016-09-18 13:15:41 -0700125_auto_null = object()
Ethan Furmanc16595e2016-09-10 23:36:59 -0700126class auto:
127 """
128 Instances are replaced with an appropriate value in Enum class suites.
129 """
Ethan Furman3515dcc2016-09-18 13:15:41 -0700130 value = _auto_null
Ethan Furmanc16595e2016-09-10 23:36:59 -0700131
Ethan Furmanc314e602021-01-12 23:47:57 -0800132class property(DynamicClassAttribute):
133 """
134 This is a descriptor, used to define attributes that act differently
135 when accessed through an enum member and through an enum class.
136 Instance access is the same as property(), but access to an attribute
137 through the enum class will instead look in the class' _member_map_ for
138 a corresponding enum member.
139 """
140
141 def __get__(self, instance, ownerclass=None):
142 if instance is None:
143 try:
144 return ownerclass._member_map_[self.name]
145 except KeyError:
Ethan Furman7aaeb2a2021-01-25 14:26:19 -0800146 raise AttributeError(
Ethan Furmand65b9032021-02-08 17:32:38 -0800147 '%s: no class attribute %r' % (ownerclass.__name__, self.name)
Ethan Furman7aaeb2a2021-01-25 14:26:19 -0800148 )
Ethan Furmanc314e602021-01-12 23:47:57 -0800149 else:
150 if self.fget is None:
Ethan Furmand65b9032021-02-08 17:32:38 -0800151 # check for member
152 if self.name in ownerclass._member_map_:
153 import warnings
154 warnings.warn(
155 "accessing one member from another is not supported, "
Ethan Furman44e580f2021-03-03 09:54:30 -0800156 " and will be disabled in 3.12",
Ethan Furmand65b9032021-02-08 17:32:38 -0800157 DeprecationWarning,
158 stacklevel=2,
159 )
160 return ownerclass._member_map_[self.name]
Ethan Furman7aaeb2a2021-01-25 14:26:19 -0800161 raise AttributeError(
Ethan Furmand65b9032021-02-08 17:32:38 -0800162 '%s: no instance attribute %r' % (ownerclass.__name__, self.name)
Ethan Furman7aaeb2a2021-01-25 14:26:19 -0800163 )
Ethan Furmanc314e602021-01-12 23:47:57 -0800164 else:
165 return self.fget(instance)
166
167 def __set__(self, instance, value):
168 if self.fset is None:
Ethan Furman7aaeb2a2021-01-25 14:26:19 -0800169 raise AttributeError(
Ethan Furmand65b9032021-02-08 17:32:38 -0800170 "%s: cannot set instance attribute %r" % (self.clsname, self.name)
Ethan Furman7aaeb2a2021-01-25 14:26:19 -0800171 )
Ethan Furmanc314e602021-01-12 23:47:57 -0800172 else:
173 return self.fset(instance, value)
174
175 def __delete__(self, instance):
176 if self.fdel is None:
Ethan Furman7aaeb2a2021-01-25 14:26:19 -0800177 raise AttributeError(
Ethan Furmand65b9032021-02-08 17:32:38 -0800178 "%s: cannot delete instance attribute %r" % (self.clsname, self.name)
Ethan Furman7aaeb2a2021-01-25 14:26:19 -0800179 )
Ethan Furmanc314e602021-01-12 23:47:57 -0800180 else:
181 return self.fdel(instance)
182
183 def __set_name__(self, ownerclass, name):
184 self.name = name
185 self.clsname = ownerclass.__name__
186
187
188class _proto_member:
189 """
190 intermediate step for enum members between class execution and final creation
191 """
192
193 def __init__(self, value):
194 self.value = value
195
196 def __set_name__(self, enum_class, member_name):
197 """
198 convert each quasi-member into an instance of the new enum class
199 """
200 # first step: remove ourself from enum_class
201 delattr(enum_class, member_name)
202 # second step: create member based on enum_class
203 value = self.value
204 if not isinstance(value, tuple):
205 args = (value, )
206 else:
207 args = value
208 if enum_class._member_type_ is tuple: # special case for tuple enums
209 args = (args, ) # wrap it one more time
210 if not enum_class._use_args_:
211 enum_member = enum_class._new_member_(enum_class)
212 if not hasattr(enum_member, '_value_'):
213 enum_member._value_ = value
214 else:
215 enum_member = enum_class._new_member_(enum_class, *args)
216 if not hasattr(enum_member, '_value_'):
217 if enum_class._member_type_ is object:
218 enum_member._value_ = value
219 else:
Ethan Furman7aaeb2a2021-01-25 14:26:19 -0800220 try:
221 enum_member._value_ = enum_class._member_type_(*args)
222 except Exception as exc:
223 raise TypeError(
224 '_value_ not set in __new__, unable to create it'
225 ) from None
Ethan Furmanc314e602021-01-12 23:47:57 -0800226 value = enum_member._value_
227 enum_member._name_ = member_name
228 enum_member.__objclass__ = enum_class
229 enum_member.__init__(*args)
Ethan Furman7aaeb2a2021-01-25 14:26:19 -0800230 enum_member._sort_order_ = len(enum_class._member_names_)
Ethan Furmanc314e602021-01-12 23:47:57 -0800231 # If another member with the same value was already defined, the
232 # new member becomes an alias to the existing one.
233 for name, canonical_member in enum_class._member_map_.items():
234 if canonical_member._value_ == enum_member._value_:
235 enum_member = canonical_member
236 break
237 else:
Ethan Furman7aaeb2a2021-01-25 14:26:19 -0800238 # this could still be an alias if the value is multi-bit and the
239 # class is a flag class
240 if (
241 Flag is None
242 or not issubclass(enum_class, Flag)
243 ):
244 # no other instances found, record this member in _member_names_
245 enum_class._member_names_.append(member_name)
246 elif (
247 Flag is not None
248 and issubclass(enum_class, Flag)
249 and _is_single_bit(value)
250 ):
251 # no other instances found, record this member in _member_names_
252 enum_class._member_names_.append(member_name)
Ethan Furmanc314e602021-01-12 23:47:57 -0800253 # get redirect in place before adding to _member_map_
254 # but check for other instances in parent classes first
255 need_override = False
256 descriptor = None
257 for base in enum_class.__mro__[1:]:
258 descriptor = base.__dict__.get(member_name)
259 if descriptor is not None:
260 if isinstance(descriptor, (property, DynamicClassAttribute)):
261 break
262 else:
263 need_override = True
264 # keep looking for an enum.property
265 if descriptor and not need_override:
266 # previous enum.property found, no further action needed
267 pass
268 else:
269 redirect = property()
270 redirect.__set_name__(enum_class, member_name)
271 if descriptor and need_override:
272 # previous enum.property found, but some other inherited attribute
273 # is in the way; copy fget, fset, fdel to this one
274 redirect.fget = descriptor.fget
275 redirect.fset = descriptor.fset
276 redirect.fdel = descriptor.fdel
277 setattr(enum_class, member_name, redirect)
278 # now add to _member_map_ (even aliases)
279 enum_class._member_map_[member_name] = enum_member
280 try:
281 # This may fail if value is not hashable. We can't add the value
282 # to the map, and by-value lookups for this value will be
283 # linear.
Ethan Furman7aaeb2a2021-01-25 14:26:19 -0800284 enum_class._value2member_map_.setdefault(value, enum_member)
Ethan Furmanc314e602021-01-12 23:47:57 -0800285 except TypeError:
Ethan Furman6bd92882021-04-27 13:05:08 -0700286 # keep track of the value in a list so containment checks are quick
287 enum_class._unhashable_values_.append(value)
Ethan Furmanc314e602021-01-12 23:47:57 -0800288
Ethan Furman101e0742013-09-15 12:34:36 -0700289
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700290class _EnumDict(dict):
Ethan Furman6d3dfee2020-12-08 12:26:56 -0800291 """
292 Track enum member order and ensure member names are not reused.
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700293
Ethan Furmanb7751062021-03-30 21:17:26 -0700294 EnumType will use the names found in self._member_names as the
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700295 enumeration member names.
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700296 """
297 def __init__(self):
298 super().__init__()
299 self._member_names = []
Ethan Furmanc16595e2016-09-10 23:36:59 -0700300 self._last_values = []
Ethan Furmana4b1bb42018-01-22 07:56:37 -0800301 self._ignore = []
Ethan Onstottd9a43e22020-04-28 13:20:55 -0400302 self._auto_called = False
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700303
304 def __setitem__(self, key, value):
Ethan Furman6d3dfee2020-12-08 12:26:56 -0800305 """
306 Changes anything not dundered or not a descriptor.
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700307
308 If an enum member name is used twice, an error is raised; duplicate
309 values are not checked for.
310
311 Single underscore (sunder) names are reserved.
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700312 """
Ethan Furman7cf0aad2020-12-09 17:12:11 -0800313 if _is_private(self._cls_name, key):
314 # do nothing, name will be a normal attribute
315 pass
316 elif _is_sunder(key):
Ethan Furmanee47e5c2016-08-31 00:12:15 -0700317 if key not in (
Ethan Furman0dca5eb2021-04-15 06:49:54 -0700318 '_order_',
Ethan Furmana4b1bb42018-01-22 07:56:37 -0800319 '_generate_next_value_', '_missing_', '_ignore_',
Ethan Furman7aaeb2a2021-01-25 14:26:19 -0800320 '_iter_member_', '_iter_member_by_value_', '_iter_member_by_def_',
Ethan Furmanee47e5c2016-08-31 00:12:15 -0700321 ):
Ethan Furman6d3dfee2020-12-08 12:26:56 -0800322 raise ValueError(
323 '_sunder_ names, such as %r, are reserved for future Enum use'
324 % (key, )
325 )
Ethan Furmanc16595e2016-09-10 23:36:59 -0700326 if key == '_generate_next_value_':
Ethan Onstottd9a43e22020-04-28 13:20:55 -0400327 # check if members already defined as auto()
328 if self._auto_called:
329 raise TypeError("_generate_next_value_ must be defined before members")
Ethan Furmanb7751062021-03-30 21:17:26 -0700330 _gnv = value.__func__ if isinstance(value, staticmethod) else value
331 setattr(self, '_generate_next_value', _gnv)
Ethan Furmana4b1bb42018-01-22 07:56:37 -0800332 elif key == '_ignore_':
333 if isinstance(value, str):
334 value = value.replace(',',' ').split()
335 else:
336 value = list(value)
337 self._ignore = value
338 already = set(value) & set(self._member_names)
339 if already:
Ethan Furman6d3dfee2020-12-08 12:26:56 -0800340 raise ValueError(
341 '_ignore_ cannot specify already set names: %r'
342 % (already, )
343 )
Ethan Furman101e0742013-09-15 12:34:36 -0700344 elif _is_dunder(key):
Ethan Furmane8e61272016-08-20 07:19:31 -0700345 if key == '__order__':
346 key = '_order_'
Ethan Furman101e0742013-09-15 12:34:36 -0700347 elif key in self._member_names:
348 # descriptor overwriting an enum?
Ethan Furmana6582872020-12-10 13:07:00 -0800349 raise TypeError('%r already defined as: %r' % (key, self[key]))
Ethan Furmana4b1bb42018-01-22 07:56:37 -0800350 elif key in self._ignore:
351 pass
Ethan Furman101e0742013-09-15 12:34:36 -0700352 elif not _is_descriptor(value):
353 if key in self:
354 # enum overwriting a descriptor?
Ethan Furmanee47e5c2016-08-31 00:12:15 -0700355 raise TypeError('%r already defined as: %r' % (key, self[key]))
Ethan Furmanc16595e2016-09-10 23:36:59 -0700356 if isinstance(value, auto):
Ethan Furman3515dcc2016-09-18 13:15:41 -0700357 if value.value == _auto_null:
Ethan Furman6d3dfee2020-12-08 12:26:56 -0800358 value.value = self._generate_next_value(
Ethan Furman7aaeb2a2021-01-25 14:26:19 -0800359 key, 1, len(self._member_names), self._last_values[:],
Ethan Furman6d3dfee2020-12-08 12:26:56 -0800360 )
Ethan Furmanfc23a942020-09-16 12:37:54 -0700361 self._auto_called = True
Ethan Furman3515dcc2016-09-18 13:15:41 -0700362 value = value.value
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700363 self._member_names.append(key)
Ethan Furmanc16595e2016-09-10 23:36:59 -0700364 self._last_values.append(value)
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700365 super().__setitem__(key, value)
366
Ethan Furmana6582872020-12-10 13:07:00 -0800367 def update(self, members, **more_members):
368 try:
369 for name in members.keys():
370 self[name] = members[name]
371 except AttributeError:
372 for name, value in members:
373 self[name] = value
374 for name, value in more_members.items():
375 self[name] = value
376
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700377
Ethan Furmanb7751062021-03-30 21:17:26 -0700378class EnumType(type):
Ethan Furman6d3dfee2020-12-08 12:26:56 -0800379 """
380 Metaclass for Enum
381 """
Ethan Furman7aaeb2a2021-01-25 14:26:19 -0800382
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700383 @classmethod
Ethan Furman6ec0ade2020-12-24 10:05:02 -0800384 def __prepare__(metacls, cls, bases, **kwds):
Ethan Furman3064dbf2020-09-16 07:11:57 -0700385 # check that previous enum members do not exist
386 metacls._check_for_existing_members(cls, bases)
Ethan Furmanee47e5c2016-08-31 00:12:15 -0700387 # create the namespace dict
388 enum_dict = _EnumDict()
Ethan Furman7cf0aad2020-12-09 17:12:11 -0800389 enum_dict._cls_name = cls
Ethan Furmanee47e5c2016-08-31 00:12:15 -0700390 # inherit previous flags and _generate_next_value_ function
Ethan Furman3064dbf2020-09-16 07:11:57 -0700391 member_type, first_enum = metacls._get_mixins_(cls, bases)
Ethan Furmanee47e5c2016-08-31 00:12:15 -0700392 if first_enum is not None:
Ethan Furman6d3dfee2020-12-08 12:26:56 -0800393 enum_dict['_generate_next_value_'] = getattr(
394 first_enum, '_generate_next_value_', None,
395 )
Ethan Furmanee47e5c2016-08-31 00:12:15 -0700396 return enum_dict
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700397
Ethan Furmana02cb472021-04-21 10:20:44 -0700398 def __new__(metacls, cls, bases, classdict, *, boundary=None, _simple=False, **kwds):
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700399 # an Enum class is final once enumeration items have been defined; it
400 # cannot be mixed with other types (int, float, etc.) if it has an
401 # inherited __new__ unless a new __new__ is defined (or the resulting
402 # class will fail).
Ethan Furmana4b1bb42018-01-22 07:56:37 -0800403 #
404 # remove any keys listed in _ignore_
Ethan Furmana02cb472021-04-21 10:20:44 -0700405 if _simple:
406 return super().__new__(metacls, cls, bases, classdict, **kwds)
Ethan Furmana4b1bb42018-01-22 07:56:37 -0800407 classdict.setdefault('_ignore_', []).append('_ignore_')
408 ignore = classdict['_ignore_']
409 for key in ignore:
410 classdict.pop(key, None)
Ethan Furmanc314e602021-01-12 23:47:57 -0800411 #
412 # grab member names
413 member_names = classdict._member_names
414 #
415 # check for illegal enum names (any others?)
416 invalid_names = set(member_names) & {'mro', ''}
417 if invalid_names:
418 raise ValueError('Invalid enum member name: {0}'.format(
419 ','.join(invalid_names)))
420 #
421 # adjust the sunders
422 _order_ = classdict.pop('_order_', None)
423 # convert to normal dict
424 classdict = dict(classdict.items())
425 #
426 # data type of member and the controlling Enum class
Ethan Furman3064dbf2020-09-16 07:11:57 -0700427 member_type, first_enum = metacls._get_mixins_(cls, bases)
Ethan Furmanc2667362020-12-07 00:17:31 -0800428 __new__, save_new, use_args = metacls._find_new_(
429 classdict, member_type, first_enum,
430 )
Ethan Furmanc314e602021-01-12 23:47:57 -0800431 classdict['_new_member_'] = __new__
432 classdict['_use_args_'] = use_args
433 #
434 # convert future enum members into temporary _proto_members
Ethan Furman7aaeb2a2021-01-25 14:26:19 -0800435 # and record integer values in case this will be a Flag
436 flag_mask = 0
Ethan Furmanc314e602021-01-12 23:47:57 -0800437 for name in member_names:
Ethan Furman7aaeb2a2021-01-25 14:26:19 -0800438 value = classdict[name]
439 if isinstance(value, int):
440 flag_mask |= value
441 classdict[name] = _proto_member(value)
Ethan Furmanc314e602021-01-12 23:47:57 -0800442 #
Ethan Furman7aaeb2a2021-01-25 14:26:19 -0800443 # house-keeping structures
Ethan Furmanc314e602021-01-12 23:47:57 -0800444 classdict['_member_names_'] = []
445 classdict['_member_map_'] = {}
446 classdict['_value2member_map_'] = {}
Ethan Furman6bd92882021-04-27 13:05:08 -0700447 classdict['_unhashable_values_'] = []
Ethan Furmanc314e602021-01-12 23:47:57 -0800448 classdict['_member_type_'] = member_type
449 #
Ethan Furman7aaeb2a2021-01-25 14:26:19 -0800450 # Flag structures (will be removed if final class is not a Flag
451 classdict['_boundary_'] = (
452 boundary
453 or getattr(first_enum, '_boundary_', None)
454 )
455 classdict['_flag_mask_'] = flag_mask
456 classdict['_all_bits_'] = 2 ** ((flag_mask).bit_length()) - 1
457 classdict['_inverted_'] = None
458 #
Ethan Furmanc314e602021-01-12 23:47:57 -0800459 # create a default docstring if one has not been provided
460 if '__doc__' not in classdict:
461 classdict['__doc__'] = 'An enumeration.'
462 try:
463 exc = None
464 enum_class = super().__new__(metacls, cls, bases, classdict, **kwds)
465 except RuntimeError as e:
466 # any exceptions raised by member.__new__ will get converted to a
467 # RuntimeError, so get that original exception back and raise it instead
468 exc = e.__cause__ or e
469 if exc is not None:
470 raise exc
471 #
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700472 # double check that repr and friends are not the mixin's or various
473 # things break (such as pickle)
Ethan Furman22415ad2020-09-15 16:28:25 -0700474 # however, if the method is defined in the Enum itself, don't replace
475 # it
Ethan Furmandc870522014-02-18 12:37:12 -0800476 for name in ('__repr__', '__str__', '__format__', '__reduce_ex__'):
Ethan Furman22415ad2020-09-15 16:28:25 -0700477 if name in classdict:
478 continue
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700479 class_method = getattr(enum_class, name)
480 obj_method = getattr(member_type, name, None)
481 enum_method = getattr(first_enum, name, None)
482 if obj_method is not None and obj_method is class_method:
483 setattr(enum_class, name, enum_method)
Ethan Furmanc314e602021-01-12 23:47:57 -0800484 #
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700485 # replace any other __new__ with our own (as long as Enum is not None,
486 # anyway) -- again, this is to support pickle
487 if Enum is not None:
488 # if the user defined their own __new__, save it before it gets
489 # clobbered in case they subclass later
490 if save_new:
491 enum_class.__new_member__ = __new__
492 enum_class.__new__ = Enum.__new__
Ethan Furmanc314e602021-01-12 23:47:57 -0800493 #
Ethan Furmane8e61272016-08-20 07:19:31 -0700494 # py3 support for definition order (helps keep py2/py3 code in sync)
Ethan Furman7aaeb2a2021-01-25 14:26:19 -0800495 #
496 # _order_ checking is spread out into three/four steps
497 # - if enum_class is a Flag:
498 # - remove any non-single-bit flags from _order_
499 # - remove any aliases from _order_
500 # - check that _order_ and _member_names_ match
501 #
502 # step 1: ensure we have a list
Ethan Furmane8e61272016-08-20 07:19:31 -0700503 if _order_ is not None:
504 if isinstance(_order_, str):
505 _order_ = _order_.replace(',', ' ').split()
Ethan Furman7aaeb2a2021-01-25 14:26:19 -0800506 #
507 # remove Flag structures if final class is not a Flag
508 if (
509 Flag is None and cls != 'Flag'
510 or Flag is not None and not issubclass(enum_class, Flag)
511 ):
512 delattr(enum_class, '_boundary_')
513 delattr(enum_class, '_flag_mask_')
514 delattr(enum_class, '_all_bits_')
515 delattr(enum_class, '_inverted_')
516 elif Flag is not None and issubclass(enum_class, Flag):
517 # ensure _all_bits_ is correct and there are no missing flags
518 single_bit_total = 0
519 multi_bit_total = 0
520 for flag in enum_class._member_map_.values():
521 flag_value = flag._value_
522 if _is_single_bit(flag_value):
523 single_bit_total |= flag_value
524 else:
525 # multi-bit flags are considered aliases
526 multi_bit_total |= flag_value
Ethan Furman7aaeb2a2021-01-25 14:26:19 -0800527 enum_class._flag_mask_ = single_bit_total
528 #
529 # set correct __iter__
530 member_list = [m._value_ for m in enum_class]
531 if member_list != sorted(member_list):
532 enum_class._iter_member_ = enum_class._iter_member_by_def_
533 if _order_:
534 # _order_ step 2: remove any items from _order_ that are not single-bit
535 _order_ = [
536 o
537 for o in _order_
538 if o not in enum_class._member_map_ or _is_single_bit(enum_class[o]._value_)
539 ]
540 #
541 if _order_:
542 # _order_ step 3: remove aliases from _order_
543 _order_ = [
544 o
545 for o in _order_
546 if (
547 o not in enum_class._member_map_
548 or
549 (o in enum_class._member_map_ and o in enum_class._member_names_)
550 )]
551 # _order_ step 4: verify that _order_ and _member_names_ match
Ethan Furmane8e61272016-08-20 07:19:31 -0700552 if _order_ != enum_class._member_names_:
Ethan Furman7aaeb2a2021-01-25 14:26:19 -0800553 raise TypeError(
554 'member order does not match _order_:\n%r\n%r'
555 % (enum_class._member_names_, _order_)
556 )
Ethan Furmanc314e602021-01-12 23:47:57 -0800557 #
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700558 return enum_class
559
Ethan Furman5de67b12016-04-13 23:52:09 -0700560 def __bool__(self):
561 """
562 classes/types should always be True.
563 """
564 return True
565
Ethan Furman7aaeb2a2021-01-25 14:26:19 -0800566 def __call__(cls, value, names=None, *, module=None, qualname=None, type=None, start=1, boundary=None):
Ethan Furman6d3dfee2020-12-08 12:26:56 -0800567 """
568 Either returns an existing member, or creates a new enum class.
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700569
570 This method is used both when an enum class is given a value to match
571 to an enumeration member (i.e. Color(3)) and for the functional API
Ethan Furman23bb6f42016-11-21 09:22:05 -0800572 (i.e. Color = Enum('Color', names='RED GREEN BLUE')).
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700573
Ethan Furman2da95042014-03-03 12:42:52 -0800574 When used for the functional API:
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700575
Ethan Furman2da95042014-03-03 12:42:52 -0800576 `value` will be the name of the new class.
577
578 `names` should be either a string of white-space/comma delimited names
Ethan Furmand9925a12014-09-16 20:35:55 -0700579 (values will start at `start`), or an iterator/mapping of name, value pairs.
Ethan Furman2da95042014-03-03 12:42:52 -0800580
581 `module` should be set to the module this class is being created in;
582 if it is not set, an attempt to find that module will be made, but if
583 it fails the class will not be picklable.
584
585 `qualname` should be set to the actual location this class can be found
586 at in its module; by default it is set to the global scope. If this is
587 not correct, unpickling will fail in some circumstances.
588
589 `type`, if set, will be mixed in as the first base class.
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700590 """
591 if names is None: # simple value lookup
592 return cls.__new__(cls, value)
593 # otherwise, functional API: we're creating a new Enum type
Ethan Furman6d3dfee2020-12-08 12:26:56 -0800594 return cls._create_(
595 value,
596 names,
597 module=module,
598 qualname=qualname,
599 type=type,
600 start=start,
Ethan Furman7aaeb2a2021-01-25 14:26:19 -0800601 boundary=boundary,
Ethan Furman6d3dfee2020-12-08 12:26:56 -0800602 )
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700603
604 def __contains__(cls, member):
Rahul Jha94306522018-09-10 23:51:04 +0530605 if not isinstance(member, Enum):
Ethan Furman6bd92882021-04-27 13:05:08 -0700606 import warnings
607 warnings.warn(
608 "in 3.12 __contains__ will no longer raise TypeError, but will return True or\n"
609 "False depending on whether the value is a member or the value of a member",
610 DeprecationWarning,
611 stacklevel=2,
612 )
Rahul Jha94306522018-09-10 23:51:04 +0530613 raise TypeError(
614 "unsupported operand type(s) for 'in': '%s' and '%s'" % (
615 type(member).__qualname__, cls.__class__.__qualname__))
Ethan Furman0081f232014-09-16 17:31:23 -0700616 return isinstance(member, cls) and member._name_ in cls._member_map_
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700617
Ethan Furman64a99722013-09-22 16:18:19 -0700618 def __delattr__(cls, attr):
619 # nicer error message when someone tries to delete an attribute
620 # (see issue19025).
621 if attr in cls._member_map_:
Ethan Furman6d3dfee2020-12-08 12:26:56 -0800622 raise AttributeError("%s: cannot delete Enum member %r." % (cls.__name__, attr))
Ethan Furman64a99722013-09-22 16:18:19 -0700623 super().__delattr__(attr)
624
Ethan Furman388a3922013-08-12 06:51:41 -0700625 def __dir__(self):
Ethan Furman6d3dfee2020-12-08 12:26:56 -0800626 return (
627 ['__class__', '__doc__', '__members__', '__module__']
628 + self._member_names_
629 )
Ethan Furman388a3922013-08-12 06:51:41 -0700630
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700631 def __getattr__(cls, name):
Ethan Furman6d3dfee2020-12-08 12:26:56 -0800632 """
633 Return the enum member matching `name`
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700634
635 We use __getattr__ instead of descriptors or inserting into the enum
636 class' __dict__ in order to support `name` and `value` being both
637 properties for enum members (which live in the class' __dict__) and
638 enum members themselves.
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700639 """
640 if _is_dunder(name):
641 raise AttributeError(name)
642 try:
Ethan Furman520ad572013-07-19 19:47:21 -0700643 return cls._member_map_[name]
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700644 except KeyError:
645 raise AttributeError(name) from None
646
647 def __getitem__(cls, name):
Ethan Furman520ad572013-07-19 19:47:21 -0700648 return cls._member_map_[name]
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700649
650 def __iter__(cls):
Ethan Furman6d3dfee2020-12-08 12:26:56 -0800651 """
652 Returns members in definition order.
653 """
Ethan Furman520ad572013-07-19 19:47:21 -0700654 return (cls._member_map_[name] for name in cls._member_names_)
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700655
656 def __len__(cls):
Ethan Furman520ad572013-07-19 19:47:21 -0700657 return len(cls._member_names_)
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700658
Ethan Furmanc314e602021-01-12 23:47:57 -0800659 @_bltin_property
Ethan Furman2131a4a2013-09-14 18:11:24 -0700660 def __members__(cls):
Ethan Furman6d3dfee2020-12-08 12:26:56 -0800661 """
662 Returns a mapping of member name->value.
Ethan Furman2131a4a2013-09-14 18:11:24 -0700663
664 This mapping lists all enum members, including aliases. Note that this
665 is a read-only view of the internal mapping.
Ethan Furman2131a4a2013-09-14 18:11:24 -0700666 """
667 return MappingProxyType(cls._member_map_)
668
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700669 def __repr__(cls):
Ethan Furman74964862021-06-10 07:24:20 -0700670 if Flag is not None and issubclass(cls, Flag):
671 return "<flag %r>" % cls.__name__
672 else:
673 return "<enum %r>" % cls.__name__
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700674
Ethan Furman2131a4a2013-09-14 18:11:24 -0700675 def __reversed__(cls):
Ethan Furman6d3dfee2020-12-08 12:26:56 -0800676 """
677 Returns members in reverse definition order.
678 """
Ethan Furman2131a4a2013-09-14 18:11:24 -0700679 return (cls._member_map_[name] for name in reversed(cls._member_names_))
680
Ethan Furmanf203f2d2013-09-06 07:16:48 -0700681 def __setattr__(cls, name, value):
Ethan Furman6d3dfee2020-12-08 12:26:56 -0800682 """
683 Block attempts to reassign Enum members.
Ethan Furmanf203f2d2013-09-06 07:16:48 -0700684
685 A simple assignment to the class namespace only changes one of the
686 several possible ways to get an Enum member from the Enum class,
687 resulting in an inconsistent Enumeration.
Ethan Furmanf203f2d2013-09-06 07:16:48 -0700688 """
689 member_map = cls.__dict__.get('_member_map_', {})
690 if name in member_map:
Ethan Furmana02cb472021-04-21 10:20:44 -0700691 raise AttributeError('Cannot reassign member %r.' % (name, ))
Ethan Furmanf203f2d2013-09-06 07:16:48 -0700692 super().__setattr__(name, value)
693
Ethan Furman7aaeb2a2021-01-25 14:26:19 -0800694 def _create_(cls, class_name, names, *, module=None, qualname=None, type=None, start=1, boundary=None):
Ethan Furman6d3dfee2020-12-08 12:26:56 -0800695 """
696 Convenience method to create a new Enum class.
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700697
698 `names` can be:
699
700 * A string containing member names, separated either with spaces or
Ethan Furmand9925a12014-09-16 20:35:55 -0700701 commas. Values are incremented by 1 from `start`.
702 * An iterable of member names. Values are incremented by 1 from `start`.
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700703 * An iterable of (member name, value) pairs.
Ethan Furmand9925a12014-09-16 20:35:55 -0700704 * A mapping of member name -> value pairs.
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700705 """
706 metacls = cls.__class__
707 bases = (cls, ) if type is None else (type, cls)
Ethan Furman3064dbf2020-09-16 07:11:57 -0700708 _, first_enum = cls._get_mixins_(cls, bases)
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700709 classdict = metacls.__prepare__(class_name, bases)
710
711 # special processing needed for names?
712 if isinstance(names, str):
713 names = names.replace(',', ' ').split()
Dong-hee Nadcc8ce42017-06-22 01:52:32 +0900714 if isinstance(names, (tuple, list)) and names and isinstance(names[0], str):
Ethan Furmanee47e5c2016-08-31 00:12:15 -0700715 original_names, names = names, []
Ethan Furmanc16595e2016-09-10 23:36:59 -0700716 last_values = []
Ethan Furmanee47e5c2016-08-31 00:12:15 -0700717 for count, name in enumerate(original_names):
Ethan Furmanc16595e2016-09-10 23:36:59 -0700718 value = first_enum._generate_next_value_(name, start, count, last_values[:])
719 last_values.append(value)
720 names.append((name, value))
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700721
722 # Here, names is either an iterable of (name, value) or a mapping.
723 for item in names:
724 if isinstance(item, str):
725 member_name, member_value = item, names[item]
726 else:
727 member_name, member_value = item
728 classdict[member_name] = member_value
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700729
730 # TODO: replace the frame hack if a blessed way to know the calling
731 # module is ever developed
732 if module is None:
733 try:
734 module = sys._getframe(2).f_globals['__name__']
Pablo Galindo293dd232019-11-19 21:34:03 +0000735 except (AttributeError, ValueError, KeyError):
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700736 pass
737 if module is None:
Ethan Furmanc314e602021-01-12 23:47:57 -0800738 _make_class_unpicklable(classdict)
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700739 else:
Ethan Furmanc314e602021-01-12 23:47:57 -0800740 classdict['__module__'] = module
Ethan Furmanca1b7942014-02-08 11:36:27 -0800741 if qualname is not None:
Ethan Furmanc314e602021-01-12 23:47:57 -0800742 classdict['__qualname__'] = qualname
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700743
Ethan Furman7aaeb2a2021-01-25 14:26:19 -0800744 return metacls.__new__(metacls, class_name, bases, classdict, boundary=boundary)
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700745
Ethan Furmana02cb472021-04-21 10:20:44 -0700746 def _convert_(cls, name, module, filter, source=None, *, boundary=None):
747
orlnub1230fb9fad2018-09-12 20:28:53 +0300748 """
749 Create a new Enum subclass that replaces a collection of global constants
750 """
751 # convert all constants from source (or module) that pass filter() to
752 # a new Enum called name, and export the enum and its members back to
753 # module;
754 # also, replace the __reduce_ex__ method so unpickling works in
755 # previous Python versions
Ethan Furmanb7751062021-03-30 21:17:26 -0700756 module_globals = sys.modules[module].__dict__
orlnub1230fb9fad2018-09-12 20:28:53 +0300757 if source:
Ethan Furmanb7751062021-03-30 21:17:26 -0700758 source = source.__dict__
orlnub1230fb9fad2018-09-12 20:28:53 +0300759 else:
760 source = module_globals
761 # _value2member_map_ is populated in the same order every time
762 # for a consistent reverse mapping of number to name when there
763 # are multiple names for the same number.
764 members = [
765 (name, value)
766 for name, value in source.items()
767 if filter(name)]
768 try:
769 # sort by value
770 members.sort(key=lambda t: (t[1], t[0]))
771 except TypeError:
772 # unless some values aren't comparable, in which case sort by name
773 members.sort(key=lambda t: t[0])
Ethan Furmana02cb472021-04-21 10:20:44 -0700774 body = {t[0]: t[1] for t in members}
775 body['__module__'] = module
776 tmp_cls = type(name, (object, ), body)
777 cls = _simple_enum(etype=cls, boundary=boundary or KEEP)(tmp_cls)
Miss Islington (bot)b6131322021-06-10 16:37:27 -0700778 cls.__reduce_ex__ = _reduce_ex_by_global_name
Ethan Furmanb7751062021-03-30 21:17:26 -0700779 global_enum(cls)
orlnub1230fb9fad2018-09-12 20:28:53 +0300780 module_globals[name] = cls
781 return cls
782
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700783 @staticmethod
Ethan Furman3064dbf2020-09-16 07:11:57 -0700784 def _check_for_existing_members(class_name, bases):
785 for chain in bases:
786 for base in chain.__mro__:
787 if issubclass(base, Enum) and base._member_names_:
Ethan Furman6d3dfee2020-12-08 12:26:56 -0800788 raise TypeError(
789 "%s: cannot extend enumeration %r"
790 % (class_name, base.__name__)
791 )
Ethan Furman3064dbf2020-09-16 07:11:57 -0700792
793 @staticmethod
794 def _get_mixins_(class_name, bases):
Ethan Furman6d3dfee2020-12-08 12:26:56 -0800795 """
796 Returns the type for creating enum members, and the first inherited
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700797 enum class.
798
799 bases: the tuple of bases that was given to __new__
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700800 """
801 if not bases:
802 return object, Enum
803
Ethan Furman5bdab642018-09-21 19:03:09 -0700804 def _find_data_type(bases):
Miss Islington (bot)01286012021-06-10 15:01:03 -0700805 data_types = set()
Ethan Furman5bdab642018-09-21 19:03:09 -0700806 for chain in bases:
Ethan Furmanbff01f32020-09-15 15:56:26 -0700807 candidate = None
Ethan Furman5bdab642018-09-21 19:03:09 -0700808 for base in chain.__mro__:
809 if base is object:
810 continue
Ethan Furmanc2667362020-12-07 00:17:31 -0800811 elif issubclass(base, Enum):
812 if base._member_type_ is not object:
Miss Islington (bot)01286012021-06-10 15:01:03 -0700813 data_types.add(base._member_type_)
Ethan Furmanc2667362020-12-07 00:17:31 -0800814 break
Ethan Furman5bdab642018-09-21 19:03:09 -0700815 elif '__new__' in base.__dict__:
Ethan Furmancd453852018-10-05 23:29:36 -0700816 if issubclass(base, Enum):
Ethan Furman5bdab642018-09-21 19:03:09 -0700817 continue
Miss Islington (bot)01286012021-06-10 15:01:03 -0700818 data_types.add(candidate or base)
Ethan Furmanbff01f32020-09-15 15:56:26 -0700819 break
Ethan Furmanc2667362020-12-07 00:17:31 -0800820 else:
Ethan Furmanbff01f32020-09-15 15:56:26 -0700821 candidate = base
822 if len(data_types) > 1:
Ethan Furman3064dbf2020-09-16 07:11:57 -0700823 raise TypeError('%r: too many data types: %r' % (class_name, data_types))
Ethan Furmanbff01f32020-09-15 15:56:26 -0700824 elif data_types:
Miss Islington (bot)01286012021-06-10 15:01:03 -0700825 return data_types.pop()
Ethan Furmanbff01f32020-09-15 15:56:26 -0700826 else:
827 return None
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700828
Ethan Furman5bdab642018-09-21 19:03:09 -0700829 # ensure final parent class is an Enum derivative, find any concrete
830 # data type, and check that Enum has no members
831 first_enum = bases[-1]
832 if not issubclass(first_enum, Enum):
833 raise TypeError("new enumerations should be created as "
834 "`EnumName([mixin_type, ...] [data_type,] enum_type)`")
835 member_type = _find_data_type(bases) or object
836 if first_enum._member_names_:
837 raise TypeError("Cannot extend enumerations")
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700838 return member_type, first_enum
839
840 @staticmethod
841 def _find_new_(classdict, member_type, first_enum):
Ethan Furman6d3dfee2020-12-08 12:26:56 -0800842 """
843 Returns the __new__ to be used for creating the enum members.
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700844
845 classdict: the class dictionary given to __new__
846 member_type: the data type whose __new__ will be used by default
847 first_enum: enumeration to check for an overriding __new__
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700848 """
849 # now find the correct __new__, checking to see of one was defined
850 # by the user; also check earlier enum classes in case a __new__ was
851 # saved as __new_member__
852 __new__ = classdict.get('__new__', None)
853
854 # should __new__ be saved as __new_member__ later?
Ethan Furmana02cb472021-04-21 10:20:44 -0700855 save_new = first_enum is not None and __new__ is not None
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700856
857 if __new__ is None:
858 # check all possibles for __new_member__ before falling back to
859 # __new__
860 for method in ('__new_member__', '__new__'):
861 for possible in (member_type, first_enum):
862 target = getattr(possible, method, None)
863 if target not in {
864 None,
865 None.__new__,
866 object.__new__,
867 Enum.__new__,
868 }:
869 __new__ = target
870 break
871 if __new__ is not None:
872 break
873 else:
874 __new__ = object.__new__
875
876 # if a non-object.__new__ is used then whatever value/tuple was
877 # assigned to the enum member name will be passed to __new__ and to the
878 # new enum member's __init__
Ethan Furmana02cb472021-04-21 10:20:44 -0700879 if first_enum is None or __new__ in (Enum.__new__, object.__new__):
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700880 use_args = False
881 else:
882 use_args = True
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700883 return __new__, save_new, use_args
Ethan Furmanb7751062021-03-30 21:17:26 -0700884EnumMeta = EnumType
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700885
886
Ethan Furmanb7751062021-03-30 21:17:26 -0700887class Enum(metaclass=EnumType):
Ethan Furman6d3dfee2020-12-08 12:26:56 -0800888 """
889 Generic enumeration.
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700890
891 Derive from this class to define new enumerations.
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700892 """
Ethan Furman7aaeb2a2021-01-25 14:26:19 -0800893
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700894 def __new__(cls, value):
895 # all enum instances are actually created during class construction
896 # without calling this method; this method is called by the metaclass'
897 # __call__ (i.e. Color(3) ), and by pickle
898 if type(value) is cls:
Ethan Furman23bb6f42016-11-21 09:22:05 -0800899 # For lookups like Color(Color.RED)
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700900 return value
901 # by-value search for a matching enum member
902 # see if it's in the reverse mapping (for hashable values)
Ethan Furman2aa27322013-07-19 19:35:56 -0700903 try:
Andrew Svetlov34ae04f2018-12-26 20:45:33 +0200904 return cls._value2member_map_[value]
905 except KeyError:
906 # Not found, no need to do long O(n) search
907 pass
Ethan Furman2aa27322013-07-19 19:35:56 -0700908 except TypeError:
909 # not there, now do long search -- O(n) behavior
Ethan Furman520ad572013-07-19 19:47:21 -0700910 for member in cls._member_map_.values():
Ethan Furman0081f232014-09-16 17:31:23 -0700911 if member._value_ == value:
Ethan Furman2aa27322013-07-19 19:35:56 -0700912 return member
Ethan Furmanee47e5c2016-08-31 00:12:15 -0700913 # still not found -- try _missing_ hook
Ethan Furman019f0a02018-09-12 11:43:34 -0700914 try:
915 exc = None
916 result = cls._missing_(value)
917 except Exception as e:
918 exc = e
919 result = None
Ethan Furman8c14f5a2021-04-12 08:51:20 -0700920 try:
921 if isinstance(result, cls):
922 return result
923 elif (
924 Flag is not None and issubclass(cls, Flag)
925 and cls._boundary_ is EJECT and isinstance(result, int)
926 ):
927 return result
928 else:
929 ve_exc = ValueError("%r is not a valid %s" % (value, cls.__qualname__))
930 if result is None and exc is None:
931 raise ve_exc
932 elif exc is None:
933 exc = TypeError(
934 'error in %s._missing_: returned %r instead of None or a valid member'
935 % (cls.__name__, result)
936 )
937 if not isinstance(exc, ValueError):
938 exc.__context__ = ve_exc
939 raise exc
940 finally:
941 # ensure all variables that could hold an exception are destroyed
942 exc = None
943 ve_exc = None
Ethan Furmanee47e5c2016-08-31 00:12:15 -0700944
Ethan Furmanc16595e2016-09-10 23:36:59 -0700945 def _generate_next_value_(name, start, count, last_values):
Ethan Furman6d3dfee2020-12-08 12:26:56 -0800946 """
947 Generate the next value when not given.
948
949 name: the name of the member
950 start: the initial start value or None
951 count: the number of existing members
952 last_value: the last value assigned or None
953 """
Ethan Furmanc16595e2016-09-10 23:36:59 -0700954 for last_value in reversed(last_values):
955 try:
956 return last_value + 1
957 except TypeError:
958 pass
959 else:
Ethan Furmanee47e5c2016-08-31 00:12:15 -0700960 return start
Ethan Furmanc16595e2016-09-10 23:36:59 -0700961
Ethan Furmanee47e5c2016-08-31 00:12:15 -0700962 @classmethod
963 def _missing_(cls, value):
Ethan Furmanc95ad7a2020-09-16 10:26:50 -0700964 return None
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700965
966 def __repr__(self):
Ethan Furmanb7751062021-03-30 21:17:26 -0700967 return "%s.%s" % ( self.__class__.__name__, self._name_)
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700968
969 def __str__(self):
Ethan Furmanb7751062021-03-30 21:17:26 -0700970 return "%s" % (self._name_, )
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700971
Ethan Furman388a3922013-08-12 06:51:41 -0700972 def __dir__(self):
Ethan Furman6d3dfee2020-12-08 12:26:56 -0800973 """
974 Returns all members and all public methods
975 """
Ethan Furman0ae550b2014-10-14 08:58:32 -0700976 added_behavior = [
977 m
978 for cls in self.__class__.mro()
979 for m in cls.__dict__
Ethan Furman354ecf12015-03-11 08:43:12 -0700980 if m[0] != '_' and m not in self._member_map_
Angelin BOOZ68526fe2020-09-21 15:11:06 +0200981 ] + [m for m in self.__dict__ if m[0] != '_']
Ethan Furmanec5f8eb2014-10-21 13:40:35 -0700982 return (['__class__', '__doc__', '__module__'] + added_behavior)
Ethan Furman388a3922013-08-12 06:51:41 -0700983
Ethan Furmanec15a822013-08-31 19:17:41 -0700984 def __format__(self, format_spec):
Ethan Furman6d3dfee2020-12-08 12:26:56 -0800985 """
986 Returns format using actual value type unless __str__ has been overridden.
987 """
Ethan Furmanec15a822013-08-31 19:17:41 -0700988 # mixed-in Enums should use the mixed-in type's __format__, otherwise
989 # we can get strange results with the Enum name showing up instead of
990 # the value
991
thatneat2f19e822019-07-04 11:28:37 -0700992 # pure Enum branch, or branch with __str__ explicitly overridden
Ethan Furman37440ee2020-12-08 11:14:10 -0800993 str_overridden = type(self).__str__ not in (Enum.__str__, Flag.__str__)
thatneat2f19e822019-07-04 11:28:37 -0700994 if self._member_type_ is object or str_overridden:
Ethan Furmanec15a822013-08-31 19:17:41 -0700995 cls = str
996 val = str(self)
997 # mix-in branch
998 else:
Ethan Furman6bd92882021-04-27 13:05:08 -0700999 if not format_spec or format_spec in ('{}','{:}'):
1000 import warnings
1001 warnings.warn(
1002 "in 3.12 format() will use the enum member, not the enum member's value;\n"
Pablo Galindo9a42d502021-05-01 20:26:09 +01001003 "use a format specifier, such as :d for an IntEnum member, to maintain "
Ethan Furman6bd92882021-04-27 13:05:08 -07001004 "the current display",
1005 DeprecationWarning,
1006 stacklevel=2,
1007 )
Ethan Furmanec15a822013-08-31 19:17:41 -07001008 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):
Miss Islington (bot)b6131322021-06-10 16:37:27 -07001016 return getattr, (self.__class__, self._name_)
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
Miss Islington (bot)b6131322021-06-10 16:37:27 -07001077def _reduce_ex_by_global_name(self, proto):
Ethan Furman24e837f2015-03-18 17:27:57 -07001078 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_)
Ethan Furman74964862021-06-10 07:24:20 -07001288 if isinstance(self._inverted_, self.__class__):
1289 self._inverted_._inverted_ = self
Ethan Furman7aaeb2a2021-01-25 14:26:19 -08001290 return self._inverted_
Ethan Furmanee47e5c2016-08-31 00:12:15 -07001291
1292
Ethan Furman7aaeb2a2021-01-25 14:26:19 -08001293class IntFlag(int, Flag, boundary=EJECT):
Ethan Furman6d3dfee2020-12-08 12:26:56 -08001294 """
1295 Support for integer-based Flags
1296 """
Ethan Furmanee47e5c2016-08-31 00:12:15 -07001297
Ethan Furmanee47e5c2016-08-31 00:12:15 -07001298 def __or__(self, other):
Ethan Furman7aaeb2a2021-01-25 14:26:19 -08001299 if isinstance(other, self.__class__):
1300 other = other._value_
1301 elif isinstance(other, int):
1302 other = other
1303 else:
Ethan Furmanee47e5c2016-08-31 00:12:15 -07001304 return NotImplemented
Ethan Furman7aaeb2a2021-01-25 14:26:19 -08001305 value = self._value_
1306 return self.__class__(value | other)
Ethan Furmanee47e5c2016-08-31 00:12:15 -07001307
1308 def __and__(self, other):
Ethan Furman7aaeb2a2021-01-25 14:26:19 -08001309 if isinstance(other, self.__class__):
1310 other = other._value_
1311 elif isinstance(other, int):
1312 other = other
1313 else:
Ethan Furmanee47e5c2016-08-31 00:12:15 -07001314 return NotImplemented
Ethan Furman7aaeb2a2021-01-25 14:26:19 -08001315 value = self._value_
1316 return self.__class__(value & other)
Ethan Furmanee47e5c2016-08-31 00:12:15 -07001317
1318 def __xor__(self, other):
Ethan Furman7aaeb2a2021-01-25 14:26:19 -08001319 if isinstance(other, self.__class__):
1320 other = other._value_
1321 elif isinstance(other, int):
1322 other = other
1323 else:
Ethan Furmanee47e5c2016-08-31 00:12:15 -07001324 return NotImplemented
Ethan Furman7aaeb2a2021-01-25 14:26:19 -08001325 value = self._value_
1326 return self.__class__(value ^ other)
Ethan Furmanee47e5c2016-08-31 00:12:15 -07001327
1328 __ror__ = __or__
1329 __rand__ = __and__
1330 __rxor__ = __xor__
Ethan Furman7aaeb2a2021-01-25 14:26:19 -08001331 __invert__ = Flag.__invert__
Ethan Furmanee47e5c2016-08-31 00:12:15 -07001332
1333def _high_bit(value):
Ethan Furman6d3dfee2020-12-08 12:26:56 -08001334 """
1335 returns index of highest bit, or -1 if value is zero or negative
1336 """
Ethan Furman3515dcc2016-09-18 13:15:41 -07001337 return value.bit_length() - 1
Ethan Furmanee47e5c2016-08-31 00:12:15 -07001338
Ethan Furmanf24bb352013-07-18 17:05:39 -07001339def unique(enumeration):
Ethan Furman6d3dfee2020-12-08 12:26:56 -08001340 """
1341 Class decorator for enumerations ensuring unique member values.
1342 """
Ethan Furmanf24bb352013-07-18 17:05:39 -07001343 duplicates = []
1344 for name, member in enumeration.__members__.items():
1345 if name != member.name:
1346 duplicates.append((name, member.name))
1347 if duplicates:
1348 alias_details = ', '.join(
1349 ["%s -> %s" % (alias, name) for (alias, name) in duplicates])
1350 raise ValueError('duplicate values found in %r: %s' %
1351 (enumeration, alias_details))
1352 return enumeration
Ethan Furman3515dcc2016-09-18 13:15:41 -07001353
Ethan Furman7aaeb2a2021-01-25 14:26:19 -08001354def _power_of_two(value):
1355 if value < 1:
1356 return False
1357 return value == 2 ** _high_bit(value)
Ethan Furmanb7751062021-03-30 21:17:26 -07001358
1359def global_enum_repr(self):
1360 return '%s.%s' % (self.__class__.__module__, self._name_)
1361
1362def global_flag_repr(self):
1363 module = self.__class__.__module__
1364 cls_name = self.__class__.__name__
1365 if self._name_ is None:
1366 return "%x" % (module, cls_name, self._value_)
1367 if _is_single_bit(self):
1368 return '%s.%s' % (module, self._name_)
1369 if self._boundary_ is not FlagBoundary.KEEP:
1370 return module + module.join(self.name.split('|'))
1371 else:
1372 name = []
1373 for n in self._name_.split('|'):
1374 if n.startswith('0'):
1375 name.append(n)
1376 else:
1377 name.append('%s.%s' % (module, n))
1378 return '|'.join(name)
1379
1380
1381def global_enum(cls):
1382 """
1383 decorator that makes the repr() of an enum member reference its module
1384 instead of its class; also exports all members to the enum's module's
1385 global namespace
1386 """
1387 if issubclass(cls, Flag):
1388 cls.__repr__ = global_flag_repr
1389 else:
1390 cls.__repr__ = global_enum_repr
1391 sys.modules[cls.__module__].__dict__.update(cls.__members__)
1392 return cls
Ethan Furmana02cb472021-04-21 10:20:44 -07001393
1394def _simple_enum(etype=Enum, *, boundary=None, use_args=None):
1395 """
1396 Class decorator that converts a normal class into an :class:`Enum`. No
1397 safety checks are done, and some advanced behavior (such as
1398 :func:`__init_subclass__`) is not available. Enum creation can be faster
1399 using :func:`simple_enum`.
1400
1401 >>> from enum import Enum, _simple_enum
1402 >>> @_simple_enum(Enum)
1403 ... class Color:
1404 ... RED = auto()
1405 ... GREEN = auto()
1406 ... BLUE = auto()
1407 >>> Color
1408 <enum 'Color'>
1409 """
1410 def convert_class(cls):
1411 nonlocal use_args
1412 cls_name = cls.__name__
1413 if use_args is None:
1414 use_args = etype._use_args_
1415 __new__ = cls.__dict__.get('__new__')
1416 if __new__ is not None:
1417 new_member = __new__.__func__
1418 else:
1419 new_member = etype._member_type_.__new__
1420 attrs = {}
1421 body = {}
1422 if __new__ is not None:
1423 body['__new_member__'] = new_member
1424 body['_new_member_'] = new_member
1425 body['_use_args_'] = use_args
1426 body['_generate_next_value_'] = gnv = etype._generate_next_value_
1427 body['_member_names_'] = member_names = []
1428 body['_member_map_'] = member_map = {}
1429 body['_value2member_map_'] = value2member_map = {}
Ethan Furman6bd92882021-04-27 13:05:08 -07001430 body['_unhashable_values_'] = []
Ethan Furmana02cb472021-04-21 10:20:44 -07001431 body['_member_type_'] = member_type = etype._member_type_
1432 if issubclass(etype, Flag):
1433 body['_boundary_'] = boundary or etype._boundary_
1434 body['_flag_mask_'] = None
1435 body['_all_bits_'] = None
1436 body['_inverted_'] = None
1437 for name, obj in cls.__dict__.items():
1438 if name in ('__dict__', '__weakref__'):
1439 continue
1440 if _is_dunder(name) or _is_private(cls_name, name) or _is_sunder(name) or _is_descriptor(obj):
1441 body[name] = obj
1442 else:
1443 attrs[name] = obj
1444 if cls.__dict__.get('__doc__') is None:
1445 body['__doc__'] = 'An enumeration.'
1446 #
1447 # double check that repr and friends are not the mixin's or various
1448 # things break (such as pickle)
1449 # however, if the method is defined in the Enum itself, don't replace
1450 # it
1451 enum_class = type(cls_name, (etype, ), body, boundary=boundary, _simple=True)
1452 for name in ('__repr__', '__str__', '__format__', '__reduce_ex__'):
1453 if name in body:
1454 continue
1455 class_method = getattr(enum_class, name)
1456 obj_method = getattr(member_type, name, None)
1457 enum_method = getattr(etype, name, None)
1458 if obj_method is not None and obj_method is class_method:
1459 setattr(enum_class, name, enum_method)
1460 gnv_last_values = []
1461 if issubclass(enum_class, Flag):
1462 # Flag / IntFlag
1463 single_bits = multi_bits = 0
1464 for name, value in attrs.items():
1465 if isinstance(value, auto) and auto.value is _auto_null:
1466 value = gnv(name, 1, len(member_names), gnv_last_values)
1467 if value in value2member_map:
1468 # an alias to an existing member
1469 redirect = property()
1470 redirect.__set_name__(enum_class, name)
1471 setattr(enum_class, name, redirect)
1472 member_map[name] = value2member_map[value]
1473 else:
1474 # create the member
1475 if use_args:
1476 if not isinstance(value, tuple):
1477 value = (value, )
1478 member = new_member(enum_class, *value)
1479 value = value[0]
1480 else:
1481 member = new_member(enum_class)
1482 if __new__ is None:
1483 member._value_ = value
1484 member._name_ = name
1485 member.__objclass__ = enum_class
1486 member.__init__(value)
1487 redirect = property()
1488 redirect.__set_name__(enum_class, name)
1489 setattr(enum_class, name, redirect)
1490 member_map[name] = member
1491 member._sort_order_ = len(member_names)
1492 value2member_map[value] = member
1493 if _is_single_bit(value):
1494 # not a multi-bit alias, record in _member_names_ and _flag_mask_
1495 member_names.append(name)
1496 single_bits |= value
1497 else:
1498 multi_bits |= value
1499 gnv_last_values.append(value)
1500 enum_class._flag_mask_ = single_bits
1501 enum_class._all_bits_ = 2 ** ((single_bits|multi_bits).bit_length()) - 1
1502 # set correct __iter__
1503 member_list = [m._value_ for m in enum_class]
1504 if member_list != sorted(member_list):
1505 enum_class._iter_member_ = enum_class._iter_member_by_def_
1506 else:
1507 # Enum / IntEnum / StrEnum
1508 for name, value in attrs.items():
1509 if isinstance(value, auto):
1510 if value.value is _auto_null:
1511 value.value = gnv(name, 1, len(member_names), gnv_last_values)
1512 value = value.value
1513 if value in value2member_map:
1514 # an alias to an existing member
1515 redirect = property()
1516 redirect.__set_name__(enum_class, name)
1517 setattr(enum_class, name, redirect)
1518 member_map[name] = value2member_map[value]
1519 else:
1520 # create the member
1521 if use_args:
1522 if not isinstance(value, tuple):
1523 value = (value, )
1524 member = new_member(enum_class, *value)
1525 value = value[0]
1526 else:
1527 member = new_member(enum_class)
1528 if __new__ is None:
1529 member._value_ = value
1530 member._name_ = name
1531 member.__objclass__ = enum_class
1532 member.__init__(value)
1533 member._sort_order_ = len(member_names)
1534 redirect = property()
1535 redirect.__set_name__(enum_class, name)
1536 setattr(enum_class, name, redirect)
1537 member_map[name] = member
1538 value2member_map[value] = member
1539 member_names.append(name)
1540 gnv_last_values.append(value)
1541 if '__new__' in body:
1542 enum_class.__new_member__ = enum_class.__new__
1543 enum_class.__new__ = Enum.__new__
1544 return enum_class
1545 return convert_class
1546
Ethan Furman74964862021-06-10 07:24:20 -07001547@_simple_enum(StrEnum)
1548class EnumCheck:
1549 """
1550 various conditions to check an enumeration for
1551 """
1552 CONTINUOUS = "no skipped integer values"
1553 NAMED_FLAGS = "multi-flag aliases may not contain unnamed flags"
1554 UNIQUE = "one name per value"
1555CONTINUOUS, NAMED_FLAGS, UNIQUE = EnumCheck
1556
1557
1558class verify:
1559 """
1560 Check an enumeration for various constraints. (see EnumCheck)
1561 """
1562 def __init__(self, *checks):
1563 self.checks = checks
1564 def __call__(self, enumeration):
1565 checks = self.checks
1566 cls_name = enumeration.__name__
1567 if Flag is not None and issubclass(enumeration, Flag):
1568 enum_type = 'flag'
1569 elif issubclass(enumeration, Enum):
1570 enum_type = 'enum'
1571 else:
1572 raise TypeError("the 'verify' decorator only works with Enum and Flag")
1573 for check in checks:
1574 if check is UNIQUE:
1575 # check for duplicate names
1576 duplicates = []
1577 for name, member in enumeration.__members__.items():
1578 if name != member.name:
1579 duplicates.append((name, member.name))
1580 if duplicates:
1581 alias_details = ', '.join(
1582 ["%s -> %s" % (alias, name) for (alias, name) in duplicates])
1583 raise ValueError('aliases found in %r: %s' %
1584 (enumeration, alias_details))
1585 elif check is CONTINUOUS:
1586 values = set(e.value for e in enumeration)
1587 if len(values) < 2:
1588 continue
1589 low, high = min(values), max(values)
1590 missing = []
1591 if enum_type == 'flag':
1592 # check for powers of two
1593 for i in range(_high_bit(low)+1, _high_bit(high)):
1594 if 2**i not in values:
1595 missing.append(2**i)
1596 elif enum_type == 'enum':
1597 # check for powers of one
1598 for i in range(low+1, high):
1599 if i not in values:
1600 missing.append(i)
1601 else:
1602 raise Exception('verify: unknown type %r' % enum_type)
1603 if missing:
1604 raise ValueError('invalid %s %r: missing values %s' % (
1605 enum_type, cls_name, ', '.join((str(m) for m in missing)))
1606 )
1607 elif check is NAMED_FLAGS:
1608 # examine each alias and check for unnamed flags
1609 member_names = enumeration._member_names_
1610 member_values = [m.value for m in enumeration]
1611 missing = []
1612 for name, alias in enumeration._member_map_.items():
1613 if name in member_names:
1614 # not an alias
1615 continue
1616 values = list(_iter_bits_lsb(alias.value))
1617 missed = [v for v in values if v not in member_values]
1618 if missed:
1619 plural = ('', 's')[len(missed) > 1]
1620 a = ('a ', '')[len(missed) > 1]
1621 missing.append('%r is missing %snamed flag%s for value%s %s' % (
1622 name, a, plural, plural,
1623 ', '.join(str(v) for v in missed)
1624 ))
1625 if missing:
1626 raise ValueError(
1627 'invalid Flag %r: %s'
1628 % (cls_name, '; '.join(missing))
1629 )
1630 return enumeration
1631
Ethan Furmana02cb472021-04-21 10:20:44 -07001632def _test_simple_enum(checked_enum, simple_enum):
1633 """
1634 A function that can be used to test an enum created with :func:`_simple_enum`
1635 against the version created by subclassing :class:`Enum`::
1636
1637 >>> from enum import Enum, _simple_enum, _test_simple_enum
1638 >>> @_simple_enum(Enum)
1639 ... class Color:
1640 ... RED = auto()
1641 ... GREEN = auto()
1642 ... BLUE = auto()
1643 >>> class CheckedColor(Enum):
1644 ... RED = auto()
1645 ... GREEN = auto()
1646 ... BLUE = auto()
1647 >>> _test_simple_enum(CheckedColor, Color)
1648
1649 If differences are found, a :exc:`TypeError` is raised.
1650 """
1651 failed = []
1652 if checked_enum.__dict__ != simple_enum.__dict__:
1653 checked_dict = checked_enum.__dict__
1654 checked_keys = list(checked_dict.keys())
1655 simple_dict = simple_enum.__dict__
1656 simple_keys = list(simple_dict.keys())
1657 member_names = set(
1658 list(checked_enum._member_map_.keys())
1659 + list(simple_enum._member_map_.keys())
1660 )
1661 for key in set(checked_keys + simple_keys):
1662 if key in ('__module__', '_member_map_', '_value2member_map_'):
1663 # keys known to be different
1664 continue
1665 elif key in member_names:
1666 # members are checked below
1667 continue
1668 elif key not in simple_keys:
1669 failed.append("missing key: %r" % (key, ))
1670 elif key not in checked_keys:
1671 failed.append("extra key: %r" % (key, ))
1672 else:
1673 checked_value = checked_dict[key]
1674 simple_value = simple_dict[key]
1675 if callable(checked_value):
1676 continue
1677 if key == '__doc__':
1678 # remove all spaces/tabs
1679 compressed_checked_value = checked_value.replace(' ','').replace('\t','')
1680 compressed_simple_value = simple_value.replace(' ','').replace('\t','')
1681 if compressed_checked_value != compressed_simple_value:
1682 failed.append("%r:\n %s\n %s" % (
1683 key,
1684 "checked -> %r" % (checked_value, ),
1685 "simple -> %r" % (simple_value, ),
1686 ))
1687 elif checked_value != simple_value:
1688 failed.append("%r:\n %s\n %s" % (
1689 key,
1690 "checked -> %r" % (checked_value, ),
1691 "simple -> %r" % (simple_value, ),
1692 ))
1693 failed.sort()
1694 for name in member_names:
1695 failed_member = []
1696 if name not in simple_keys:
1697 failed.append('missing member from simple enum: %r' % name)
1698 elif name not in checked_keys:
1699 failed.append('extra member in simple enum: %r' % name)
1700 else:
1701 checked_member_dict = checked_enum[name].__dict__
1702 checked_member_keys = list(checked_member_dict.keys())
1703 simple_member_dict = simple_enum[name].__dict__
1704 simple_member_keys = list(simple_member_dict.keys())
1705 for key in set(checked_member_keys + simple_member_keys):
Ethan Furman6c681e12021-04-23 19:08:22 -07001706 if key in ('__module__', '__objclass__', '_inverted_'):
1707 # keys known to be different or absent
Ethan Furmana02cb472021-04-21 10:20:44 -07001708 continue
1709 elif key not in simple_member_keys:
1710 failed_member.append("missing key %r not in the simple enum member %r" % (key, name))
1711 elif key not in checked_member_keys:
1712 failed_member.append("extra key %r in simple enum member %r" % (key, name))
1713 else:
1714 checked_value = checked_member_dict[key]
1715 simple_value = simple_member_dict[key]
1716 if checked_value != simple_value:
1717 failed_member.append("%r:\n %s\n %s" % (
1718 key,
1719 "checked member -> %r" % (checked_value, ),
1720 "simple member -> %r" % (simple_value, ),
1721 ))
1722 if failed_member:
1723 failed.append('%r member mismatch:\n %s' % (
1724 name, '\n '.join(failed_member),
1725 ))
1726 for method in (
1727 '__str__', '__repr__', '__reduce_ex__', '__format__',
1728 '__getnewargs_ex__', '__getnewargs__', '__reduce_ex__', '__reduce__'
1729 ):
1730 if method in simple_keys and method in checked_keys:
1731 # cannot compare functions, and it exists in both, so we're good
1732 continue
1733 elif method not in simple_keys and method not in checked_keys:
1734 # method is inherited -- check it out
1735 checked_method = getattr(checked_enum, method, None)
1736 simple_method = getattr(simple_enum, method, None)
1737 if hasattr(checked_method, '__func__'):
1738 checked_method = checked_method.__func__
1739 simple_method = simple_method.__func__
1740 if checked_method != simple_method:
1741 failed.append("%r: %-30s %s" % (
1742 method,
1743 "checked -> %r" % (checked_method, ),
1744 "simple -> %r" % (simple_method, ),
1745 ))
1746 else:
1747 # if the method existed in only one of the enums, it will have been caught
1748 # in the first checks above
1749 pass
1750 if failed:
1751 raise TypeError('enum mismatch:\n %s' % '\n '.join(failed))
1752
1753def _old_convert_(etype, name, module, filter, source=None, *, boundary=None):
1754 """
1755 Create a new Enum subclass that replaces a collection of global constants
1756 """
1757 # convert all constants from source (or module) that pass filter() to
1758 # a new Enum called name, and export the enum and its members back to
1759 # module;
1760 # also, replace the __reduce_ex__ method so unpickling works in
1761 # previous Python versions
1762 module_globals = sys.modules[module].__dict__
1763 if source:
1764 source = source.__dict__
1765 else:
1766 source = module_globals
1767 # _value2member_map_ is populated in the same order every time
1768 # for a consistent reverse mapping of number to name when there
1769 # are multiple names for the same number.
1770 members = [
1771 (name, value)
1772 for name, value in source.items()
1773 if filter(name)]
1774 try:
1775 # sort by value
1776 members.sort(key=lambda t: (t[1], t[0]))
1777 except TypeError:
1778 # unless some values aren't comparable, in which case sort by name
1779 members.sort(key=lambda t: t[0])
1780 cls = etype(name, members, module=module, boundary=boundary or KEEP)
Miss Islington (bot)b6131322021-06-10 16:37:27 -07001781 cls.__reduce_ex__ = _reduce_ex_by_global_name
Ethan Furmana02cb472021-04-21 10:20:44 -07001782 cls.__repr__ = global_enum_repr
1783 return cls