blob: 49c46ea86dbacfa2b111fbb25c1ab588209a4a63 [file] [log] [blame]
Ethan Furman6b3d64a2013-06-14 16:55:46 -07001import sys
Ethan Furmane03ea372013-09-25 07:14:41 -07002from types import MappingProxyType, DynamicClassAttribute
Miss Islington (bot)0a186b12021-06-11 02:58:57 -07003from operator import or_ as _or_
4from functools import reduce
Ethan Furman7aaeb2a2021-01-25 14:26:19 -08005from builtins import property as _bltin_property, bin as _bltin_bin
Ethan Furman6b3d64a2013-06-14 16:55:46 -07006
Ethan Furmane5754ab2015-09-17 22:03:52 -07007
Ethan Furmanc16595e2016-09-10 23:36:59 -07008__all__ = [
Ethan Furmanb7751062021-03-30 21:17:26 -07009 'EnumType', 'EnumMeta',
Ethan Furman0063ff42020-09-21 17:23:13 -070010 'Enum', 'IntEnum', 'StrEnum', 'Flag', 'IntFlag',
Ethan Furman74964862021-06-10 07:24:20 -070011 'auto', 'unique', 'property', 'verify',
Ethan Furman7aaeb2a2021-01-25 14:26:19 -080012 'FlagBoundary', 'STRICT', 'CONFORM', 'EJECT', 'KEEP',
Ethan Furmanb7751062021-03-30 21:17:26 -070013 'global_flag_repr', 'global_enum_repr', 'global_enum',
Ethan Furman74964862021-06-10 07:24:20 -070014 'EnumCheck', 'CONTINUOUS', 'NAMED_FLAGS', 'UNIQUE',
Ethan Furmanc16595e2016-09-10 23:36:59 -070015 ]
Ethan Furman6b3d64a2013-06-14 16:55:46 -070016
17
Ethan Furman7aaeb2a2021-01-25 14:26:19 -080018# Dummy value for Enum and Flag as there are explicit checks for them
19# before they have been created.
Ethan Furmanb7751062021-03-30 21:17:26 -070020# This is also why there are checks in EnumType like `if Enum is not None`
Ethan Furman7aaeb2a2021-01-25 14:26:19 -080021Enum = Flag = EJECT = None
22
Ethan Furman101e0742013-09-15 12:34:36 -070023def _is_descriptor(obj):
Ethan Furman6d3dfee2020-12-08 12:26:56 -080024 """
25 Returns True if obj is a descriptor, False otherwise.
26 """
Ethan Furman101e0742013-09-15 12:34:36 -070027 return (
28 hasattr(obj, '__get__') or
29 hasattr(obj, '__set__') or
Ethan Furman6d3dfee2020-12-08 12:26:56 -080030 hasattr(obj, '__delete__')
31 )
Ethan Furman101e0742013-09-15 12:34:36 -070032
Ethan Furman6b3d64a2013-06-14 16:55:46 -070033def _is_dunder(name):
Ethan Furman6d3dfee2020-12-08 12:26:56 -080034 """
35 Returns True if a __dunder__ name, False otherwise.
36 """
37 return (
38 len(name) > 4 and
Brennan D Baraban8b914d22019-03-03 14:09:11 -080039 name[:2] == name[-2:] == '__' and
40 name[2] != '_' and
Ethan Furman6d3dfee2020-12-08 12:26:56 -080041 name[-3] != '_'
42 )
Ethan Furman6b3d64a2013-06-14 16:55:46 -070043
44def _is_sunder(name):
Ethan Furman6d3dfee2020-12-08 12:26:56 -080045 """
46 Returns True if a _sunder_ name, False otherwise.
47 """
48 return (
49 len(name) > 2 and
Brennan D Baraban8b914d22019-03-03 14:09:11 -080050 name[0] == name[-1] == '_' and
Ethan Furman6b3d64a2013-06-14 16:55:46 -070051 name[1:2] != '_' and
Ethan Furman6d3dfee2020-12-08 12:26:56 -080052 name[-2:-1] != '_'
53 )
Ethan Furman6b3d64a2013-06-14 16:55:46 -070054
Ethan Furman7cf0aad2020-12-09 17:12:11 -080055def _is_private(cls_name, name):
56 # do not use `re` as `re` imports `enum`
57 pattern = '_%s__' % (cls_name, )
Ethan Furmanec099732021-04-15 06:58:33 -070058 pat_len = len(pattern)
Ethan Furman7cf0aad2020-12-09 17:12:11 -080059 if (
Ethan Furmanec099732021-04-15 06:58:33 -070060 len(name) > pat_len
Ethan Furman7cf0aad2020-12-09 17:12:11 -080061 and name.startswith(pattern)
Ethan Furmanec099732021-04-15 06:58:33 -070062 and name[pat_len:pat_len+1] != ['_']
Ethan Furman7cf0aad2020-12-09 17:12:11 -080063 and (name[-1] != '_' or name[-2] != '_')
64 ):
65 return True
66 else:
67 return False
68
Ethan Furman7aaeb2a2021-01-25 14:26:19 -080069def _is_single_bit(num):
70 """
71 True if only one bit set in num (should be an int)
72 """
73 if num == 0:
74 return False
75 num &= num - 1
76 return num == 0
77
Ethan Furmanc314e602021-01-12 23:47:57 -080078def _make_class_unpicklable(obj):
Ethan Furman6d3dfee2020-12-08 12:26:56 -080079 """
Ethan Furmanc314e602021-01-12 23:47:57 -080080 Make the given obj un-picklable.
81
Miss Islington (bot)f6cf38c2021-06-15 11:25:07 -070082 obj should be either a dictionary, or an Enum
Ethan Furman6d3dfee2020-12-08 12:26:56 -080083 """
Ethan Furmanca1b7942014-02-08 11:36:27 -080084 def _break_on_call_reduce(self, proto):
Ethan Furman6b3d64a2013-06-14 16:55:46 -070085 raise TypeError('%r cannot be pickled' % self)
Ethan Furmanc314e602021-01-12 23:47:57 -080086 if isinstance(obj, dict):
87 obj['__reduce_ex__'] = _break_on_call_reduce
88 obj['__module__'] = '<unknown>'
89 else:
90 setattr(obj, '__reduce_ex__', _break_on_call_reduce)
91 setattr(obj, '__module__', '<unknown>')
Ethan Furman6b3d64a2013-06-14 16:55:46 -070092
Ethan Furman7aaeb2a2021-01-25 14:26:19 -080093def _iter_bits_lsb(num):
Ethan Furman74964862021-06-10 07:24:20 -070094 # num must be an integer
95 if isinstance(num, Enum):
96 num = num.value
Ethan Furman7aaeb2a2021-01-25 14:26:19 -080097 while num:
98 b = num & (~num + 1)
99 yield b
100 num ^= b
101
Miss Islington (bot)0a186b12021-06-11 02:58:57 -0700102def show_flag_values(value):
103 return list(_iter_bits_lsb(value))
104
Ethan Furman7aaeb2a2021-01-25 14:26:19 -0800105def bin(num, max_bits=None):
106 """
107 Like built-in bin(), except negative values are represented in
108 twos-compliment, and the leading bit always indicates sign
109 (0=positive, 1=negative).
110
111 >>> bin(10)
112 '0b0 1010'
113 >>> bin(~10) # ~10 is -11
114 '0b1 0101'
115 """
116
117 ceiling = 2 ** (num).bit_length()
118 if num >= 0:
119 s = _bltin_bin(num + ceiling).replace('1', '0', 1)
120 else:
121 s = _bltin_bin(~num ^ (ceiling - 1) + ceiling)
122 sign = s[:3]
123 digits = s[3:]
124 if max_bits is not None:
125 if len(digits) < max_bits:
126 digits = (sign[-1] * max_bits + digits)[-max_bits:]
127 return "%s %s" % (sign, digits)
128
129
Ethan Furman3515dcc2016-09-18 13:15:41 -0700130_auto_null = object()
Ethan Furmanc16595e2016-09-10 23:36:59 -0700131class auto:
132 """
133 Instances are replaced with an appropriate value in Enum class suites.
134 """
Ethan Furman3515dcc2016-09-18 13:15:41 -0700135 value = _auto_null
Ethan Furmanc16595e2016-09-10 23:36:59 -0700136
Ethan Furmanc314e602021-01-12 23:47:57 -0800137class property(DynamicClassAttribute):
138 """
139 This is a descriptor, used to define attributes that act differently
140 when accessed through an enum member and through an enum class.
141 Instance access is the same as property(), but access to an attribute
142 through the enum class will instead look in the class' _member_map_ for
143 a corresponding enum member.
144 """
145
146 def __get__(self, instance, ownerclass=None):
147 if instance is None:
148 try:
149 return ownerclass._member_map_[self.name]
150 except KeyError:
Ethan Furman7aaeb2a2021-01-25 14:26:19 -0800151 raise AttributeError(
Ethan Furmand65b9032021-02-08 17:32:38 -0800152 '%s: no class attribute %r' % (ownerclass.__name__, self.name)
Ethan Furman7aaeb2a2021-01-25 14:26:19 -0800153 )
Ethan Furmanc314e602021-01-12 23:47:57 -0800154 else:
155 if self.fget is None:
Ethan Furmand65b9032021-02-08 17:32:38 -0800156 # check for member
157 if self.name in ownerclass._member_map_:
158 import warnings
159 warnings.warn(
160 "accessing one member from another is not supported, "
Ethan Furman44e580f2021-03-03 09:54:30 -0800161 " and will be disabled in 3.12",
Ethan Furmand65b9032021-02-08 17:32:38 -0800162 DeprecationWarning,
163 stacklevel=2,
164 )
165 return ownerclass._member_map_[self.name]
Ethan Furman7aaeb2a2021-01-25 14:26:19 -0800166 raise AttributeError(
Ethan Furmand65b9032021-02-08 17:32:38 -0800167 '%s: no instance attribute %r' % (ownerclass.__name__, self.name)
Ethan Furman7aaeb2a2021-01-25 14:26:19 -0800168 )
Ethan Furmanc314e602021-01-12 23:47:57 -0800169 else:
170 return self.fget(instance)
171
172 def __set__(self, instance, value):
173 if self.fset is None:
Ethan Furman7aaeb2a2021-01-25 14:26:19 -0800174 raise AttributeError(
Ethan Furmand65b9032021-02-08 17:32:38 -0800175 "%s: cannot set 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.fset(instance, value)
179
180 def __delete__(self, instance):
181 if self.fdel is None:
Ethan Furman7aaeb2a2021-01-25 14:26:19 -0800182 raise AttributeError(
Ethan Furmand65b9032021-02-08 17:32:38 -0800183 "%s: cannot delete instance attribute %r" % (self.clsname, self.name)
Ethan Furman7aaeb2a2021-01-25 14:26:19 -0800184 )
Ethan Furmanc314e602021-01-12 23:47:57 -0800185 else:
186 return self.fdel(instance)
187
188 def __set_name__(self, ownerclass, name):
189 self.name = name
190 self.clsname = ownerclass.__name__
191
192
193class _proto_member:
194 """
195 intermediate step for enum members between class execution and final creation
196 """
197
198 def __init__(self, value):
199 self.value = value
200
201 def __set_name__(self, enum_class, member_name):
202 """
203 convert each quasi-member into an instance of the new enum class
204 """
205 # first step: remove ourself from enum_class
206 delattr(enum_class, member_name)
207 # second step: create member based on enum_class
208 value = self.value
209 if not isinstance(value, tuple):
210 args = (value, )
211 else:
212 args = value
213 if enum_class._member_type_ is tuple: # special case for tuple enums
214 args = (args, ) # wrap it one more time
215 if not enum_class._use_args_:
216 enum_member = enum_class._new_member_(enum_class)
217 if not hasattr(enum_member, '_value_'):
218 enum_member._value_ = value
219 else:
220 enum_member = enum_class._new_member_(enum_class, *args)
221 if not hasattr(enum_member, '_value_'):
222 if enum_class._member_type_ is object:
223 enum_member._value_ = value
224 else:
Ethan Furman7aaeb2a2021-01-25 14:26:19 -0800225 try:
226 enum_member._value_ = enum_class._member_type_(*args)
227 except Exception as exc:
228 raise TypeError(
229 '_value_ not set in __new__, unable to create it'
230 ) from None
Ethan Furmanc314e602021-01-12 23:47:57 -0800231 value = enum_member._value_
232 enum_member._name_ = member_name
233 enum_member.__objclass__ = enum_class
234 enum_member.__init__(*args)
Ethan Furman7aaeb2a2021-01-25 14:26:19 -0800235 enum_member._sort_order_ = len(enum_class._member_names_)
Ethan Furmanc314e602021-01-12 23:47:57 -0800236 # If another member with the same value was already defined, the
237 # new member becomes an alias to the existing one.
238 for name, canonical_member in enum_class._member_map_.items():
239 if canonical_member._value_ == enum_member._value_:
240 enum_member = canonical_member
241 break
242 else:
Ethan Furman7aaeb2a2021-01-25 14:26:19 -0800243 # this could still be an alias if the value is multi-bit and the
244 # class is a flag class
245 if (
246 Flag is None
247 or not issubclass(enum_class, Flag)
248 ):
249 # no other instances found, record this member in _member_names_
250 enum_class._member_names_.append(member_name)
251 elif (
252 Flag is not None
253 and issubclass(enum_class, Flag)
254 and _is_single_bit(value)
255 ):
256 # no other instances found, record this member in _member_names_
257 enum_class._member_names_.append(member_name)
Ethan Furmanc314e602021-01-12 23:47:57 -0800258 # get redirect in place before adding to _member_map_
259 # but check for other instances in parent classes first
260 need_override = False
261 descriptor = None
262 for base in enum_class.__mro__[1:]:
263 descriptor = base.__dict__.get(member_name)
264 if descriptor is not None:
265 if isinstance(descriptor, (property, DynamicClassAttribute)):
266 break
267 else:
268 need_override = True
269 # keep looking for an enum.property
270 if descriptor and not need_override:
271 # previous enum.property found, no further action needed
272 pass
273 else:
274 redirect = property()
275 redirect.__set_name__(enum_class, member_name)
276 if descriptor and need_override:
277 # previous enum.property found, but some other inherited attribute
278 # is in the way; copy fget, fset, fdel to this one
279 redirect.fget = descriptor.fget
280 redirect.fset = descriptor.fset
281 redirect.fdel = descriptor.fdel
282 setattr(enum_class, member_name, redirect)
283 # now add to _member_map_ (even aliases)
284 enum_class._member_map_[member_name] = enum_member
285 try:
286 # This may fail if value is not hashable. We can't add the value
287 # to the map, and by-value lookups for this value will be
288 # linear.
Ethan Furman7aaeb2a2021-01-25 14:26:19 -0800289 enum_class._value2member_map_.setdefault(value, enum_member)
Ethan Furmanc314e602021-01-12 23:47:57 -0800290 except TypeError:
Ethan Furman6bd92882021-04-27 13:05:08 -0700291 # keep track of the value in a list so containment checks are quick
292 enum_class._unhashable_values_.append(value)
Ethan Furmanc314e602021-01-12 23:47:57 -0800293
Ethan Furman101e0742013-09-15 12:34:36 -0700294
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700295class _EnumDict(dict):
Ethan Furman6d3dfee2020-12-08 12:26:56 -0800296 """
297 Track enum member order and ensure member names are not reused.
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700298
Ethan Furmanb7751062021-03-30 21:17:26 -0700299 EnumType will use the names found in self._member_names as the
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700300 enumeration member names.
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700301 """
302 def __init__(self):
303 super().__init__()
304 self._member_names = []
Ethan Furmanc16595e2016-09-10 23:36:59 -0700305 self._last_values = []
Ethan Furmana4b1bb42018-01-22 07:56:37 -0800306 self._ignore = []
Ethan Onstottd9a43e22020-04-28 13:20:55 -0400307 self._auto_called = False
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700308
309 def __setitem__(self, key, value):
Ethan Furman6d3dfee2020-12-08 12:26:56 -0800310 """
311 Changes anything not dundered or not a descriptor.
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700312
313 If an enum member name is used twice, an error is raised; duplicate
314 values are not checked for.
315
316 Single underscore (sunder) names are reserved.
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700317 """
Ethan Furman7cf0aad2020-12-09 17:12:11 -0800318 if _is_private(self._cls_name, key):
319 # do nothing, name will be a normal attribute
320 pass
321 elif _is_sunder(key):
Ethan Furmanee47e5c2016-08-31 00:12:15 -0700322 if key not in (
Ethan Furman0dca5eb2021-04-15 06:49:54 -0700323 '_order_',
Ethan Furmana4b1bb42018-01-22 07:56:37 -0800324 '_generate_next_value_', '_missing_', '_ignore_',
Ethan Furman7aaeb2a2021-01-25 14:26:19 -0800325 '_iter_member_', '_iter_member_by_value_', '_iter_member_by_def_',
Ethan Furmanee47e5c2016-08-31 00:12:15 -0700326 ):
Ethan Furman6d3dfee2020-12-08 12:26:56 -0800327 raise ValueError(
328 '_sunder_ names, such as %r, are reserved for future Enum use'
329 % (key, )
330 )
Ethan Furmanc16595e2016-09-10 23:36:59 -0700331 if key == '_generate_next_value_':
Ethan Onstottd9a43e22020-04-28 13:20:55 -0400332 # check if members already defined as auto()
333 if self._auto_called:
334 raise TypeError("_generate_next_value_ must be defined before members")
Ethan Furmanb7751062021-03-30 21:17:26 -0700335 _gnv = value.__func__ if isinstance(value, staticmethod) else value
336 setattr(self, '_generate_next_value', _gnv)
Ethan Furmana4b1bb42018-01-22 07:56:37 -0800337 elif key == '_ignore_':
338 if isinstance(value, str):
339 value = value.replace(',',' ').split()
340 else:
341 value = list(value)
342 self._ignore = value
343 already = set(value) & set(self._member_names)
344 if already:
Ethan Furman6d3dfee2020-12-08 12:26:56 -0800345 raise ValueError(
346 '_ignore_ cannot specify already set names: %r'
347 % (already, )
348 )
Ethan Furman101e0742013-09-15 12:34:36 -0700349 elif _is_dunder(key):
Ethan Furmane8e61272016-08-20 07:19:31 -0700350 if key == '__order__':
351 key = '_order_'
Ethan Furman101e0742013-09-15 12:34:36 -0700352 elif key in self._member_names:
353 # descriptor overwriting an enum?
Ethan Furmana6582872020-12-10 13:07:00 -0800354 raise TypeError('%r already defined as: %r' % (key, self[key]))
Ethan Furmana4b1bb42018-01-22 07:56:37 -0800355 elif key in self._ignore:
356 pass
Ethan Furman101e0742013-09-15 12:34:36 -0700357 elif not _is_descriptor(value):
358 if key in self:
359 # enum overwriting a descriptor?
Ethan Furmanee47e5c2016-08-31 00:12:15 -0700360 raise TypeError('%r already defined as: %r' % (key, self[key]))
Ethan Furmanc16595e2016-09-10 23:36:59 -0700361 if isinstance(value, auto):
Ethan Furman3515dcc2016-09-18 13:15:41 -0700362 if value.value == _auto_null:
Ethan Furman6d3dfee2020-12-08 12:26:56 -0800363 value.value = self._generate_next_value(
Ethan Furman7aaeb2a2021-01-25 14:26:19 -0800364 key, 1, len(self._member_names), self._last_values[:],
Ethan Furman6d3dfee2020-12-08 12:26:56 -0800365 )
Ethan Furmanfc23a942020-09-16 12:37:54 -0700366 self._auto_called = True
Ethan Furman3515dcc2016-09-18 13:15:41 -0700367 value = value.value
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700368 self._member_names.append(key)
Ethan Furmanc16595e2016-09-10 23:36:59 -0700369 self._last_values.append(value)
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700370 super().__setitem__(key, value)
371
Ethan Furmana6582872020-12-10 13:07:00 -0800372 def update(self, members, **more_members):
373 try:
374 for name in members.keys():
375 self[name] = members[name]
376 except AttributeError:
377 for name, value in members:
378 self[name] = value
379 for name, value in more_members.items():
380 self[name] = value
381
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700382
Ethan Furmanb7751062021-03-30 21:17:26 -0700383class EnumType(type):
Ethan Furman6d3dfee2020-12-08 12:26:56 -0800384 """
385 Metaclass for Enum
386 """
Ethan Furman7aaeb2a2021-01-25 14:26:19 -0800387
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700388 @classmethod
Ethan Furman6ec0ade2020-12-24 10:05:02 -0800389 def __prepare__(metacls, cls, bases, **kwds):
Ethan Furman3064dbf2020-09-16 07:11:57 -0700390 # check that previous enum members do not exist
391 metacls._check_for_existing_members(cls, bases)
Ethan Furmanee47e5c2016-08-31 00:12:15 -0700392 # create the namespace dict
393 enum_dict = _EnumDict()
Ethan Furman7cf0aad2020-12-09 17:12:11 -0800394 enum_dict._cls_name = cls
Ethan Furmanee47e5c2016-08-31 00:12:15 -0700395 # inherit previous flags and _generate_next_value_ function
Ethan Furman3064dbf2020-09-16 07:11:57 -0700396 member_type, first_enum = metacls._get_mixins_(cls, bases)
Ethan Furmanee47e5c2016-08-31 00:12:15 -0700397 if first_enum is not None:
Ethan Furman6d3dfee2020-12-08 12:26:56 -0800398 enum_dict['_generate_next_value_'] = getattr(
399 first_enum, '_generate_next_value_', None,
400 )
Ethan Furmanee47e5c2016-08-31 00:12:15 -0700401 return enum_dict
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700402
Ethan Furmana02cb472021-04-21 10:20:44 -0700403 def __new__(metacls, cls, bases, classdict, *, boundary=None, _simple=False, **kwds):
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700404 # an Enum class is final once enumeration items have been defined; it
405 # cannot be mixed with other types (int, float, etc.) if it has an
406 # inherited __new__ unless a new __new__ is defined (or the resulting
407 # class will fail).
Ethan Furmana4b1bb42018-01-22 07:56:37 -0800408 #
409 # remove any keys listed in _ignore_
Ethan Furmana02cb472021-04-21 10:20:44 -0700410 if _simple:
411 return super().__new__(metacls, cls, bases, classdict, **kwds)
Ethan Furmana4b1bb42018-01-22 07:56:37 -0800412 classdict.setdefault('_ignore_', []).append('_ignore_')
413 ignore = classdict['_ignore_']
414 for key in ignore:
415 classdict.pop(key, None)
Ethan Furmanc314e602021-01-12 23:47:57 -0800416 #
417 # grab member names
418 member_names = classdict._member_names
419 #
420 # check for illegal enum names (any others?)
421 invalid_names = set(member_names) & {'mro', ''}
422 if invalid_names:
423 raise ValueError('Invalid enum member name: {0}'.format(
424 ','.join(invalid_names)))
425 #
426 # adjust the sunders
427 _order_ = classdict.pop('_order_', None)
428 # convert to normal dict
429 classdict = dict(classdict.items())
430 #
431 # data type of member and the controlling Enum class
Ethan Furman3064dbf2020-09-16 07:11:57 -0700432 member_type, first_enum = metacls._get_mixins_(cls, bases)
Ethan Furmanc2667362020-12-07 00:17:31 -0800433 __new__, save_new, use_args = metacls._find_new_(
434 classdict, member_type, first_enum,
435 )
Ethan Furmanc314e602021-01-12 23:47:57 -0800436 classdict['_new_member_'] = __new__
437 classdict['_use_args_'] = use_args
438 #
439 # convert future enum members into temporary _proto_members
Ethan Furman7aaeb2a2021-01-25 14:26:19 -0800440 # and record integer values in case this will be a Flag
441 flag_mask = 0
Ethan Furmanc314e602021-01-12 23:47:57 -0800442 for name in member_names:
Ethan Furman7aaeb2a2021-01-25 14:26:19 -0800443 value = classdict[name]
444 if isinstance(value, int):
445 flag_mask |= value
446 classdict[name] = _proto_member(value)
Ethan Furmanc314e602021-01-12 23:47:57 -0800447 #
Ethan Furman7aaeb2a2021-01-25 14:26:19 -0800448 # house-keeping structures
Ethan Furmanc314e602021-01-12 23:47:57 -0800449 classdict['_member_names_'] = []
450 classdict['_member_map_'] = {}
451 classdict['_value2member_map_'] = {}
Ethan Furman6bd92882021-04-27 13:05:08 -0700452 classdict['_unhashable_values_'] = []
Ethan Furmanc314e602021-01-12 23:47:57 -0800453 classdict['_member_type_'] = member_type
454 #
Ethan Furman7aaeb2a2021-01-25 14:26:19 -0800455 # Flag structures (will be removed if final class is not a Flag
456 classdict['_boundary_'] = (
457 boundary
458 or getattr(first_enum, '_boundary_', None)
459 )
460 classdict['_flag_mask_'] = flag_mask
461 classdict['_all_bits_'] = 2 ** ((flag_mask).bit_length()) - 1
462 classdict['_inverted_'] = None
463 #
Ethan Furmanc314e602021-01-12 23:47:57 -0800464 # create a default docstring if one has not been provided
465 if '__doc__' not in classdict:
466 classdict['__doc__'] = 'An enumeration.'
467 try:
468 exc = None
469 enum_class = super().__new__(metacls, cls, bases, classdict, **kwds)
470 except RuntimeError as e:
471 # any exceptions raised by member.__new__ will get converted to a
472 # RuntimeError, so get that original exception back and raise it instead
473 exc = e.__cause__ or e
474 if exc is not None:
475 raise exc
476 #
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700477 # double check that repr and friends are not the mixin's or various
478 # things break (such as pickle)
Ethan Furman22415ad2020-09-15 16:28:25 -0700479 # however, if the method is defined in the Enum itself, don't replace
480 # it
Ethan Furmandc870522014-02-18 12:37:12 -0800481 for name in ('__repr__', '__str__', '__format__', '__reduce_ex__'):
Ethan Furman22415ad2020-09-15 16:28:25 -0700482 if name in classdict:
483 continue
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700484 class_method = getattr(enum_class, name)
485 obj_method = getattr(member_type, name, None)
486 enum_method = getattr(first_enum, name, None)
487 if obj_method is not None and obj_method is class_method:
488 setattr(enum_class, name, enum_method)
Ethan Furmanc314e602021-01-12 23:47:57 -0800489 #
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700490 # replace any other __new__ with our own (as long as Enum is not None,
491 # anyway) -- again, this is to support pickle
492 if Enum is not None:
493 # if the user defined their own __new__, save it before it gets
494 # clobbered in case they subclass later
495 if save_new:
496 enum_class.__new_member__ = __new__
497 enum_class.__new__ = Enum.__new__
Ethan Furmanc314e602021-01-12 23:47:57 -0800498 #
Ethan Furmane8e61272016-08-20 07:19:31 -0700499 # py3 support for definition order (helps keep py2/py3 code in sync)
Ethan Furman7aaeb2a2021-01-25 14:26:19 -0800500 #
501 # _order_ checking is spread out into three/four steps
502 # - if enum_class is a Flag:
503 # - remove any non-single-bit flags from _order_
504 # - remove any aliases from _order_
505 # - check that _order_ and _member_names_ match
506 #
507 # step 1: ensure we have a list
Ethan Furmane8e61272016-08-20 07:19:31 -0700508 if _order_ is not None:
509 if isinstance(_order_, str):
510 _order_ = _order_.replace(',', ' ').split()
Ethan Furman7aaeb2a2021-01-25 14:26:19 -0800511 #
512 # remove Flag structures if final class is not a Flag
513 if (
514 Flag is None and cls != 'Flag'
515 or Flag is not None and not issubclass(enum_class, Flag)
516 ):
517 delattr(enum_class, '_boundary_')
518 delattr(enum_class, '_flag_mask_')
519 delattr(enum_class, '_all_bits_')
520 delattr(enum_class, '_inverted_')
521 elif Flag is not None and issubclass(enum_class, Flag):
522 # ensure _all_bits_ is correct and there are no missing flags
523 single_bit_total = 0
524 multi_bit_total = 0
525 for flag in enum_class._member_map_.values():
526 flag_value = flag._value_
527 if _is_single_bit(flag_value):
528 single_bit_total |= flag_value
529 else:
530 # multi-bit flags are considered aliases
531 multi_bit_total |= flag_value
Ethan Furman7aaeb2a2021-01-25 14:26:19 -0800532 enum_class._flag_mask_ = single_bit_total
533 #
534 # set correct __iter__
535 member_list = [m._value_ for m in enum_class]
536 if member_list != sorted(member_list):
537 enum_class._iter_member_ = enum_class._iter_member_by_def_
538 if _order_:
539 # _order_ step 2: remove any items from _order_ that are not single-bit
540 _order_ = [
541 o
542 for o in _order_
543 if o not in enum_class._member_map_ or _is_single_bit(enum_class[o]._value_)
544 ]
545 #
546 if _order_:
547 # _order_ step 3: remove aliases from _order_
548 _order_ = [
549 o
550 for o in _order_
551 if (
552 o not in enum_class._member_map_
553 or
554 (o in enum_class._member_map_ and o in enum_class._member_names_)
555 )]
556 # _order_ step 4: verify that _order_ and _member_names_ match
Ethan Furmane8e61272016-08-20 07:19:31 -0700557 if _order_ != enum_class._member_names_:
Ethan Furman7aaeb2a2021-01-25 14:26:19 -0800558 raise TypeError(
559 'member order does not match _order_:\n%r\n%r'
560 % (enum_class._member_names_, _order_)
561 )
Ethan Furmanc314e602021-01-12 23:47:57 -0800562 #
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700563 return enum_class
564
Ethan Furman5de67b12016-04-13 23:52:09 -0700565 def __bool__(self):
566 """
567 classes/types should always be True.
568 """
569 return True
570
Ethan Furman7aaeb2a2021-01-25 14:26:19 -0800571 def __call__(cls, value, names=None, *, module=None, qualname=None, type=None, start=1, boundary=None):
Ethan Furman6d3dfee2020-12-08 12:26:56 -0800572 """
573 Either returns an existing member, or creates a new enum class.
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700574
575 This method is used both when an enum class is given a value to match
576 to an enumeration member (i.e. Color(3)) and for the functional API
Ethan Furman23bb6f42016-11-21 09:22:05 -0800577 (i.e. Color = Enum('Color', names='RED GREEN BLUE')).
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700578
Ethan Furman2da95042014-03-03 12:42:52 -0800579 When used for the functional API:
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700580
Ethan Furman2da95042014-03-03 12:42:52 -0800581 `value` will be the name of the new class.
582
583 `names` should be either a string of white-space/comma delimited names
Ethan Furmand9925a12014-09-16 20:35:55 -0700584 (values will start at `start`), or an iterator/mapping of name, value pairs.
Ethan Furman2da95042014-03-03 12:42:52 -0800585
586 `module` should be set to the module this class is being created in;
587 if it is not set, an attempt to find that module will be made, but if
588 it fails the class will not be picklable.
589
590 `qualname` should be set to the actual location this class can be found
591 at in its module; by default it is set to the global scope. If this is
592 not correct, unpickling will fail in some circumstances.
593
594 `type`, if set, will be mixed in as the first base class.
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700595 """
596 if names is None: # simple value lookup
597 return cls.__new__(cls, value)
598 # otherwise, functional API: we're creating a new Enum type
Ethan Furman6d3dfee2020-12-08 12:26:56 -0800599 return cls._create_(
600 value,
601 names,
602 module=module,
603 qualname=qualname,
604 type=type,
605 start=start,
Ethan Furman7aaeb2a2021-01-25 14:26:19 -0800606 boundary=boundary,
Ethan Furman6d3dfee2020-12-08 12:26:56 -0800607 )
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700608
609 def __contains__(cls, member):
Rahul Jha94306522018-09-10 23:51:04 +0530610 if not isinstance(member, Enum):
Ethan Furman6bd92882021-04-27 13:05:08 -0700611 import warnings
612 warnings.warn(
613 "in 3.12 __contains__ will no longer raise TypeError, but will return True or\n"
614 "False depending on whether the value is a member or the value of a member",
615 DeprecationWarning,
616 stacklevel=2,
617 )
Rahul Jha94306522018-09-10 23:51:04 +0530618 raise TypeError(
619 "unsupported operand type(s) for 'in': '%s' and '%s'" % (
620 type(member).__qualname__, cls.__class__.__qualname__))
Ethan Furman0081f232014-09-16 17:31:23 -0700621 return isinstance(member, cls) and member._name_ in cls._member_map_
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700622
Ethan Furman64a99722013-09-22 16:18:19 -0700623 def __delattr__(cls, attr):
624 # nicer error message when someone tries to delete an attribute
625 # (see issue19025).
626 if attr in cls._member_map_:
Ethan Furman6d3dfee2020-12-08 12:26:56 -0800627 raise AttributeError("%s: cannot delete Enum member %r." % (cls.__name__, attr))
Ethan Furman64a99722013-09-22 16:18:19 -0700628 super().__delattr__(attr)
629
Ethan Furman388a3922013-08-12 06:51:41 -0700630 def __dir__(self):
Ethan Furman6d3dfee2020-12-08 12:26:56 -0800631 return (
632 ['__class__', '__doc__', '__members__', '__module__']
633 + self._member_names_
634 )
Ethan Furman388a3922013-08-12 06:51:41 -0700635
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700636 def __getattr__(cls, name):
Ethan Furman6d3dfee2020-12-08 12:26:56 -0800637 """
638 Return the enum member matching `name`
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700639
640 We use __getattr__ instead of descriptors or inserting into the enum
641 class' __dict__ in order to support `name` and `value` being both
642 properties for enum members (which live in the class' __dict__) and
643 enum members themselves.
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700644 """
645 if _is_dunder(name):
646 raise AttributeError(name)
647 try:
Ethan Furman520ad572013-07-19 19:47:21 -0700648 return cls._member_map_[name]
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700649 except KeyError:
650 raise AttributeError(name) from None
651
652 def __getitem__(cls, name):
Ethan Furman520ad572013-07-19 19:47:21 -0700653 return cls._member_map_[name]
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700654
655 def __iter__(cls):
Ethan Furman6d3dfee2020-12-08 12:26:56 -0800656 """
657 Returns members in definition order.
658 """
Ethan Furman520ad572013-07-19 19:47:21 -0700659 return (cls._member_map_[name] for name in cls._member_names_)
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700660
661 def __len__(cls):
Ethan Furman520ad572013-07-19 19:47:21 -0700662 return len(cls._member_names_)
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700663
Ethan Furmanc314e602021-01-12 23:47:57 -0800664 @_bltin_property
Ethan Furman2131a4a2013-09-14 18:11:24 -0700665 def __members__(cls):
Ethan Furman6d3dfee2020-12-08 12:26:56 -0800666 """
667 Returns a mapping of member name->value.
Ethan Furman2131a4a2013-09-14 18:11:24 -0700668
669 This mapping lists all enum members, including aliases. Note that this
670 is a read-only view of the internal mapping.
Ethan Furman2131a4a2013-09-14 18:11:24 -0700671 """
672 return MappingProxyType(cls._member_map_)
673
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700674 def __repr__(cls):
Ethan Furman74964862021-06-10 07:24:20 -0700675 if Flag is not None and issubclass(cls, Flag):
676 return "<flag %r>" % cls.__name__
677 else:
678 return "<enum %r>" % cls.__name__
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700679
Ethan Furman2131a4a2013-09-14 18:11:24 -0700680 def __reversed__(cls):
Ethan Furman6d3dfee2020-12-08 12:26:56 -0800681 """
682 Returns members in reverse definition order.
683 """
Ethan Furman2131a4a2013-09-14 18:11:24 -0700684 return (cls._member_map_[name] for name in reversed(cls._member_names_))
685
Ethan Furmanf203f2d2013-09-06 07:16:48 -0700686 def __setattr__(cls, name, value):
Ethan Furman6d3dfee2020-12-08 12:26:56 -0800687 """
688 Block attempts to reassign Enum members.
Ethan Furmanf203f2d2013-09-06 07:16:48 -0700689
690 A simple assignment to the class namespace only changes one of the
691 several possible ways to get an Enum member from the Enum class,
692 resulting in an inconsistent Enumeration.
Ethan Furmanf203f2d2013-09-06 07:16:48 -0700693 """
694 member_map = cls.__dict__.get('_member_map_', {})
695 if name in member_map:
Ethan Furmana02cb472021-04-21 10:20:44 -0700696 raise AttributeError('Cannot reassign member %r.' % (name, ))
Ethan Furmanf203f2d2013-09-06 07:16:48 -0700697 super().__setattr__(name, value)
698
Ethan Furman7aaeb2a2021-01-25 14:26:19 -0800699 def _create_(cls, class_name, names, *, module=None, qualname=None, type=None, start=1, boundary=None):
Ethan Furman6d3dfee2020-12-08 12:26:56 -0800700 """
701 Convenience method to create a new Enum class.
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700702
703 `names` can be:
704
705 * A string containing member names, separated either with spaces or
Ethan Furmand9925a12014-09-16 20:35:55 -0700706 commas. Values are incremented by 1 from `start`.
707 * An iterable of member names. Values are incremented by 1 from `start`.
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700708 * An iterable of (member name, value) pairs.
Ethan Furmand9925a12014-09-16 20:35:55 -0700709 * A mapping of member name -> value pairs.
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700710 """
711 metacls = cls.__class__
712 bases = (cls, ) if type is None else (type, cls)
Ethan Furman3064dbf2020-09-16 07:11:57 -0700713 _, first_enum = cls._get_mixins_(cls, bases)
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700714 classdict = metacls.__prepare__(class_name, bases)
715
716 # special processing needed for names?
717 if isinstance(names, str):
718 names = names.replace(',', ' ').split()
Dong-hee Nadcc8ce42017-06-22 01:52:32 +0900719 if isinstance(names, (tuple, list)) and names and isinstance(names[0], str):
Ethan Furmanee47e5c2016-08-31 00:12:15 -0700720 original_names, names = names, []
Ethan Furmanc16595e2016-09-10 23:36:59 -0700721 last_values = []
Ethan Furmanee47e5c2016-08-31 00:12:15 -0700722 for count, name in enumerate(original_names):
Ethan Furmanc16595e2016-09-10 23:36:59 -0700723 value = first_enum._generate_next_value_(name, start, count, last_values[:])
724 last_values.append(value)
725 names.append((name, value))
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700726
727 # Here, names is either an iterable of (name, value) or a mapping.
728 for item in names:
729 if isinstance(item, str):
730 member_name, member_value = item, names[item]
731 else:
732 member_name, member_value = item
733 classdict[member_name] = member_value
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700734
735 # TODO: replace the frame hack if a blessed way to know the calling
736 # module is ever developed
737 if module is None:
738 try:
739 module = sys._getframe(2).f_globals['__name__']
Pablo Galindo293dd232019-11-19 21:34:03 +0000740 except (AttributeError, ValueError, KeyError):
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700741 pass
742 if module is None:
Ethan Furmanc314e602021-01-12 23:47:57 -0800743 _make_class_unpicklable(classdict)
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700744 else:
Ethan Furmanc314e602021-01-12 23:47:57 -0800745 classdict['__module__'] = module
Ethan Furmanca1b7942014-02-08 11:36:27 -0800746 if qualname is not None:
Ethan Furmanc314e602021-01-12 23:47:57 -0800747 classdict['__qualname__'] = qualname
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700748
Ethan Furman7aaeb2a2021-01-25 14:26:19 -0800749 return metacls.__new__(metacls, class_name, bases, classdict, boundary=boundary)
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700750
Ethan Furmana02cb472021-04-21 10:20:44 -0700751 def _convert_(cls, name, module, filter, source=None, *, boundary=None):
752
orlnub1230fb9fad2018-09-12 20:28:53 +0300753 """
754 Create a new Enum subclass that replaces a collection of global constants
755 """
756 # convert all constants from source (or module) that pass filter() to
757 # a new Enum called name, and export the enum and its members back to
758 # module;
759 # also, replace the __reduce_ex__ method so unpickling works in
760 # previous Python versions
Ethan Furmanb7751062021-03-30 21:17:26 -0700761 module_globals = sys.modules[module].__dict__
orlnub1230fb9fad2018-09-12 20:28:53 +0300762 if source:
Ethan Furmanb7751062021-03-30 21:17:26 -0700763 source = source.__dict__
orlnub1230fb9fad2018-09-12 20:28:53 +0300764 else:
765 source = module_globals
766 # _value2member_map_ is populated in the same order every time
767 # for a consistent reverse mapping of number to name when there
768 # are multiple names for the same number.
769 members = [
770 (name, value)
771 for name, value in source.items()
772 if filter(name)]
773 try:
774 # sort by value
775 members.sort(key=lambda t: (t[1], t[0]))
776 except TypeError:
777 # unless some values aren't comparable, in which case sort by name
778 members.sort(key=lambda t: t[0])
Ethan Furmana02cb472021-04-21 10:20:44 -0700779 body = {t[0]: t[1] for t in members}
780 body['__module__'] = module
781 tmp_cls = type(name, (object, ), body)
782 cls = _simple_enum(etype=cls, boundary=boundary or KEEP)(tmp_cls)
Miss Islington (bot)b6131322021-06-10 16:37:27 -0700783 cls.__reduce_ex__ = _reduce_ex_by_global_name
Ethan Furmanb7751062021-03-30 21:17:26 -0700784 global_enum(cls)
orlnub1230fb9fad2018-09-12 20:28:53 +0300785 module_globals[name] = cls
786 return cls
787
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700788 @staticmethod
Ethan Furman3064dbf2020-09-16 07:11:57 -0700789 def _check_for_existing_members(class_name, bases):
790 for chain in bases:
791 for base in chain.__mro__:
792 if issubclass(base, Enum) and base._member_names_:
Ethan Furman6d3dfee2020-12-08 12:26:56 -0800793 raise TypeError(
794 "%s: cannot extend enumeration %r"
795 % (class_name, base.__name__)
796 )
Ethan Furman3064dbf2020-09-16 07:11:57 -0700797
798 @staticmethod
799 def _get_mixins_(class_name, bases):
Ethan Furman6d3dfee2020-12-08 12:26:56 -0800800 """
801 Returns the type for creating enum members, and the first inherited
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700802 enum class.
803
804 bases: the tuple of bases that was given to __new__
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700805 """
806 if not bases:
807 return object, Enum
808
Ethan Furman5bdab642018-09-21 19:03:09 -0700809 def _find_data_type(bases):
Miss Islington (bot)01286012021-06-10 15:01:03 -0700810 data_types = set()
Ethan Furman5bdab642018-09-21 19:03:09 -0700811 for chain in bases:
Ethan Furmanbff01f32020-09-15 15:56:26 -0700812 candidate = None
Ethan Furman5bdab642018-09-21 19:03:09 -0700813 for base in chain.__mro__:
814 if base is object:
815 continue
Ethan Furmanc2667362020-12-07 00:17:31 -0800816 elif issubclass(base, Enum):
817 if base._member_type_ is not object:
Miss Islington (bot)01286012021-06-10 15:01:03 -0700818 data_types.add(base._member_type_)
Ethan Furmanc2667362020-12-07 00:17:31 -0800819 break
Ethan Furman5bdab642018-09-21 19:03:09 -0700820 elif '__new__' in base.__dict__:
Ethan Furmancd453852018-10-05 23:29:36 -0700821 if issubclass(base, Enum):
Ethan Furman5bdab642018-09-21 19:03:09 -0700822 continue
Miss Islington (bot)01286012021-06-10 15:01:03 -0700823 data_types.add(candidate or base)
Ethan Furmanbff01f32020-09-15 15:56:26 -0700824 break
Ethan Furmanc2667362020-12-07 00:17:31 -0800825 else:
Miss Islington (bot)0f993242021-06-15 14:07:37 -0700826 candidate = candidate or base
Ethan Furmanbff01f32020-09-15 15:56:26 -0700827 if len(data_types) > 1:
Ethan Furman3064dbf2020-09-16 07:11:57 -0700828 raise TypeError('%r: too many data types: %r' % (class_name, data_types))
Ethan Furmanbff01f32020-09-15 15:56:26 -0700829 elif data_types:
Miss Islington (bot)01286012021-06-10 15:01:03 -0700830 return data_types.pop()
Ethan Furmanbff01f32020-09-15 15:56:26 -0700831 else:
832 return None
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700833
Ethan Furman5bdab642018-09-21 19:03:09 -0700834 # ensure final parent class is an Enum derivative, find any concrete
835 # data type, and check that Enum has no members
836 first_enum = bases[-1]
837 if not issubclass(first_enum, Enum):
838 raise TypeError("new enumerations should be created as "
839 "`EnumName([mixin_type, ...] [data_type,] enum_type)`")
840 member_type = _find_data_type(bases) or object
841 if first_enum._member_names_:
842 raise TypeError("Cannot extend enumerations")
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700843 return member_type, first_enum
844
845 @staticmethod
846 def _find_new_(classdict, member_type, first_enum):
Ethan Furman6d3dfee2020-12-08 12:26:56 -0800847 """
848 Returns the __new__ to be used for creating the enum members.
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700849
850 classdict: the class dictionary given to __new__
851 member_type: the data type whose __new__ will be used by default
852 first_enum: enumeration to check for an overriding __new__
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700853 """
854 # now find the correct __new__, checking to see of one was defined
855 # by the user; also check earlier enum classes in case a __new__ was
856 # saved as __new_member__
857 __new__ = classdict.get('__new__', None)
858
859 # should __new__ be saved as __new_member__ later?
Ethan Furmana02cb472021-04-21 10:20:44 -0700860 save_new = first_enum is not None and __new__ is not None
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700861
862 if __new__ is None:
863 # check all possibles for __new_member__ before falling back to
864 # __new__
865 for method in ('__new_member__', '__new__'):
866 for possible in (member_type, first_enum):
867 target = getattr(possible, method, None)
868 if target not in {
869 None,
870 None.__new__,
871 object.__new__,
872 Enum.__new__,
873 }:
874 __new__ = target
875 break
876 if __new__ is not None:
877 break
878 else:
879 __new__ = object.__new__
880
881 # if a non-object.__new__ is used then whatever value/tuple was
882 # assigned to the enum member name will be passed to __new__ and to the
883 # new enum member's __init__
Ethan Furmana02cb472021-04-21 10:20:44 -0700884 if first_enum is None or __new__ in (Enum.__new__, object.__new__):
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700885 use_args = False
886 else:
887 use_args = True
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700888 return __new__, save_new, use_args
Ethan Furmanb7751062021-03-30 21:17:26 -0700889EnumMeta = EnumType
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700890
891
Ethan Furmanb7751062021-03-30 21:17:26 -0700892class Enum(metaclass=EnumType):
Ethan Furman6d3dfee2020-12-08 12:26:56 -0800893 """
894 Generic enumeration.
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700895
896 Derive from this class to define new enumerations.
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700897 """
Ethan Furman7aaeb2a2021-01-25 14:26:19 -0800898
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700899 def __new__(cls, value):
900 # all enum instances are actually created during class construction
901 # without calling this method; this method is called by the metaclass'
902 # __call__ (i.e. Color(3) ), and by pickle
903 if type(value) is cls:
Ethan Furman23bb6f42016-11-21 09:22:05 -0800904 # For lookups like Color(Color.RED)
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700905 return value
906 # by-value search for a matching enum member
907 # see if it's in the reverse mapping (for hashable values)
Ethan Furman2aa27322013-07-19 19:35:56 -0700908 try:
Andrew Svetlov34ae04f2018-12-26 20:45:33 +0200909 return cls._value2member_map_[value]
910 except KeyError:
911 # Not found, no need to do long O(n) search
912 pass
Ethan Furman2aa27322013-07-19 19:35:56 -0700913 except TypeError:
914 # not there, now do long search -- O(n) behavior
Ethan Furman520ad572013-07-19 19:47:21 -0700915 for member in cls._member_map_.values():
Ethan Furman0081f232014-09-16 17:31:23 -0700916 if member._value_ == value:
Ethan Furman2aa27322013-07-19 19:35:56 -0700917 return member
Ethan Furmanee47e5c2016-08-31 00:12:15 -0700918 # still not found -- try _missing_ hook
Ethan Furman019f0a02018-09-12 11:43:34 -0700919 try:
920 exc = None
921 result = cls._missing_(value)
922 except Exception as e:
923 exc = e
924 result = None
Ethan Furman8c14f5a2021-04-12 08:51:20 -0700925 try:
926 if isinstance(result, cls):
927 return result
928 elif (
929 Flag is not None and issubclass(cls, Flag)
930 and cls._boundary_ is EJECT and isinstance(result, int)
931 ):
932 return result
933 else:
934 ve_exc = ValueError("%r is not a valid %s" % (value, cls.__qualname__))
935 if result is None and exc is None:
936 raise ve_exc
937 elif exc is None:
938 exc = TypeError(
939 'error in %s._missing_: returned %r instead of None or a valid member'
940 % (cls.__name__, result)
941 )
942 if not isinstance(exc, ValueError):
943 exc.__context__ = ve_exc
944 raise exc
945 finally:
946 # ensure all variables that could hold an exception are destroyed
947 exc = None
948 ve_exc = None
Ethan Furmanee47e5c2016-08-31 00:12:15 -0700949
Ethan Furmanc16595e2016-09-10 23:36:59 -0700950 def _generate_next_value_(name, start, count, last_values):
Ethan Furman6d3dfee2020-12-08 12:26:56 -0800951 """
952 Generate the next value when not given.
953
954 name: the name of the member
955 start: the initial start value or None
956 count: the number of existing members
957 last_value: the last value assigned or None
958 """
Ethan Furmanc16595e2016-09-10 23:36:59 -0700959 for last_value in reversed(last_values):
960 try:
961 return last_value + 1
962 except TypeError:
963 pass
964 else:
Ethan Furmanee47e5c2016-08-31 00:12:15 -0700965 return start
Ethan Furmanc16595e2016-09-10 23:36:59 -0700966
Ethan Furmanee47e5c2016-08-31 00:12:15 -0700967 @classmethod
968 def _missing_(cls, value):
Ethan Furmanc95ad7a2020-09-16 10:26:50 -0700969 return None
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700970
971 def __repr__(self):
Ethan Furmanb7751062021-03-30 21:17:26 -0700972 return "%s.%s" % ( self.__class__.__name__, self._name_)
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700973
974 def __str__(self):
Ethan Furmanb7751062021-03-30 21:17:26 -0700975 return "%s" % (self._name_, )
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700976
Ethan Furman388a3922013-08-12 06:51:41 -0700977 def __dir__(self):
Ethan Furman6d3dfee2020-12-08 12:26:56 -0800978 """
979 Returns all members and all public methods
980 """
Ethan Furman0ae550b2014-10-14 08:58:32 -0700981 added_behavior = [
982 m
983 for cls in self.__class__.mro()
984 for m in cls.__dict__
Ethan Furman354ecf12015-03-11 08:43:12 -0700985 if m[0] != '_' and m not in self._member_map_
Angelin BOOZ68526fe2020-09-21 15:11:06 +0200986 ] + [m for m in self.__dict__ if m[0] != '_']
Ethan Furmanec5f8eb2014-10-21 13:40:35 -0700987 return (['__class__', '__doc__', '__module__'] + added_behavior)
Ethan Furman388a3922013-08-12 06:51:41 -0700988
Ethan Furmanec15a822013-08-31 19:17:41 -0700989 def __format__(self, format_spec):
Ethan Furman6d3dfee2020-12-08 12:26:56 -0800990 """
991 Returns format using actual value type unless __str__ has been overridden.
992 """
Ethan Furmanec15a822013-08-31 19:17:41 -0700993 # mixed-in Enums should use the mixed-in type's __format__, otherwise
994 # we can get strange results with the Enum name showing up instead of
995 # the value
996
thatneat2f19e822019-07-04 11:28:37 -0700997 # pure Enum branch, or branch with __str__ explicitly overridden
Ethan Furman37440ee2020-12-08 11:14:10 -0800998 str_overridden = type(self).__str__ not in (Enum.__str__, Flag.__str__)
thatneat2f19e822019-07-04 11:28:37 -0700999 if self._member_type_ is object or str_overridden:
Ethan Furmanec15a822013-08-31 19:17:41 -07001000 cls = str
1001 val = str(self)
1002 # mix-in branch
1003 else:
Ethan Furman6bd92882021-04-27 13:05:08 -07001004 if not format_spec or format_spec in ('{}','{:}'):
1005 import warnings
1006 warnings.warn(
1007 "in 3.12 format() will use the enum member, not the enum member's value;\n"
Pablo Galindo9a42d502021-05-01 20:26:09 +01001008 "use a format specifier, such as :d for an IntEnum member, to maintain "
Ethan Furman6bd92882021-04-27 13:05:08 -07001009 "the current display",
1010 DeprecationWarning,
1011 stacklevel=2,
1012 )
Ethan Furmanec15a822013-08-31 19:17:41 -07001013 cls = self._member_type_
Ethan Furman0081f232014-09-16 17:31:23 -07001014 val = self._value_
Ethan Furmanec15a822013-08-31 19:17:41 -07001015 return cls.__format__(val, format_spec)
1016
Ethan Furman6b3d64a2013-06-14 16:55:46 -07001017 def __hash__(self):
Ethan Furman520ad572013-07-19 19:47:21 -07001018 return hash(self._name_)
Ethan Furman6b3d64a2013-06-14 16:55:46 -07001019
Ethan Furmanca1b7942014-02-08 11:36:27 -08001020 def __reduce_ex__(self, proto):
Miss Islington (bot)b6131322021-06-10 16:37:27 -07001021 return getattr, (self.__class__, self._name_)
Ethan Furmanca1b7942014-02-08 11:36:27 -08001022
Ethan Furmanc314e602021-01-12 23:47:57 -08001023 # enum.property is used to provide access to the `name` and
1024 # `value` attributes of enum members while keeping some measure of
Ethan Furman6b3d64a2013-06-14 16:55:46 -07001025 # protection from modification, while still allowing for an enumeration
1026 # to have members named `name` and `value`. This works because enumeration
Ethan Furmanc314e602021-01-12 23:47:57 -08001027 # members are not set directly on the enum class; they are kept in a
1028 # separate structure, _member_map_, which is where enum.property looks for
1029 # them
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 name(self):
Ethan Furmanc850f342013-09-15 16:59:35 -07001033 """The name of the Enum member."""
Ethan Furman520ad572013-07-19 19:47:21 -07001034 return self._name_
Ethan Furman6b3d64a2013-06-14 16:55:46 -07001035
Ethan Furmanc314e602021-01-12 23:47:57 -08001036 @property
Ethan Furman6b3d64a2013-06-14 16:55:46 -07001037 def value(self):
Ethan Furmanc850f342013-09-15 16:59:35 -07001038 """The value of the Enum member."""
Ethan Furman520ad572013-07-19 19:47:21 -07001039 return self._value_
Ethan Furman6b3d64a2013-06-14 16:55:46 -07001040
1041
1042class IntEnum(int, Enum):
Ethan Furman0063ff42020-09-21 17:23:13 -07001043 """
1044 Enum where members are also (and must be) ints
1045 """
1046
1047
1048class StrEnum(str, Enum):
1049 """
1050 Enum where members are also (and must be) strings
1051 """
1052
1053 def __new__(cls, *values):
1054 if len(values) > 3:
1055 raise TypeError('too many arguments for str(): %r' % (values, ))
1056 if len(values) == 1:
1057 # it must be a string
1058 if not isinstance(values[0], str):
1059 raise TypeError('%r is not a string' % (values[0], ))
Ethan Furman7aaeb2a2021-01-25 14:26:19 -08001060 if len(values) >= 2:
Ethan Furman0063ff42020-09-21 17:23:13 -07001061 # check that encoding argument is a string
1062 if not isinstance(values[1], str):
1063 raise TypeError('encoding must be a string, not %r' % (values[1], ))
Ethan Furman7aaeb2a2021-01-25 14:26:19 -08001064 if len(values) == 3:
1065 # check that errors argument is a string
1066 if not isinstance(values[2], str):
1067 raise TypeError('errors must be a string, not %r' % (values[2]))
Ethan Furman0063ff42020-09-21 17:23:13 -07001068 value = str(*values)
1069 member = str.__new__(cls, value)
1070 member._value_ = value
1071 return member
Ethan Furmanf24bb352013-07-18 17:05:39 -07001072
Ethan Furmand986d162020-09-22 13:00:07 -07001073 __str__ = str.__str__
1074
Ethan Furmanefb13be2020-12-10 12:20:06 -08001075 def _generate_next_value_(name, start, count, last_values):
1076 """
1077 Return the lower-cased version of the member name.
1078 """
1079 return name.lower()
1080
Ethan Furmanf24bb352013-07-18 17:05:39 -07001081
Miss Islington (bot)b6131322021-06-10 16:37:27 -07001082def _reduce_ex_by_global_name(self, proto):
Ethan Furman24e837f2015-03-18 17:27:57 -07001083 return self.name
1084
Ethan Furman7aaeb2a2021-01-25 14:26:19 -08001085class FlagBoundary(StrEnum):
1086 """
1087 control how out of range values are handled
1088 "strict" -> error is raised [default for Flag]
1089 "conform" -> extra bits are discarded
1090 "eject" -> lose flag status [default for IntFlag]
1091 "keep" -> keep flag status and all bits
1092 """
1093 STRICT = auto()
1094 CONFORM = auto()
1095 EJECT = auto()
1096 KEEP = auto()
1097STRICT, CONFORM, EJECT, KEEP = FlagBoundary
1098
1099
1100class Flag(Enum, boundary=STRICT):
Ethan Furman6d3dfee2020-12-08 12:26:56 -08001101 """
1102 Support for flags
1103 """
Ethan Furmanc16595e2016-09-10 23:36:59 -07001104
1105 def _generate_next_value_(name, start, count, last_values):
Ethan Furmanee47e5c2016-08-31 00:12:15 -07001106 """
1107 Generate the next value when not given.
1108
1109 name: the name of the member
HongWeipengbb16fb22019-09-21 13:22:54 +08001110 start: the initial start value or None
Ethan Furmanee47e5c2016-08-31 00:12:15 -07001111 count: the number of existing members
1112 last_value: the last value assigned or None
1113 """
1114 if not count:
1115 return start if start is not None else 1
Ethan Furman7aaeb2a2021-01-25 14:26:19 -08001116 last_value = max(last_values)
1117 try:
1118 high_bit = _high_bit(last_value)
1119 except Exception:
1120 raise TypeError('Invalid Flag value: %r' % last_value) from None
Ethan Furmanee47e5c2016-08-31 00:12:15 -07001121 return 2 ** (high_bit+1)
1122
1123 @classmethod
Ethan Furman7aaeb2a2021-01-25 14:26:19 -08001124 def _iter_member_by_value_(cls, value):
1125 """
1126 Extract all members from the value in definition (i.e. increasing value) order.
1127 """
1128 for val in _iter_bits_lsb(value & cls._flag_mask_):
1129 yield cls._value2member_map_.get(val)
1130
1131 _iter_member_ = _iter_member_by_value_
1132
1133 @classmethod
1134 def _iter_member_by_def_(cls, value):
1135 """
1136 Extract all members from the value in definition order.
1137 """
1138 yield from sorted(
1139 cls._iter_member_by_value_(value),
1140 key=lambda m: m._sort_order_,
1141 )
1142
1143 @classmethod
Ethan Furmanee47e5c2016-08-31 00:12:15 -07001144 def _missing_(cls, value):
Ethan Furman6d3dfee2020-12-08 12:26:56 -08001145 """
Ethan Furman0dca5eb2021-04-15 06:49:54 -07001146 Create a composite member containing all canonical members present in `value`.
1147
1148 If non-member values are present, result depends on `_boundary_` setting.
Ethan Furman3515dcc2016-09-18 13:15:41 -07001149 """
Ethan Furman7aaeb2a2021-01-25 14:26:19 -08001150 if not isinstance(value, int):
1151 raise ValueError(
1152 "%r is not a valid %s" % (value, cls.__qualname__)
1153 )
1154 # check boundaries
1155 # - value must be in range (e.g. -16 <-> +15, i.e. ~15 <-> 15)
1156 # - value must not include any skipped flags (e.g. if bit 2 is not
1157 # defined, then 0d10 is invalid)
1158 flag_mask = cls._flag_mask_
1159 all_bits = cls._all_bits_
1160 neg_value = None
1161 if (
1162 not ~all_bits <= value <= all_bits
1163 or value & (all_bits ^ flag_mask)
1164 ):
1165 if cls._boundary_ is STRICT:
1166 max_bits = max(value.bit_length(), flag_mask.bit_length())
1167 raise ValueError(
1168 "%s: invalid value: %r\n given %s\n allowed %s" % (
1169 cls.__name__, value, bin(value, max_bits), bin(flag_mask, max_bits),
1170 ))
1171 elif cls._boundary_ is CONFORM:
1172 value = value & flag_mask
1173 elif cls._boundary_ is EJECT:
1174 return value
1175 elif cls._boundary_ is KEEP:
1176 if value < 0:
1177 value = (
1178 max(all_bits+1, 2**(value.bit_length()))
1179 + value
1180 )
1181 else:
1182 raise ValueError(
1183 'unknown flag boundary: %r' % (cls._boundary_, )
1184 )
1185 if value < 0:
1186 neg_value = value
1187 value = all_bits + 1 + value
1188 # get members and unknown
1189 unknown = value & ~flag_mask
1190 member_value = value & flag_mask
1191 if unknown and cls._boundary_ is not KEEP:
1192 raise ValueError(
1193 '%s(%r) --> unknown values %r [%s]'
1194 % (cls.__name__, value, unknown, bin(unknown))
1195 )
1196 # normal Flag?
1197 __new__ = getattr(cls, '__new_member__', None)
1198 if cls._member_type_ is object and not __new__:
Ethan Furman3515dcc2016-09-18 13:15:41 -07001199 # construct a singleton enum pseudo-member
Ethan Furmanee47e5c2016-08-31 00:12:15 -07001200 pseudo_member = object.__new__(cls)
Ethan Furman7aaeb2a2021-01-25 14:26:19 -08001201 else:
1202 pseudo_member = (__new__ or cls._member_type_.__new__)(cls, value)
Ethan Furmana02cb472021-04-21 10:20:44 -07001203 if not hasattr(pseudo_member, '_value_'):
Ethan Furmanee47e5c2016-08-31 00:12:15 -07001204 pseudo_member._value_ = value
Ethan Furman7aaeb2a2021-01-25 14:26:19 -08001205 if member_value:
1206 pseudo_member._name_ = '|'.join([
1207 m._name_ for m in cls._iter_member_(member_value)
1208 ])
1209 if unknown:
1210 pseudo_member._name_ += '|0x%x' % unknown
1211 else:
1212 pseudo_member._name_ = None
1213 # use setdefault in case another thread already created a composite
1214 # with this value, but only if all members are known
1215 # note: zero is a special case -- add it
1216 if not unknown:
Ethan Furman28cf6632017-01-24 12:12:06 -08001217 pseudo_member = cls._value2member_map_.setdefault(value, pseudo_member)
Ethan Furman7aaeb2a2021-01-25 14:26:19 -08001218 if neg_value is not None:
1219 cls._value2member_map_[neg_value] = pseudo_member
Ethan Furmanee47e5c2016-08-31 00:12:15 -07001220 return pseudo_member
1221
Ethan Furmanee47e5c2016-08-31 00:12:15 -07001222 def __contains__(self, other):
Ethan Furman6d3dfee2020-12-08 12:26:56 -08001223 """
1224 Returns True if self has at least the same flags set as other.
1225 """
Ethan Furmanee47e5c2016-08-31 00:12:15 -07001226 if not isinstance(other, self.__class__):
Rahul Jha94306522018-09-10 23:51:04 +05301227 raise TypeError(
1228 "unsupported operand type(s) for 'in': '%s' and '%s'" % (
1229 type(other).__qualname__, self.__class__.__qualname__))
Ethan Furman7aaeb2a2021-01-25 14:26:19 -08001230 if other._value_ == 0 or self._value_ == 0:
1231 return False
Ethan Furmanee47e5c2016-08-31 00:12:15 -07001232 return other._value_ & self._value_ == other._value_
1233
Ethan Furman7219e272020-09-16 13:01:00 -07001234 def __iter__(self):
Ethan Furman6d3dfee2020-12-08 12:26:56 -08001235 """
Ethan Furman7aaeb2a2021-01-25 14:26:19 -08001236 Returns flags in definition order.
Ethan Furman6d3dfee2020-12-08 12:26:56 -08001237 """
Ethan Furman7aaeb2a2021-01-25 14:26:19 -08001238 yield from self._iter_member_(self._value_)
1239
1240 def __len__(self):
1241 return self._value_.bit_count()
Ethan Furman7219e272020-09-16 13:01:00 -07001242
Ethan Furmanee47e5c2016-08-31 00:12:15 -07001243 def __repr__(self):
Ethan Furmanb7751062021-03-30 21:17:26 -07001244 cls_name = self.__class__.__name__
1245 if self._name_ is None:
1246 return "0x%x" % (self._value_, )
1247 if _is_single_bit(self._value_):
1248 return '%s.%s' % (cls_name, self._name_)
1249 if self._boundary_ is not FlagBoundary.KEEP:
1250 return '%s.' % cls_name + ('|%s.' % cls_name).join(self.name.split('|'))
Ethan Furman7aaeb2a2021-01-25 14:26:19 -08001251 else:
Ethan Furmanb7751062021-03-30 21:17:26 -07001252 name = []
1253 for n in self._name_.split('|'):
1254 if n.startswith('0'):
1255 name.append(n)
1256 else:
1257 name.append('%s.%s' % (cls_name, n))
1258 return '|'.join(name)
Ethan Furmanee47e5c2016-08-31 00:12:15 -07001259
1260 def __str__(self):
1261 cls = self.__class__
Ethan Furmanb7751062021-03-30 21:17:26 -07001262 if self._name_ is None:
1263 return '%s(%x)' % (cls.__name__, self._value_)
Ethan Furmanee47e5c2016-08-31 00:12:15 -07001264 else:
Ethan Furmanb7751062021-03-30 21:17:26 -07001265 return self._name_
Ethan Furmanee47e5c2016-08-31 00:12:15 -07001266
Ethan Furman25d94bb2016-09-02 16:32:32 -07001267 def __bool__(self):
1268 return bool(self._value_)
1269
Ethan Furmanee47e5c2016-08-31 00:12:15 -07001270 def __or__(self, other):
1271 if not isinstance(other, self.__class__):
1272 return NotImplemented
1273 return self.__class__(self._value_ | other._value_)
1274
1275 def __and__(self, other):
1276 if not isinstance(other, self.__class__):
1277 return NotImplemented
1278 return self.__class__(self._value_ & other._value_)
1279
1280 def __xor__(self, other):
1281 if not isinstance(other, self.__class__):
1282 return NotImplemented
1283 return self.__class__(self._value_ ^ other._value_)
1284
1285 def __invert__(self):
Ethan Furman7aaeb2a2021-01-25 14:26:19 -08001286 if self._inverted_ is None:
1287 if self._boundary_ is KEEP:
1288 # use all bits
1289 self._inverted_ = self.__class__(~self._value_)
1290 else:
1291 # calculate flags not in this member
1292 self._inverted_ = self.__class__(self._flag_mask_ ^ self._value_)
Ethan Furman74964862021-06-10 07:24:20 -07001293 if isinstance(self._inverted_, self.__class__):
1294 self._inverted_._inverted_ = self
Ethan Furman7aaeb2a2021-01-25 14:26:19 -08001295 return self._inverted_
Ethan Furmanee47e5c2016-08-31 00:12:15 -07001296
1297
Ethan Furman7aaeb2a2021-01-25 14:26:19 -08001298class IntFlag(int, Flag, boundary=EJECT):
Ethan Furman6d3dfee2020-12-08 12:26:56 -08001299 """
1300 Support for integer-based Flags
1301 """
Ethan Furmanee47e5c2016-08-31 00:12:15 -07001302
Ethan Furmanee47e5c2016-08-31 00:12:15 -07001303 def __or__(self, other):
Ethan Furman7aaeb2a2021-01-25 14:26:19 -08001304 if isinstance(other, self.__class__):
1305 other = other._value_
1306 elif isinstance(other, int):
1307 other = other
1308 else:
Ethan Furmanee47e5c2016-08-31 00:12:15 -07001309 return NotImplemented
Ethan Furman7aaeb2a2021-01-25 14:26:19 -08001310 value = self._value_
1311 return self.__class__(value | other)
Ethan Furmanee47e5c2016-08-31 00:12:15 -07001312
1313 def __and__(self, other):
Ethan Furman7aaeb2a2021-01-25 14:26:19 -08001314 if isinstance(other, self.__class__):
1315 other = other._value_
1316 elif isinstance(other, int):
1317 other = other
1318 else:
Ethan Furmanee47e5c2016-08-31 00:12:15 -07001319 return NotImplemented
Ethan Furman7aaeb2a2021-01-25 14:26:19 -08001320 value = self._value_
1321 return self.__class__(value & other)
Ethan Furmanee47e5c2016-08-31 00:12:15 -07001322
1323 def __xor__(self, other):
Ethan Furman7aaeb2a2021-01-25 14:26:19 -08001324 if isinstance(other, self.__class__):
1325 other = other._value_
1326 elif isinstance(other, int):
1327 other = other
1328 else:
Ethan Furmanee47e5c2016-08-31 00:12:15 -07001329 return NotImplemented
Ethan Furman7aaeb2a2021-01-25 14:26:19 -08001330 value = self._value_
1331 return self.__class__(value ^ other)
Ethan Furmanee47e5c2016-08-31 00:12:15 -07001332
1333 __ror__ = __or__
1334 __rand__ = __and__
1335 __rxor__ = __xor__
Ethan Furman7aaeb2a2021-01-25 14:26:19 -08001336 __invert__ = Flag.__invert__
Ethan Furmanee47e5c2016-08-31 00:12:15 -07001337
1338def _high_bit(value):
Ethan Furman6d3dfee2020-12-08 12:26:56 -08001339 """
1340 returns index of highest bit, or -1 if value is zero or negative
1341 """
Ethan Furman3515dcc2016-09-18 13:15:41 -07001342 return value.bit_length() - 1
Ethan Furmanee47e5c2016-08-31 00:12:15 -07001343
Ethan Furmanf24bb352013-07-18 17:05:39 -07001344def unique(enumeration):
Ethan Furman6d3dfee2020-12-08 12:26:56 -08001345 """
1346 Class decorator for enumerations ensuring unique member values.
1347 """
Ethan Furmanf24bb352013-07-18 17:05:39 -07001348 duplicates = []
1349 for name, member in enumeration.__members__.items():
1350 if name != member.name:
1351 duplicates.append((name, member.name))
1352 if duplicates:
1353 alias_details = ', '.join(
1354 ["%s -> %s" % (alias, name) for (alias, name) in duplicates])
1355 raise ValueError('duplicate values found in %r: %s' %
1356 (enumeration, alias_details))
1357 return enumeration
Ethan Furman3515dcc2016-09-18 13:15:41 -07001358
Ethan Furman7aaeb2a2021-01-25 14:26:19 -08001359def _power_of_two(value):
1360 if value < 1:
1361 return False
1362 return value == 2 ** _high_bit(value)
Ethan Furmanb7751062021-03-30 21:17:26 -07001363
1364def global_enum_repr(self):
1365 return '%s.%s' % (self.__class__.__module__, self._name_)
1366
1367def global_flag_repr(self):
1368 module = self.__class__.__module__
1369 cls_name = self.__class__.__name__
1370 if self._name_ is None:
1371 return "%x" % (module, cls_name, self._value_)
1372 if _is_single_bit(self):
1373 return '%s.%s' % (module, self._name_)
1374 if self._boundary_ is not FlagBoundary.KEEP:
1375 return module + module.join(self.name.split('|'))
1376 else:
1377 name = []
1378 for n in self._name_.split('|'):
1379 if n.startswith('0'):
1380 name.append(n)
1381 else:
1382 name.append('%s.%s' % (module, n))
1383 return '|'.join(name)
1384
1385
1386def global_enum(cls):
1387 """
1388 decorator that makes the repr() of an enum member reference its module
1389 instead of its class; also exports all members to the enum's module's
1390 global namespace
1391 """
1392 if issubclass(cls, Flag):
1393 cls.__repr__ = global_flag_repr
1394 else:
1395 cls.__repr__ = global_enum_repr
1396 sys.modules[cls.__module__].__dict__.update(cls.__members__)
1397 return cls
Ethan Furmana02cb472021-04-21 10:20:44 -07001398
1399def _simple_enum(etype=Enum, *, boundary=None, use_args=None):
1400 """
1401 Class decorator that converts a normal class into an :class:`Enum`. No
1402 safety checks are done, and some advanced behavior (such as
1403 :func:`__init_subclass__`) is not available. Enum creation can be faster
1404 using :func:`simple_enum`.
1405
1406 >>> from enum import Enum, _simple_enum
1407 >>> @_simple_enum(Enum)
1408 ... class Color:
1409 ... RED = auto()
1410 ... GREEN = auto()
1411 ... BLUE = auto()
1412 >>> Color
1413 <enum 'Color'>
1414 """
1415 def convert_class(cls):
1416 nonlocal use_args
1417 cls_name = cls.__name__
1418 if use_args is None:
1419 use_args = etype._use_args_
1420 __new__ = cls.__dict__.get('__new__')
1421 if __new__ is not None:
1422 new_member = __new__.__func__
1423 else:
1424 new_member = etype._member_type_.__new__
1425 attrs = {}
1426 body = {}
1427 if __new__ is not None:
1428 body['__new_member__'] = new_member
1429 body['_new_member_'] = new_member
1430 body['_use_args_'] = use_args
1431 body['_generate_next_value_'] = gnv = etype._generate_next_value_
1432 body['_member_names_'] = member_names = []
1433 body['_member_map_'] = member_map = {}
1434 body['_value2member_map_'] = value2member_map = {}
Ethan Furman6bd92882021-04-27 13:05:08 -07001435 body['_unhashable_values_'] = []
Ethan Furmana02cb472021-04-21 10:20:44 -07001436 body['_member_type_'] = member_type = etype._member_type_
1437 if issubclass(etype, Flag):
1438 body['_boundary_'] = boundary or etype._boundary_
1439 body['_flag_mask_'] = None
1440 body['_all_bits_'] = None
1441 body['_inverted_'] = None
1442 for name, obj in cls.__dict__.items():
1443 if name in ('__dict__', '__weakref__'):
1444 continue
1445 if _is_dunder(name) or _is_private(cls_name, name) or _is_sunder(name) or _is_descriptor(obj):
1446 body[name] = obj
1447 else:
1448 attrs[name] = obj
1449 if cls.__dict__.get('__doc__') is None:
1450 body['__doc__'] = 'An enumeration.'
1451 #
1452 # double check that repr and friends are not the mixin's or various
1453 # things break (such as pickle)
1454 # however, if the method is defined in the Enum itself, don't replace
1455 # it
1456 enum_class = type(cls_name, (etype, ), body, boundary=boundary, _simple=True)
1457 for name in ('__repr__', '__str__', '__format__', '__reduce_ex__'):
1458 if name in body:
1459 continue
1460 class_method = getattr(enum_class, name)
1461 obj_method = getattr(member_type, name, None)
1462 enum_method = getattr(etype, name, None)
1463 if obj_method is not None and obj_method is class_method:
1464 setattr(enum_class, name, enum_method)
1465 gnv_last_values = []
1466 if issubclass(enum_class, Flag):
1467 # Flag / IntFlag
1468 single_bits = multi_bits = 0
1469 for name, value in attrs.items():
1470 if isinstance(value, auto) and auto.value is _auto_null:
1471 value = gnv(name, 1, len(member_names), gnv_last_values)
1472 if value in value2member_map:
1473 # an alias to an existing member
1474 redirect = property()
1475 redirect.__set_name__(enum_class, name)
1476 setattr(enum_class, name, redirect)
1477 member_map[name] = value2member_map[value]
1478 else:
1479 # create the member
1480 if use_args:
1481 if not isinstance(value, tuple):
1482 value = (value, )
1483 member = new_member(enum_class, *value)
1484 value = value[0]
1485 else:
1486 member = new_member(enum_class)
1487 if __new__ is None:
1488 member._value_ = value
1489 member._name_ = name
1490 member.__objclass__ = enum_class
1491 member.__init__(value)
1492 redirect = property()
1493 redirect.__set_name__(enum_class, name)
1494 setattr(enum_class, name, redirect)
1495 member_map[name] = member
1496 member._sort_order_ = len(member_names)
1497 value2member_map[value] = member
1498 if _is_single_bit(value):
1499 # not a multi-bit alias, record in _member_names_ and _flag_mask_
1500 member_names.append(name)
1501 single_bits |= value
1502 else:
1503 multi_bits |= value
1504 gnv_last_values.append(value)
1505 enum_class._flag_mask_ = single_bits
1506 enum_class._all_bits_ = 2 ** ((single_bits|multi_bits).bit_length()) - 1
1507 # set correct __iter__
1508 member_list = [m._value_ for m in enum_class]
1509 if member_list != sorted(member_list):
1510 enum_class._iter_member_ = enum_class._iter_member_by_def_
1511 else:
1512 # Enum / IntEnum / StrEnum
1513 for name, value in attrs.items():
1514 if isinstance(value, auto):
1515 if value.value is _auto_null:
1516 value.value = gnv(name, 1, len(member_names), gnv_last_values)
1517 value = value.value
1518 if value in value2member_map:
1519 # an alias to an existing member
1520 redirect = property()
1521 redirect.__set_name__(enum_class, name)
1522 setattr(enum_class, name, redirect)
1523 member_map[name] = value2member_map[value]
1524 else:
1525 # create the member
1526 if use_args:
1527 if not isinstance(value, tuple):
1528 value = (value, )
1529 member = new_member(enum_class, *value)
1530 value = value[0]
1531 else:
1532 member = new_member(enum_class)
1533 if __new__ is None:
1534 member._value_ = value
1535 member._name_ = name
1536 member.__objclass__ = enum_class
1537 member.__init__(value)
1538 member._sort_order_ = len(member_names)
1539 redirect = property()
1540 redirect.__set_name__(enum_class, name)
1541 setattr(enum_class, name, redirect)
1542 member_map[name] = member
1543 value2member_map[value] = member
1544 member_names.append(name)
1545 gnv_last_values.append(value)
1546 if '__new__' in body:
1547 enum_class.__new_member__ = enum_class.__new__
1548 enum_class.__new__ = Enum.__new__
1549 return enum_class
1550 return convert_class
1551
Ethan Furman74964862021-06-10 07:24:20 -07001552@_simple_enum(StrEnum)
1553class EnumCheck:
1554 """
1555 various conditions to check an enumeration for
1556 """
1557 CONTINUOUS = "no skipped integer values"
1558 NAMED_FLAGS = "multi-flag aliases may not contain unnamed flags"
1559 UNIQUE = "one name per value"
1560CONTINUOUS, NAMED_FLAGS, UNIQUE = EnumCheck
1561
1562
1563class verify:
1564 """
1565 Check an enumeration for various constraints. (see EnumCheck)
1566 """
1567 def __init__(self, *checks):
1568 self.checks = checks
1569 def __call__(self, enumeration):
1570 checks = self.checks
1571 cls_name = enumeration.__name__
1572 if Flag is not None and issubclass(enumeration, Flag):
1573 enum_type = 'flag'
1574 elif issubclass(enumeration, Enum):
1575 enum_type = 'enum'
1576 else:
1577 raise TypeError("the 'verify' decorator only works with Enum and Flag")
1578 for check in checks:
1579 if check is UNIQUE:
1580 # check for duplicate names
1581 duplicates = []
1582 for name, member in enumeration.__members__.items():
1583 if name != member.name:
1584 duplicates.append((name, member.name))
1585 if duplicates:
1586 alias_details = ', '.join(
1587 ["%s -> %s" % (alias, name) for (alias, name) in duplicates])
1588 raise ValueError('aliases found in %r: %s' %
1589 (enumeration, alias_details))
1590 elif check is CONTINUOUS:
1591 values = set(e.value for e in enumeration)
1592 if len(values) < 2:
1593 continue
1594 low, high = min(values), max(values)
1595 missing = []
1596 if enum_type == 'flag':
1597 # check for powers of two
1598 for i in range(_high_bit(low)+1, _high_bit(high)):
1599 if 2**i not in values:
1600 missing.append(2**i)
1601 elif enum_type == 'enum':
1602 # check for powers of one
1603 for i in range(low+1, high):
1604 if i not in values:
1605 missing.append(i)
1606 else:
1607 raise Exception('verify: unknown type %r' % enum_type)
1608 if missing:
Miss Islington (bot)0a186b12021-06-11 02:58:57 -07001609 raise ValueError(('invalid %s %r: missing values %s' % (
Ethan Furman74964862021-06-10 07:24:20 -07001610 enum_type, cls_name, ', '.join((str(m) for m in missing)))
Miss Islington (bot)0a186b12021-06-11 02:58:57 -07001611 )[:256])
1612 # limit max length to protect against DOS attacks
Ethan Furman74964862021-06-10 07:24:20 -07001613 elif check is NAMED_FLAGS:
1614 # examine each alias and check for unnamed flags
1615 member_names = enumeration._member_names_
1616 member_values = [m.value for m in enumeration]
Miss Islington (bot)0a186b12021-06-11 02:58:57 -07001617 missing_names = []
1618 missing_value = 0
Ethan Furman74964862021-06-10 07:24:20 -07001619 for name, alias in enumeration._member_map_.items():
1620 if name in member_names:
1621 # not an alias
1622 continue
1623 values = list(_iter_bits_lsb(alias.value))
1624 missed = [v for v in values if v not in member_values]
1625 if missed:
Miss Islington (bot)0a186b12021-06-11 02:58:57 -07001626 missing_names.append(name)
1627 missing_value |= reduce(_or_, missed)
1628 if missing_names:
1629 if len(missing_names) == 1:
1630 alias = 'alias %s is missing' % missing_names[0]
1631 else:
1632 alias = 'aliases %s and %s are missing' % (
1633 ', '.join(missing_names[:-1]), missing_names[-1]
1634 )
1635 if _is_single_bit(missing_value):
1636 value = 'value 0x%x' % missing_value
1637 else:
1638 value = 'combined values of 0x%x' % missing_value
Ethan Furman74964862021-06-10 07:24:20 -07001639 raise ValueError(
Miss Islington (bot)0a186b12021-06-11 02:58:57 -07001640 'invalid Flag %r: %s %s [use `enum.show_flag_values(value)` for details]'
1641 % (cls_name, alias, value)
Ethan Furman74964862021-06-10 07:24:20 -07001642 )
1643 return enumeration
1644
Ethan Furmana02cb472021-04-21 10:20:44 -07001645def _test_simple_enum(checked_enum, simple_enum):
1646 """
1647 A function that can be used to test an enum created with :func:`_simple_enum`
1648 against the version created by subclassing :class:`Enum`::
1649
1650 >>> from enum import Enum, _simple_enum, _test_simple_enum
1651 >>> @_simple_enum(Enum)
1652 ... class Color:
1653 ... RED = auto()
1654 ... GREEN = auto()
1655 ... BLUE = auto()
1656 >>> class CheckedColor(Enum):
1657 ... RED = auto()
1658 ... GREEN = auto()
1659 ... BLUE = auto()
1660 >>> _test_simple_enum(CheckedColor, Color)
1661
1662 If differences are found, a :exc:`TypeError` is raised.
1663 """
1664 failed = []
1665 if checked_enum.__dict__ != simple_enum.__dict__:
1666 checked_dict = checked_enum.__dict__
1667 checked_keys = list(checked_dict.keys())
1668 simple_dict = simple_enum.__dict__
1669 simple_keys = list(simple_dict.keys())
1670 member_names = set(
1671 list(checked_enum._member_map_.keys())
1672 + list(simple_enum._member_map_.keys())
1673 )
1674 for key in set(checked_keys + simple_keys):
1675 if key in ('__module__', '_member_map_', '_value2member_map_'):
1676 # keys known to be different
1677 continue
1678 elif key in member_names:
1679 # members are checked below
1680 continue
1681 elif key not in simple_keys:
1682 failed.append("missing key: %r" % (key, ))
1683 elif key not in checked_keys:
1684 failed.append("extra key: %r" % (key, ))
1685 else:
1686 checked_value = checked_dict[key]
1687 simple_value = simple_dict[key]
1688 if callable(checked_value):
1689 continue
1690 if key == '__doc__':
1691 # remove all spaces/tabs
1692 compressed_checked_value = checked_value.replace(' ','').replace('\t','')
1693 compressed_simple_value = simple_value.replace(' ','').replace('\t','')
1694 if compressed_checked_value != compressed_simple_value:
1695 failed.append("%r:\n %s\n %s" % (
1696 key,
1697 "checked -> %r" % (checked_value, ),
1698 "simple -> %r" % (simple_value, ),
1699 ))
1700 elif checked_value != simple_value:
1701 failed.append("%r:\n %s\n %s" % (
1702 key,
1703 "checked -> %r" % (checked_value, ),
1704 "simple -> %r" % (simple_value, ),
1705 ))
1706 failed.sort()
1707 for name in member_names:
1708 failed_member = []
1709 if name not in simple_keys:
1710 failed.append('missing member from simple enum: %r' % name)
1711 elif name not in checked_keys:
1712 failed.append('extra member in simple enum: %r' % name)
1713 else:
1714 checked_member_dict = checked_enum[name].__dict__
1715 checked_member_keys = list(checked_member_dict.keys())
1716 simple_member_dict = simple_enum[name].__dict__
1717 simple_member_keys = list(simple_member_dict.keys())
1718 for key in set(checked_member_keys + simple_member_keys):
Ethan Furman6c681e12021-04-23 19:08:22 -07001719 if key in ('__module__', '__objclass__', '_inverted_'):
1720 # keys known to be different or absent
Ethan Furmana02cb472021-04-21 10:20:44 -07001721 continue
1722 elif key not in simple_member_keys:
1723 failed_member.append("missing key %r not in the simple enum member %r" % (key, name))
1724 elif key not in checked_member_keys:
1725 failed_member.append("extra key %r in simple enum member %r" % (key, name))
1726 else:
1727 checked_value = checked_member_dict[key]
1728 simple_value = simple_member_dict[key]
1729 if checked_value != simple_value:
1730 failed_member.append("%r:\n %s\n %s" % (
1731 key,
1732 "checked member -> %r" % (checked_value, ),
1733 "simple member -> %r" % (simple_value, ),
1734 ))
1735 if failed_member:
1736 failed.append('%r member mismatch:\n %s' % (
1737 name, '\n '.join(failed_member),
1738 ))
1739 for method in (
1740 '__str__', '__repr__', '__reduce_ex__', '__format__',
1741 '__getnewargs_ex__', '__getnewargs__', '__reduce_ex__', '__reduce__'
1742 ):
1743 if method in simple_keys and method in checked_keys:
1744 # cannot compare functions, and it exists in both, so we're good
1745 continue
1746 elif method not in simple_keys and method not in checked_keys:
1747 # method is inherited -- check it out
1748 checked_method = getattr(checked_enum, method, None)
1749 simple_method = getattr(simple_enum, method, None)
1750 if hasattr(checked_method, '__func__'):
1751 checked_method = checked_method.__func__
1752 simple_method = simple_method.__func__
1753 if checked_method != simple_method:
1754 failed.append("%r: %-30s %s" % (
1755 method,
1756 "checked -> %r" % (checked_method, ),
1757 "simple -> %r" % (simple_method, ),
1758 ))
1759 else:
1760 # if the method existed in only one of the enums, it will have been caught
1761 # in the first checks above
1762 pass
1763 if failed:
1764 raise TypeError('enum mismatch:\n %s' % '\n '.join(failed))
1765
1766def _old_convert_(etype, name, module, filter, source=None, *, boundary=None):
1767 """
1768 Create a new Enum subclass that replaces a collection of global constants
1769 """
1770 # convert all constants from source (or module) that pass filter() to
1771 # a new Enum called name, and export the enum and its members back to
1772 # module;
1773 # also, replace the __reduce_ex__ method so unpickling works in
1774 # previous Python versions
1775 module_globals = sys.modules[module].__dict__
1776 if source:
1777 source = source.__dict__
1778 else:
1779 source = module_globals
1780 # _value2member_map_ is populated in the same order every time
1781 # for a consistent reverse mapping of number to name when there
1782 # are multiple names for the same number.
1783 members = [
1784 (name, value)
1785 for name, value in source.items()
1786 if filter(name)]
1787 try:
1788 # sort by value
1789 members.sort(key=lambda t: (t[1], t[0]))
1790 except TypeError:
1791 # unless some values aren't comparable, in which case sort by name
1792 members.sort(key=lambda t: t[0])
1793 cls = etype(name, members, module=module, boundary=boundary or KEEP)
Miss Islington (bot)b6131322021-06-10 16:37:27 -07001794 cls.__reduce_ex__ = _reduce_ex_by_global_name
Ethan Furmana02cb472021-04-21 10:20:44 -07001795 cls.__repr__ = global_enum_repr
1796 return cls