blob: 01f431001a9bf063d99e81d0a98bc75aee255823 [file] [log] [blame]
Ethan Furman6b3d64a2013-06-14 16:55:46 -07001import sys
Ethan Furmane03ea372013-09-25 07:14:41 -07002from types import MappingProxyType, DynamicClassAttribute
Ethan Furman7aaeb2a2021-01-25 14:26:19 -08003from builtins import property as _bltin_property, bin as _bltin_bin
Ethan Furman6b3d64a2013-06-14 16:55:46 -07004
Ethan Furmane5754ab2015-09-17 22:03:52 -07005
Ethan Furmanc16595e2016-09-10 23:36:59 -07006__all__ = [
Ethan Furmanb7751062021-03-30 21:17:26 -07007 'EnumType', 'EnumMeta',
Ethan Furman0063ff42020-09-21 17:23:13 -07008 'Enum', 'IntEnum', 'StrEnum', 'Flag', 'IntFlag',
Ethan Furmanc16595e2016-09-10 23:36:59 -07009 'auto', 'unique',
Ethan Furmanc314e602021-01-12 23:47:57 -080010 'property',
Ethan Furman7aaeb2a2021-01-25 14:26:19 -080011 'FlagBoundary', 'STRICT', 'CONFORM', 'EJECT', 'KEEP',
Ethan Furmanb7751062021-03-30 21:17:26 -070012 'global_flag_repr', 'global_enum_repr', 'global_enum',
Ethan Furmanc16595e2016-09-10 23:36:59 -070013 ]
Ethan Furman6b3d64a2013-06-14 16:55:46 -070014
15
Ethan Furman7aaeb2a2021-01-25 14:26:19 -080016# Dummy value for Enum and Flag as there are explicit checks for them
17# before they have been created.
Ethan Furmanb7751062021-03-30 21:17:26 -070018# This is also why there are checks in EnumType like `if Enum is not None`
Ethan Furman7aaeb2a2021-01-25 14:26:19 -080019Enum = Flag = EJECT = None
20
Ethan Furman101e0742013-09-15 12:34:36 -070021def _is_descriptor(obj):
Ethan Furman6d3dfee2020-12-08 12:26:56 -080022 """
23 Returns True if obj is a descriptor, False otherwise.
24 """
Ethan Furman101e0742013-09-15 12:34:36 -070025 return (
26 hasattr(obj, '__get__') or
27 hasattr(obj, '__set__') or
Ethan Furman6d3dfee2020-12-08 12:26:56 -080028 hasattr(obj, '__delete__')
29 )
Ethan Furman101e0742013-09-15 12:34:36 -070030
Ethan Furman6b3d64a2013-06-14 16:55:46 -070031def _is_dunder(name):
Ethan Furman6d3dfee2020-12-08 12:26:56 -080032 """
33 Returns True if a __dunder__ name, False otherwise.
34 """
35 return (
36 len(name) > 4 and
Brennan D Baraban8b914d22019-03-03 14:09:11 -080037 name[:2] == name[-2:] == '__' and
38 name[2] != '_' and
Ethan Furman6d3dfee2020-12-08 12:26:56 -080039 name[-3] != '_'
40 )
Ethan Furman6b3d64a2013-06-14 16:55:46 -070041
42def _is_sunder(name):
Ethan Furman6d3dfee2020-12-08 12:26:56 -080043 """
44 Returns True if a _sunder_ name, False otherwise.
45 """
46 return (
47 len(name) > 2 and
Brennan D Baraban8b914d22019-03-03 14:09:11 -080048 name[0] == name[-1] == '_' and
Ethan Furman6b3d64a2013-06-14 16:55:46 -070049 name[1:2] != '_' and
Ethan Furman6d3dfee2020-12-08 12:26:56 -080050 name[-2:-1] != '_'
51 )
Ethan Furman6b3d64a2013-06-14 16:55:46 -070052
Ethan Furman7cf0aad2020-12-09 17:12:11 -080053def _is_private(cls_name, name):
54 # do not use `re` as `re` imports `enum`
55 pattern = '_%s__' % (cls_name, )
Ethan Furmanec099732021-04-15 06:58:33 -070056 pat_len = len(pattern)
Ethan Furman7cf0aad2020-12-09 17:12:11 -080057 if (
Ethan Furmanec099732021-04-15 06:58:33 -070058 len(name) > pat_len
Ethan Furman7cf0aad2020-12-09 17:12:11 -080059 and name.startswith(pattern)
Ethan Furmanec099732021-04-15 06:58:33 -070060 and name[pat_len:pat_len+1] != ['_']
Ethan Furman7cf0aad2020-12-09 17:12:11 -080061 and (name[-1] != '_' or name[-2] != '_')
62 ):
63 return True
64 else:
65 return False
66
Ethan Furman7aaeb2a2021-01-25 14:26:19 -080067def _is_single_bit(num):
68 """
69 True if only one bit set in num (should be an int)
70 """
71 if num == 0:
72 return False
73 num &= num - 1
74 return num == 0
75
Ethan Furmanc314e602021-01-12 23:47:57 -080076def _make_class_unpicklable(obj):
Ethan Furman6d3dfee2020-12-08 12:26:56 -080077 """
Ethan Furmanc314e602021-01-12 23:47:57 -080078 Make the given obj un-picklable.
79
80 obj should be either a dictionary, on an Enum
Ethan Furman6d3dfee2020-12-08 12:26:56 -080081 """
Ethan Furmanca1b7942014-02-08 11:36:27 -080082 def _break_on_call_reduce(self, proto):
Ethan Furman6b3d64a2013-06-14 16:55:46 -070083 raise TypeError('%r cannot be pickled' % self)
Ethan Furmanc314e602021-01-12 23:47:57 -080084 if isinstance(obj, dict):
85 obj['__reduce_ex__'] = _break_on_call_reduce
86 obj['__module__'] = '<unknown>'
87 else:
88 setattr(obj, '__reduce_ex__', _break_on_call_reduce)
89 setattr(obj, '__module__', '<unknown>')
Ethan Furman6b3d64a2013-06-14 16:55:46 -070090
Ethan Furman7aaeb2a2021-01-25 14:26:19 -080091def _iter_bits_lsb(num):
92 while num:
93 b = num & (~num + 1)
94 yield b
95 num ^= b
96
97def bin(num, max_bits=None):
98 """
99 Like built-in bin(), except negative values are represented in
100 twos-compliment, and the leading bit always indicates sign
101 (0=positive, 1=negative).
102
103 >>> bin(10)
104 '0b0 1010'
105 >>> bin(~10) # ~10 is -11
106 '0b1 0101'
107 """
108
109 ceiling = 2 ** (num).bit_length()
110 if num >= 0:
111 s = _bltin_bin(num + ceiling).replace('1', '0', 1)
112 else:
113 s = _bltin_bin(~num ^ (ceiling - 1) + ceiling)
114 sign = s[:3]
115 digits = s[3:]
116 if max_bits is not None:
117 if len(digits) < max_bits:
118 digits = (sign[-1] * max_bits + digits)[-max_bits:]
119 return "%s %s" % (sign, digits)
120
121
Ethan Furman3515dcc2016-09-18 13:15:41 -0700122_auto_null = object()
Ethan Furmanc16595e2016-09-10 23:36:59 -0700123class auto:
124 """
125 Instances are replaced with an appropriate value in Enum class suites.
126 """
Ethan Furman3515dcc2016-09-18 13:15:41 -0700127 value = _auto_null
Ethan Furmanc16595e2016-09-10 23:36:59 -0700128
Ethan Furmanc314e602021-01-12 23:47:57 -0800129class property(DynamicClassAttribute):
130 """
131 This is a descriptor, used to define attributes that act differently
132 when accessed through an enum member and through an enum class.
133 Instance access is the same as property(), but access to an attribute
134 through the enum class will instead look in the class' _member_map_ for
135 a corresponding enum member.
136 """
137
138 def __get__(self, instance, ownerclass=None):
139 if instance is None:
140 try:
141 return ownerclass._member_map_[self.name]
142 except KeyError:
Ethan Furman7aaeb2a2021-01-25 14:26:19 -0800143 raise AttributeError(
Ethan Furmand65b9032021-02-08 17:32:38 -0800144 '%s: no class attribute %r' % (ownerclass.__name__, self.name)
Ethan Furman7aaeb2a2021-01-25 14:26:19 -0800145 )
Ethan Furmanc314e602021-01-12 23:47:57 -0800146 else:
147 if self.fget is None:
Ethan Furmand65b9032021-02-08 17:32:38 -0800148 # check for member
149 if self.name in ownerclass._member_map_:
150 import warnings
151 warnings.warn(
152 "accessing one member from another is not supported, "
Ethan Furman44e580f2021-03-03 09:54:30 -0800153 " and will be disabled in 3.12",
Ethan Furmand65b9032021-02-08 17:32:38 -0800154 DeprecationWarning,
155 stacklevel=2,
156 )
157 return ownerclass._member_map_[self.name]
Ethan Furman7aaeb2a2021-01-25 14:26:19 -0800158 raise AttributeError(
Ethan Furmand65b9032021-02-08 17:32:38 -0800159 '%s: no instance attribute %r' % (ownerclass.__name__, self.name)
Ethan Furman7aaeb2a2021-01-25 14:26:19 -0800160 )
Ethan Furmanc314e602021-01-12 23:47:57 -0800161 else:
162 return self.fget(instance)
163
164 def __set__(self, instance, value):
165 if self.fset is None:
Ethan Furman7aaeb2a2021-01-25 14:26:19 -0800166 raise AttributeError(
Ethan Furmand65b9032021-02-08 17:32:38 -0800167 "%s: cannot set instance attribute %r" % (self.clsname, self.name)
Ethan Furman7aaeb2a2021-01-25 14:26:19 -0800168 )
Ethan Furmanc314e602021-01-12 23:47:57 -0800169 else:
170 return self.fset(instance, value)
171
172 def __delete__(self, instance):
173 if self.fdel is None:
Ethan Furman7aaeb2a2021-01-25 14:26:19 -0800174 raise AttributeError(
Ethan Furmand65b9032021-02-08 17:32:38 -0800175 "%s: cannot delete instance attribute %r" % (self.clsname, self.name)
Ethan Furman7aaeb2a2021-01-25 14:26:19 -0800176 )
Ethan Furmanc314e602021-01-12 23:47:57 -0800177 else:
178 return self.fdel(instance)
179
180 def __set_name__(self, ownerclass, name):
181 self.name = name
182 self.clsname = ownerclass.__name__
183
184
185class _proto_member:
186 """
187 intermediate step for enum members between class execution and final creation
188 """
189
190 def __init__(self, value):
191 self.value = value
192
193 def __set_name__(self, enum_class, member_name):
194 """
195 convert each quasi-member into an instance of the new enum class
196 """
197 # first step: remove ourself from enum_class
198 delattr(enum_class, member_name)
199 # second step: create member based on enum_class
200 value = self.value
201 if not isinstance(value, tuple):
202 args = (value, )
203 else:
204 args = value
205 if enum_class._member_type_ is tuple: # special case for tuple enums
206 args = (args, ) # wrap it one more time
207 if not enum_class._use_args_:
208 enum_member = enum_class._new_member_(enum_class)
209 if not hasattr(enum_member, '_value_'):
210 enum_member._value_ = value
211 else:
212 enum_member = enum_class._new_member_(enum_class, *args)
213 if not hasattr(enum_member, '_value_'):
214 if enum_class._member_type_ is object:
215 enum_member._value_ = value
216 else:
Ethan Furman7aaeb2a2021-01-25 14:26:19 -0800217 try:
218 enum_member._value_ = enum_class._member_type_(*args)
219 except Exception as exc:
220 raise TypeError(
221 '_value_ not set in __new__, unable to create it'
222 ) from None
Ethan Furmanc314e602021-01-12 23:47:57 -0800223 value = enum_member._value_
224 enum_member._name_ = member_name
225 enum_member.__objclass__ = enum_class
226 enum_member.__init__(*args)
Ethan Furman7aaeb2a2021-01-25 14:26:19 -0800227 enum_member._sort_order_ = len(enum_class._member_names_)
Ethan Furmanc314e602021-01-12 23:47:57 -0800228 # If another member with the same value was already defined, the
229 # new member becomes an alias to the existing one.
230 for name, canonical_member in enum_class._member_map_.items():
231 if canonical_member._value_ == enum_member._value_:
232 enum_member = canonical_member
233 break
234 else:
Ethan Furman7aaeb2a2021-01-25 14:26:19 -0800235 # this could still be an alias if the value is multi-bit and the
236 # class is a flag class
237 if (
238 Flag is None
239 or not issubclass(enum_class, Flag)
240 ):
241 # no other instances found, record this member in _member_names_
242 enum_class._member_names_.append(member_name)
243 elif (
244 Flag is not None
245 and issubclass(enum_class, Flag)
246 and _is_single_bit(value)
247 ):
248 # no other instances found, record this member in _member_names_
249 enum_class._member_names_.append(member_name)
Ethan Furmanc314e602021-01-12 23:47:57 -0800250 # get redirect in place before adding to _member_map_
251 # but check for other instances in parent classes first
252 need_override = False
253 descriptor = None
254 for base in enum_class.__mro__[1:]:
255 descriptor = base.__dict__.get(member_name)
256 if descriptor is not None:
257 if isinstance(descriptor, (property, DynamicClassAttribute)):
258 break
259 else:
260 need_override = True
261 # keep looking for an enum.property
262 if descriptor and not need_override:
263 # previous enum.property found, no further action needed
264 pass
265 else:
266 redirect = property()
267 redirect.__set_name__(enum_class, member_name)
268 if descriptor and need_override:
269 # previous enum.property found, but some other inherited attribute
270 # is in the way; copy fget, fset, fdel to this one
271 redirect.fget = descriptor.fget
272 redirect.fset = descriptor.fset
273 redirect.fdel = descriptor.fdel
274 setattr(enum_class, member_name, redirect)
275 # now add to _member_map_ (even aliases)
276 enum_class._member_map_[member_name] = enum_member
277 try:
278 # This may fail if value is not hashable. We can't add the value
279 # to the map, and by-value lookups for this value will be
280 # linear.
Ethan Furman7aaeb2a2021-01-25 14:26:19 -0800281 enum_class._value2member_map_.setdefault(value, enum_member)
Ethan Furmanc314e602021-01-12 23:47:57 -0800282 except TypeError:
Ethan Furman6bd92882021-04-27 13:05:08 -0700283 # keep track of the value in a list so containment checks are quick
284 enum_class._unhashable_values_.append(value)
Ethan Furmanc314e602021-01-12 23:47:57 -0800285
Ethan Furman101e0742013-09-15 12:34:36 -0700286
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700287class _EnumDict(dict):
Ethan Furman6d3dfee2020-12-08 12:26:56 -0800288 """
289 Track enum member order and ensure member names are not reused.
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700290
Ethan Furmanb7751062021-03-30 21:17:26 -0700291 EnumType will use the names found in self._member_names as the
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700292 enumeration member names.
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700293 """
294 def __init__(self):
295 super().__init__()
296 self._member_names = []
Ethan Furmanc16595e2016-09-10 23:36:59 -0700297 self._last_values = []
Ethan Furmana4b1bb42018-01-22 07:56:37 -0800298 self._ignore = []
Ethan Onstottd9a43e22020-04-28 13:20:55 -0400299 self._auto_called = False
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700300
301 def __setitem__(self, key, value):
Ethan Furman6d3dfee2020-12-08 12:26:56 -0800302 """
303 Changes anything not dundered or not a descriptor.
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700304
305 If an enum member name is used twice, an error is raised; duplicate
306 values are not checked for.
307
308 Single underscore (sunder) names are reserved.
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700309 """
Ethan Furman7cf0aad2020-12-09 17:12:11 -0800310 if _is_private(self._cls_name, key):
311 # do nothing, name will be a normal attribute
312 pass
313 elif _is_sunder(key):
Ethan Furmanee47e5c2016-08-31 00:12:15 -0700314 if key not in (
Ethan Furman0dca5eb2021-04-15 06:49:54 -0700315 '_order_',
Ethan Furmana4b1bb42018-01-22 07:56:37 -0800316 '_generate_next_value_', '_missing_', '_ignore_',
Ethan Furman7aaeb2a2021-01-25 14:26:19 -0800317 '_iter_member_', '_iter_member_by_value_', '_iter_member_by_def_',
Ethan Furmanee47e5c2016-08-31 00:12:15 -0700318 ):
Ethan Furman6d3dfee2020-12-08 12:26:56 -0800319 raise ValueError(
320 '_sunder_ names, such as %r, are reserved for future Enum use'
321 % (key, )
322 )
Ethan Furmanc16595e2016-09-10 23:36:59 -0700323 if key == '_generate_next_value_':
Ethan Onstottd9a43e22020-04-28 13:20:55 -0400324 # check if members already defined as auto()
325 if self._auto_called:
326 raise TypeError("_generate_next_value_ must be defined before members")
Ethan Furmanb7751062021-03-30 21:17:26 -0700327 _gnv = value.__func__ if isinstance(value, staticmethod) else value
328 setattr(self, '_generate_next_value', _gnv)
Ethan Furmana4b1bb42018-01-22 07:56:37 -0800329 elif key == '_ignore_':
330 if isinstance(value, str):
331 value = value.replace(',',' ').split()
332 else:
333 value = list(value)
334 self._ignore = value
335 already = set(value) & set(self._member_names)
336 if already:
Ethan Furman6d3dfee2020-12-08 12:26:56 -0800337 raise ValueError(
338 '_ignore_ cannot specify already set names: %r'
339 % (already, )
340 )
Ethan Furman101e0742013-09-15 12:34:36 -0700341 elif _is_dunder(key):
Ethan Furmane8e61272016-08-20 07:19:31 -0700342 if key == '__order__':
343 key = '_order_'
Ethan Furman101e0742013-09-15 12:34:36 -0700344 elif key in self._member_names:
345 # descriptor overwriting an enum?
Ethan Furmana6582872020-12-10 13:07:00 -0800346 raise TypeError('%r already defined as: %r' % (key, self[key]))
Ethan Furmana4b1bb42018-01-22 07:56:37 -0800347 elif key in self._ignore:
348 pass
Ethan Furman101e0742013-09-15 12:34:36 -0700349 elif not _is_descriptor(value):
350 if key in self:
351 # enum overwriting a descriptor?
Ethan Furmanee47e5c2016-08-31 00:12:15 -0700352 raise TypeError('%r already defined as: %r' % (key, self[key]))
Ethan Furmanc16595e2016-09-10 23:36:59 -0700353 if isinstance(value, auto):
Ethan Furman3515dcc2016-09-18 13:15:41 -0700354 if value.value == _auto_null:
Ethan Furman6d3dfee2020-12-08 12:26:56 -0800355 value.value = self._generate_next_value(
Ethan Furman7aaeb2a2021-01-25 14:26:19 -0800356 key, 1, len(self._member_names), self._last_values[:],
Ethan Furman6d3dfee2020-12-08 12:26:56 -0800357 )
Ethan Furmanfc23a942020-09-16 12:37:54 -0700358 self._auto_called = True
Ethan Furman3515dcc2016-09-18 13:15:41 -0700359 value = value.value
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700360 self._member_names.append(key)
Ethan Furmanc16595e2016-09-10 23:36:59 -0700361 self._last_values.append(value)
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700362 super().__setitem__(key, value)
363
Ethan Furmana6582872020-12-10 13:07:00 -0800364 def update(self, members, **more_members):
365 try:
366 for name in members.keys():
367 self[name] = members[name]
368 except AttributeError:
369 for name, value in members:
370 self[name] = value
371 for name, value in more_members.items():
372 self[name] = value
373
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700374
Ethan Furmanb7751062021-03-30 21:17:26 -0700375class EnumType(type):
Ethan Furman6d3dfee2020-12-08 12:26:56 -0800376 """
377 Metaclass for Enum
378 """
Ethan Furman7aaeb2a2021-01-25 14:26:19 -0800379
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700380 @classmethod
Ethan Furman6ec0ade2020-12-24 10:05:02 -0800381 def __prepare__(metacls, cls, bases, **kwds):
Ethan Furman3064dbf2020-09-16 07:11:57 -0700382 # check that previous enum members do not exist
383 metacls._check_for_existing_members(cls, bases)
Ethan Furmanee47e5c2016-08-31 00:12:15 -0700384 # create the namespace dict
385 enum_dict = _EnumDict()
Ethan Furman7cf0aad2020-12-09 17:12:11 -0800386 enum_dict._cls_name = cls
Ethan Furmanee47e5c2016-08-31 00:12:15 -0700387 # inherit previous flags and _generate_next_value_ function
Ethan Furman3064dbf2020-09-16 07:11:57 -0700388 member_type, first_enum = metacls._get_mixins_(cls, bases)
Ethan Furmanee47e5c2016-08-31 00:12:15 -0700389 if first_enum is not None:
Ethan Furman6d3dfee2020-12-08 12:26:56 -0800390 enum_dict['_generate_next_value_'] = getattr(
391 first_enum, '_generate_next_value_', None,
392 )
Ethan Furmanee47e5c2016-08-31 00:12:15 -0700393 return enum_dict
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700394
Ethan Furmana02cb472021-04-21 10:20:44 -0700395 def __new__(metacls, cls, bases, classdict, *, boundary=None, _simple=False, **kwds):
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700396 # an Enum class is final once enumeration items have been defined; it
397 # cannot be mixed with other types (int, float, etc.) if it has an
398 # inherited __new__ unless a new __new__ is defined (or the resulting
399 # class will fail).
Ethan Furmana4b1bb42018-01-22 07:56:37 -0800400 #
401 # remove any keys listed in _ignore_
Ethan Furmana02cb472021-04-21 10:20:44 -0700402 if _simple:
403 return super().__new__(metacls, cls, bases, classdict, **kwds)
Ethan Furmana4b1bb42018-01-22 07:56:37 -0800404 classdict.setdefault('_ignore_', []).append('_ignore_')
405 ignore = classdict['_ignore_']
406 for key in ignore:
407 classdict.pop(key, None)
Ethan Furmanc314e602021-01-12 23:47:57 -0800408 #
409 # grab member names
410 member_names = classdict._member_names
411 #
412 # check for illegal enum names (any others?)
413 invalid_names = set(member_names) & {'mro', ''}
414 if invalid_names:
415 raise ValueError('Invalid enum member name: {0}'.format(
416 ','.join(invalid_names)))
417 #
418 # adjust the sunders
419 _order_ = classdict.pop('_order_', None)
420 # convert to normal dict
421 classdict = dict(classdict.items())
422 #
423 # data type of member and the controlling Enum class
Ethan Furman3064dbf2020-09-16 07:11:57 -0700424 member_type, first_enum = metacls._get_mixins_(cls, bases)
Ethan Furmanc2667362020-12-07 00:17:31 -0800425 __new__, save_new, use_args = metacls._find_new_(
426 classdict, member_type, first_enum,
427 )
Ethan Furmanc314e602021-01-12 23:47:57 -0800428 classdict['_new_member_'] = __new__
429 classdict['_use_args_'] = use_args
430 #
431 # convert future enum members into temporary _proto_members
Ethan Furman7aaeb2a2021-01-25 14:26:19 -0800432 # and record integer values in case this will be a Flag
433 flag_mask = 0
Ethan Furmanc314e602021-01-12 23:47:57 -0800434 for name in member_names:
Ethan Furman7aaeb2a2021-01-25 14:26:19 -0800435 value = classdict[name]
436 if isinstance(value, int):
437 flag_mask |= value
438 classdict[name] = _proto_member(value)
Ethan Furmanc314e602021-01-12 23:47:57 -0800439 #
Ethan Furman7aaeb2a2021-01-25 14:26:19 -0800440 # house-keeping structures
Ethan Furmanc314e602021-01-12 23:47:57 -0800441 classdict['_member_names_'] = []
442 classdict['_member_map_'] = {}
443 classdict['_value2member_map_'] = {}
Ethan Furman6bd92882021-04-27 13:05:08 -0700444 classdict['_unhashable_values_'] = []
Ethan Furmanc314e602021-01-12 23:47:57 -0800445 classdict['_member_type_'] = member_type
446 #
Ethan Furman7aaeb2a2021-01-25 14:26:19 -0800447 # Flag structures (will be removed if final class is not a Flag
448 classdict['_boundary_'] = (
449 boundary
450 or getattr(first_enum, '_boundary_', None)
451 )
452 classdict['_flag_mask_'] = flag_mask
453 classdict['_all_bits_'] = 2 ** ((flag_mask).bit_length()) - 1
454 classdict['_inverted_'] = None
455 #
Ethan Furman2da95042014-03-03 12:42:52 -0800456 # If a custom type is mixed into the Enum, and it does not know how
457 # to pickle itself, pickle.dumps will succeed but pickle.loads will
458 # fail. Rather than have the error show up later and possibly far
459 # from the source, sabotage the pickle protocol for this class so
460 # that pickle.dumps also fails.
461 #
462 # However, if the new class implements its own __reduce_ex__, do not
463 # sabotage -- it's on them to make sure it works correctly. We use
464 # __reduce_ex__ instead of any of the others as it is preferred by
465 # pickle over __reduce__, and it handles all pickle protocols.
466 if '__reduce_ex__' not in classdict:
Ethan Furmandc870522014-02-18 12:37:12 -0800467 if member_type is not object:
468 methods = ('__getnewargs_ex__', '__getnewargs__',
469 '__reduce_ex__', '__reduce__')
Ethan Furman2da95042014-03-03 12:42:52 -0800470 if not any(m in member_type.__dict__ for m in methods):
Ethan Furmanc314e602021-01-12 23:47:57 -0800471 _make_class_unpicklable(classdict)
472 #
473 # create a default docstring if one has not been provided
474 if '__doc__' not in classdict:
475 classdict['__doc__'] = 'An enumeration.'
476 try:
477 exc = None
478 enum_class = super().__new__(metacls, cls, bases, classdict, **kwds)
479 except RuntimeError as e:
480 # any exceptions raised by member.__new__ will get converted to a
481 # RuntimeError, so get that original exception back and raise it instead
482 exc = e.__cause__ or e
483 if exc is not None:
484 raise exc
485 #
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700486 # double check that repr and friends are not the mixin's or various
487 # things break (such as pickle)
Ethan Furman22415ad2020-09-15 16:28:25 -0700488 # however, if the method is defined in the Enum itself, don't replace
489 # it
Ethan Furmandc870522014-02-18 12:37:12 -0800490 for name in ('__repr__', '__str__', '__format__', '__reduce_ex__'):
Ethan Furman22415ad2020-09-15 16:28:25 -0700491 if name in classdict:
492 continue
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700493 class_method = getattr(enum_class, name)
494 obj_method = getattr(member_type, name, None)
495 enum_method = getattr(first_enum, name, None)
496 if obj_method is not None and obj_method is class_method:
497 setattr(enum_class, name, enum_method)
Ethan Furmanc314e602021-01-12 23:47:57 -0800498 #
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700499 # replace any other __new__ with our own (as long as Enum is not None,
500 # anyway) -- again, this is to support pickle
501 if Enum is not None:
502 # if the user defined their own __new__, save it before it gets
503 # clobbered in case they subclass later
504 if save_new:
505 enum_class.__new_member__ = __new__
506 enum_class.__new__ = Enum.__new__
Ethan Furmanc314e602021-01-12 23:47:57 -0800507 #
Ethan Furmane8e61272016-08-20 07:19:31 -0700508 # py3 support for definition order (helps keep py2/py3 code in sync)
Ethan Furman7aaeb2a2021-01-25 14:26:19 -0800509 #
510 # _order_ checking is spread out into three/four steps
511 # - if enum_class is a Flag:
512 # - remove any non-single-bit flags from _order_
513 # - remove any aliases from _order_
514 # - check that _order_ and _member_names_ match
515 #
516 # step 1: ensure we have a list
Ethan Furmane8e61272016-08-20 07:19:31 -0700517 if _order_ is not None:
518 if isinstance(_order_, str):
519 _order_ = _order_.replace(',', ' ').split()
Ethan Furman7aaeb2a2021-01-25 14:26:19 -0800520 #
521 # remove Flag structures if final class is not a Flag
522 if (
523 Flag is None and cls != 'Flag'
524 or Flag is not None and not issubclass(enum_class, Flag)
525 ):
526 delattr(enum_class, '_boundary_')
527 delattr(enum_class, '_flag_mask_')
528 delattr(enum_class, '_all_bits_')
529 delattr(enum_class, '_inverted_')
530 elif Flag is not None and issubclass(enum_class, Flag):
531 # ensure _all_bits_ is correct and there are no missing flags
532 single_bit_total = 0
533 multi_bit_total = 0
534 for flag in enum_class._member_map_.values():
535 flag_value = flag._value_
536 if _is_single_bit(flag_value):
537 single_bit_total |= flag_value
538 else:
539 # multi-bit flags are considered aliases
540 multi_bit_total |= flag_value
541 if enum_class._boundary_ is not KEEP:
542 missed = list(_iter_bits_lsb(multi_bit_total & ~single_bit_total))
543 if missed:
544 raise TypeError(
545 'invalid Flag %r -- missing values: %s'
546 % (cls, ', '.join((str(i) for i in missed)))
547 )
548 enum_class._flag_mask_ = single_bit_total
549 #
550 # set correct __iter__
551 member_list = [m._value_ for m in enum_class]
552 if member_list != sorted(member_list):
553 enum_class._iter_member_ = enum_class._iter_member_by_def_
554 if _order_:
555 # _order_ step 2: remove any items from _order_ that are not single-bit
556 _order_ = [
557 o
558 for o in _order_
559 if o not in enum_class._member_map_ or _is_single_bit(enum_class[o]._value_)
560 ]
561 #
562 if _order_:
563 # _order_ step 3: remove aliases from _order_
564 _order_ = [
565 o
566 for o in _order_
567 if (
568 o not in enum_class._member_map_
569 or
570 (o in enum_class._member_map_ and o in enum_class._member_names_)
571 )]
572 # _order_ step 4: verify that _order_ and _member_names_ match
Ethan Furmane8e61272016-08-20 07:19:31 -0700573 if _order_ != enum_class._member_names_:
Ethan Furman7aaeb2a2021-01-25 14:26:19 -0800574 raise TypeError(
575 'member order does not match _order_:\n%r\n%r'
576 % (enum_class._member_names_, _order_)
577 )
Ethan Furmanc314e602021-01-12 23:47:57 -0800578 #
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700579 return enum_class
580
Ethan Furman5de67b12016-04-13 23:52:09 -0700581 def __bool__(self):
582 """
583 classes/types should always be True.
584 """
585 return True
586
Ethan Furman7aaeb2a2021-01-25 14:26:19 -0800587 def __call__(cls, value, names=None, *, module=None, qualname=None, type=None, start=1, boundary=None):
Ethan Furman6d3dfee2020-12-08 12:26:56 -0800588 """
589 Either returns an existing member, or creates a new enum class.
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700590
591 This method is used both when an enum class is given a value to match
592 to an enumeration member (i.e. Color(3)) and for the functional API
Ethan Furman23bb6f42016-11-21 09:22:05 -0800593 (i.e. Color = Enum('Color', names='RED GREEN BLUE')).
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700594
Ethan Furman2da95042014-03-03 12:42:52 -0800595 When used for the functional API:
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700596
Ethan Furman2da95042014-03-03 12:42:52 -0800597 `value` will be the name of the new class.
598
599 `names` should be either a string of white-space/comma delimited names
Ethan Furmand9925a12014-09-16 20:35:55 -0700600 (values will start at `start`), or an iterator/mapping of name, value pairs.
Ethan Furman2da95042014-03-03 12:42:52 -0800601
602 `module` should be set to the module this class is being created in;
603 if it is not set, an attempt to find that module will be made, but if
604 it fails the class will not be picklable.
605
606 `qualname` should be set to the actual location this class can be found
607 at in its module; by default it is set to the global scope. If this is
608 not correct, unpickling will fail in some circumstances.
609
610 `type`, if set, will be mixed in as the first base class.
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700611 """
612 if names is None: # simple value lookup
613 return cls.__new__(cls, value)
614 # otherwise, functional API: we're creating a new Enum type
Ethan Furman6d3dfee2020-12-08 12:26:56 -0800615 return cls._create_(
616 value,
617 names,
618 module=module,
619 qualname=qualname,
620 type=type,
621 start=start,
Ethan Furman7aaeb2a2021-01-25 14:26:19 -0800622 boundary=boundary,
Ethan Furman6d3dfee2020-12-08 12:26:56 -0800623 )
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700624
625 def __contains__(cls, member):
Rahul Jha94306522018-09-10 23:51:04 +0530626 if not isinstance(member, Enum):
Ethan Furman6bd92882021-04-27 13:05:08 -0700627 import warnings
628 warnings.warn(
629 "in 3.12 __contains__ will no longer raise TypeError, but will return True or\n"
630 "False depending on whether the value is a member or the value of a member",
631 DeprecationWarning,
632 stacklevel=2,
633 )
Rahul Jha94306522018-09-10 23:51:04 +0530634 raise TypeError(
635 "unsupported operand type(s) for 'in': '%s' and '%s'" % (
636 type(member).__qualname__, cls.__class__.__qualname__))
Ethan Furman0081f232014-09-16 17:31:23 -0700637 return isinstance(member, cls) and member._name_ in cls._member_map_
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700638
Ethan Furman64a99722013-09-22 16:18:19 -0700639 def __delattr__(cls, attr):
640 # nicer error message when someone tries to delete an attribute
641 # (see issue19025).
642 if attr in cls._member_map_:
Ethan Furman6d3dfee2020-12-08 12:26:56 -0800643 raise AttributeError("%s: cannot delete Enum member %r." % (cls.__name__, attr))
Ethan Furman64a99722013-09-22 16:18:19 -0700644 super().__delattr__(attr)
645
Ethan Furman388a3922013-08-12 06:51:41 -0700646 def __dir__(self):
Ethan Furman6d3dfee2020-12-08 12:26:56 -0800647 return (
648 ['__class__', '__doc__', '__members__', '__module__']
649 + self._member_names_
650 )
Ethan Furman388a3922013-08-12 06:51:41 -0700651
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700652 def __getattr__(cls, name):
Ethan Furman6d3dfee2020-12-08 12:26:56 -0800653 """
654 Return the enum member matching `name`
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700655
656 We use __getattr__ instead of descriptors or inserting into the enum
657 class' __dict__ in order to support `name` and `value` being both
658 properties for enum members (which live in the class' __dict__) and
659 enum members themselves.
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700660 """
661 if _is_dunder(name):
662 raise AttributeError(name)
663 try:
Ethan Furman520ad572013-07-19 19:47:21 -0700664 return cls._member_map_[name]
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700665 except KeyError:
666 raise AttributeError(name) from None
667
668 def __getitem__(cls, name):
Ethan Furman520ad572013-07-19 19:47:21 -0700669 return cls._member_map_[name]
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700670
671 def __iter__(cls):
Ethan Furman6d3dfee2020-12-08 12:26:56 -0800672 """
673 Returns members in definition order.
674 """
Ethan Furman520ad572013-07-19 19:47:21 -0700675 return (cls._member_map_[name] for name in cls._member_names_)
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700676
677 def __len__(cls):
Ethan Furman520ad572013-07-19 19:47:21 -0700678 return len(cls._member_names_)
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700679
Ethan Furmanc314e602021-01-12 23:47:57 -0800680 @_bltin_property
Ethan Furman2131a4a2013-09-14 18:11:24 -0700681 def __members__(cls):
Ethan Furman6d3dfee2020-12-08 12:26:56 -0800682 """
683 Returns a mapping of member name->value.
Ethan Furman2131a4a2013-09-14 18:11:24 -0700684
685 This mapping lists all enum members, including aliases. Note that this
686 is a read-only view of the internal mapping.
Ethan Furman2131a4a2013-09-14 18:11:24 -0700687 """
688 return MappingProxyType(cls._member_map_)
689
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700690 def __repr__(cls):
691 return "<enum %r>" % cls.__name__
692
Ethan Furman2131a4a2013-09-14 18:11:24 -0700693 def __reversed__(cls):
Ethan Furman6d3dfee2020-12-08 12:26:56 -0800694 """
695 Returns members in reverse definition order.
696 """
Ethan Furman2131a4a2013-09-14 18:11:24 -0700697 return (cls._member_map_[name] for name in reversed(cls._member_names_))
698
Ethan Furmanf203f2d2013-09-06 07:16:48 -0700699 def __setattr__(cls, name, value):
Ethan Furman6d3dfee2020-12-08 12:26:56 -0800700 """
701 Block attempts to reassign Enum members.
Ethan Furmanf203f2d2013-09-06 07:16:48 -0700702
703 A simple assignment to the class namespace only changes one of the
704 several possible ways to get an Enum member from the Enum class,
705 resulting in an inconsistent Enumeration.
Ethan Furmanf203f2d2013-09-06 07:16:48 -0700706 """
707 member_map = cls.__dict__.get('_member_map_', {})
708 if name in member_map:
Ethan Furmana02cb472021-04-21 10:20:44 -0700709 raise AttributeError('Cannot reassign member %r.' % (name, ))
Ethan Furmanf203f2d2013-09-06 07:16:48 -0700710 super().__setattr__(name, value)
711
Ethan Furman7aaeb2a2021-01-25 14:26:19 -0800712 def _create_(cls, class_name, names, *, module=None, qualname=None, type=None, start=1, boundary=None):
Ethan Furman6d3dfee2020-12-08 12:26:56 -0800713 """
714 Convenience method to create a new Enum class.
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700715
716 `names` can be:
717
718 * A string containing member names, separated either with spaces or
Ethan Furmand9925a12014-09-16 20:35:55 -0700719 commas. Values are incremented by 1 from `start`.
720 * An iterable of member names. Values are incremented by 1 from `start`.
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700721 * An iterable of (member name, value) pairs.
Ethan Furmand9925a12014-09-16 20:35:55 -0700722 * A mapping of member name -> value pairs.
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700723 """
724 metacls = cls.__class__
725 bases = (cls, ) if type is None else (type, cls)
Ethan Furman3064dbf2020-09-16 07:11:57 -0700726 _, first_enum = cls._get_mixins_(cls, bases)
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700727 classdict = metacls.__prepare__(class_name, bases)
728
729 # special processing needed for names?
730 if isinstance(names, str):
731 names = names.replace(',', ' ').split()
Dong-hee Nadcc8ce42017-06-22 01:52:32 +0900732 if isinstance(names, (tuple, list)) and names and isinstance(names[0], str):
Ethan Furmanee47e5c2016-08-31 00:12:15 -0700733 original_names, names = names, []
Ethan Furmanc16595e2016-09-10 23:36:59 -0700734 last_values = []
Ethan Furmanee47e5c2016-08-31 00:12:15 -0700735 for count, name in enumerate(original_names):
Ethan Furmanc16595e2016-09-10 23:36:59 -0700736 value = first_enum._generate_next_value_(name, start, count, last_values[:])
737 last_values.append(value)
738 names.append((name, value))
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700739
740 # Here, names is either an iterable of (name, value) or a mapping.
741 for item in names:
742 if isinstance(item, str):
743 member_name, member_value = item, names[item]
744 else:
745 member_name, member_value = item
746 classdict[member_name] = member_value
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700747
748 # TODO: replace the frame hack if a blessed way to know the calling
749 # module is ever developed
750 if module is None:
751 try:
752 module = sys._getframe(2).f_globals['__name__']
Pablo Galindo293dd232019-11-19 21:34:03 +0000753 except (AttributeError, ValueError, KeyError):
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700754 pass
755 if module is None:
Ethan Furmanc314e602021-01-12 23:47:57 -0800756 _make_class_unpicklable(classdict)
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700757 else:
Ethan Furmanc314e602021-01-12 23:47:57 -0800758 classdict['__module__'] = module
Ethan Furmanca1b7942014-02-08 11:36:27 -0800759 if qualname is not None:
Ethan Furmanc314e602021-01-12 23:47:57 -0800760 classdict['__qualname__'] = qualname
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700761
Ethan Furman7aaeb2a2021-01-25 14:26:19 -0800762 return metacls.__new__(metacls, class_name, bases, classdict, boundary=boundary)
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700763
Ethan Furmana02cb472021-04-21 10:20:44 -0700764 def _convert_(cls, name, module, filter, source=None, *, boundary=None):
765
orlnub1230fb9fad2018-09-12 20:28:53 +0300766 """
767 Create a new Enum subclass that replaces a collection of global constants
768 """
769 # convert all constants from source (or module) that pass filter() to
770 # a new Enum called name, and export the enum and its members back to
771 # module;
772 # also, replace the __reduce_ex__ method so unpickling works in
773 # previous Python versions
Ethan Furmanb7751062021-03-30 21:17:26 -0700774 module_globals = sys.modules[module].__dict__
orlnub1230fb9fad2018-09-12 20:28:53 +0300775 if source:
Ethan Furmanb7751062021-03-30 21:17:26 -0700776 source = source.__dict__
orlnub1230fb9fad2018-09-12 20:28:53 +0300777 else:
778 source = module_globals
779 # _value2member_map_ is populated in the same order every time
780 # for a consistent reverse mapping of number to name when there
781 # are multiple names for the same number.
782 members = [
783 (name, value)
784 for name, value in source.items()
785 if filter(name)]
786 try:
787 # sort by value
788 members.sort(key=lambda t: (t[1], t[0]))
789 except TypeError:
790 # unless some values aren't comparable, in which case sort by name
791 members.sort(key=lambda t: t[0])
Ethan Furmana02cb472021-04-21 10:20:44 -0700792 body = {t[0]: t[1] for t in members}
793 body['__module__'] = module
794 tmp_cls = type(name, (object, ), body)
795 cls = _simple_enum(etype=cls, boundary=boundary or KEEP)(tmp_cls)
orlnub1230fb9fad2018-09-12 20:28:53 +0300796 cls.__reduce_ex__ = _reduce_ex_by_name
Ethan Furmanb7751062021-03-30 21:17:26 -0700797 global_enum(cls)
orlnub1230fb9fad2018-09-12 20:28:53 +0300798 module_globals[name] = cls
799 return cls
800
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700801 @staticmethod
Ethan Furman3064dbf2020-09-16 07:11:57 -0700802 def _check_for_existing_members(class_name, bases):
803 for chain in bases:
804 for base in chain.__mro__:
805 if issubclass(base, Enum) and base._member_names_:
Ethan Furman6d3dfee2020-12-08 12:26:56 -0800806 raise TypeError(
807 "%s: cannot extend enumeration %r"
808 % (class_name, base.__name__)
809 )
Ethan Furman3064dbf2020-09-16 07:11:57 -0700810
811 @staticmethod
812 def _get_mixins_(class_name, bases):
Ethan Furman6d3dfee2020-12-08 12:26:56 -0800813 """
814 Returns the type for creating enum members, and the first inherited
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700815 enum class.
816
817 bases: the tuple of bases that was given to __new__
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700818 """
819 if not bases:
820 return object, Enum
821
Ethan Furman5bdab642018-09-21 19:03:09 -0700822 def _find_data_type(bases):
Ethan Furmanbff01f32020-09-15 15:56:26 -0700823 data_types = []
Ethan Furman5bdab642018-09-21 19:03:09 -0700824 for chain in bases:
Ethan Furmanbff01f32020-09-15 15:56:26 -0700825 candidate = None
Ethan Furman5bdab642018-09-21 19:03:09 -0700826 for base in chain.__mro__:
827 if base is object:
828 continue
Ethan Furmanc2667362020-12-07 00:17:31 -0800829 elif issubclass(base, Enum):
830 if base._member_type_ is not object:
831 data_types.append(base._member_type_)
832 break
Ethan Furman5bdab642018-09-21 19:03:09 -0700833 elif '__new__' in base.__dict__:
Ethan Furmancd453852018-10-05 23:29:36 -0700834 if issubclass(base, Enum):
Ethan Furman5bdab642018-09-21 19:03:09 -0700835 continue
Ethan Furmanbff01f32020-09-15 15:56:26 -0700836 data_types.append(candidate or base)
837 break
Ethan Furmanc2667362020-12-07 00:17:31 -0800838 else:
Ethan Furmanbff01f32020-09-15 15:56:26 -0700839 candidate = base
840 if len(data_types) > 1:
Ethan Furman3064dbf2020-09-16 07:11:57 -0700841 raise TypeError('%r: too many data types: %r' % (class_name, data_types))
Ethan Furmanbff01f32020-09-15 15:56:26 -0700842 elif data_types:
843 return data_types[0]
844 else:
845 return None
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700846
Ethan Furman5bdab642018-09-21 19:03:09 -0700847 # ensure final parent class is an Enum derivative, find any concrete
848 # data type, and check that Enum has no members
849 first_enum = bases[-1]
850 if not issubclass(first_enum, Enum):
851 raise TypeError("new enumerations should be created as "
852 "`EnumName([mixin_type, ...] [data_type,] enum_type)`")
853 member_type = _find_data_type(bases) or object
854 if first_enum._member_names_:
855 raise TypeError("Cannot extend enumerations")
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700856 return member_type, first_enum
857
858 @staticmethod
859 def _find_new_(classdict, member_type, first_enum):
Ethan Furman6d3dfee2020-12-08 12:26:56 -0800860 """
861 Returns the __new__ to be used for creating the enum members.
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700862
863 classdict: the class dictionary given to __new__
864 member_type: the data type whose __new__ will be used by default
865 first_enum: enumeration to check for an overriding __new__
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700866 """
867 # now find the correct __new__, checking to see of one was defined
868 # by the user; also check earlier enum classes in case a __new__ was
869 # saved as __new_member__
870 __new__ = classdict.get('__new__', None)
871
872 # should __new__ be saved as __new_member__ later?
Ethan Furmana02cb472021-04-21 10:20:44 -0700873 save_new = first_enum is not None and __new__ is not None
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700874
875 if __new__ is None:
876 # check all possibles for __new_member__ before falling back to
877 # __new__
878 for method in ('__new_member__', '__new__'):
879 for possible in (member_type, first_enum):
880 target = getattr(possible, method, None)
881 if target not in {
882 None,
883 None.__new__,
884 object.__new__,
885 Enum.__new__,
886 }:
887 __new__ = target
888 break
889 if __new__ is not None:
890 break
891 else:
892 __new__ = object.__new__
893
894 # if a non-object.__new__ is used then whatever value/tuple was
895 # assigned to the enum member name will be passed to __new__ and to the
896 # new enum member's __init__
Ethan Furmana02cb472021-04-21 10:20:44 -0700897 if first_enum is None or __new__ in (Enum.__new__, object.__new__):
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700898 use_args = False
899 else:
900 use_args = True
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700901 return __new__, save_new, use_args
Ethan Furmanb7751062021-03-30 21:17:26 -0700902EnumMeta = EnumType
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700903
904
Ethan Furmanb7751062021-03-30 21:17:26 -0700905class Enum(metaclass=EnumType):
Ethan Furman6d3dfee2020-12-08 12:26:56 -0800906 """
907 Generic enumeration.
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700908
909 Derive from this class to define new enumerations.
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700910 """
Ethan Furman7aaeb2a2021-01-25 14:26:19 -0800911
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700912 def __new__(cls, value):
913 # all enum instances are actually created during class construction
914 # without calling this method; this method is called by the metaclass'
915 # __call__ (i.e. Color(3) ), and by pickle
916 if type(value) is cls:
Ethan Furman23bb6f42016-11-21 09:22:05 -0800917 # For lookups like Color(Color.RED)
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700918 return value
919 # by-value search for a matching enum member
920 # see if it's in the reverse mapping (for hashable values)
Ethan Furman2aa27322013-07-19 19:35:56 -0700921 try:
Andrew Svetlov34ae04f2018-12-26 20:45:33 +0200922 return cls._value2member_map_[value]
923 except KeyError:
924 # Not found, no need to do long O(n) search
925 pass
Ethan Furman2aa27322013-07-19 19:35:56 -0700926 except TypeError:
927 # not there, now do long search -- O(n) behavior
Ethan Furman520ad572013-07-19 19:47:21 -0700928 for member in cls._member_map_.values():
Ethan Furman0081f232014-09-16 17:31:23 -0700929 if member._value_ == value:
Ethan Furman2aa27322013-07-19 19:35:56 -0700930 return member
Ethan Furmanee47e5c2016-08-31 00:12:15 -0700931 # still not found -- try _missing_ hook
Ethan Furman019f0a02018-09-12 11:43:34 -0700932 try:
933 exc = None
934 result = cls._missing_(value)
935 except Exception as e:
936 exc = e
937 result = None
Ethan Furman8c14f5a2021-04-12 08:51:20 -0700938 try:
939 if isinstance(result, cls):
940 return result
941 elif (
942 Flag is not None and issubclass(cls, Flag)
943 and cls._boundary_ is EJECT and isinstance(result, int)
944 ):
945 return result
946 else:
947 ve_exc = ValueError("%r is not a valid %s" % (value, cls.__qualname__))
948 if result is None and exc is None:
949 raise ve_exc
950 elif exc is None:
951 exc = TypeError(
952 'error in %s._missing_: returned %r instead of None or a valid member'
953 % (cls.__name__, result)
954 )
955 if not isinstance(exc, ValueError):
956 exc.__context__ = ve_exc
957 raise exc
958 finally:
959 # ensure all variables that could hold an exception are destroyed
960 exc = None
961 ve_exc = None
Ethan Furmanee47e5c2016-08-31 00:12:15 -0700962
Ethan Furmanc16595e2016-09-10 23:36:59 -0700963 def _generate_next_value_(name, start, count, last_values):
Ethan Furman6d3dfee2020-12-08 12:26:56 -0800964 """
965 Generate the next value when not given.
966
967 name: the name of the member
968 start: the initial start value or None
969 count: the number of existing members
970 last_value: the last value assigned or None
971 """
Ethan Furmanc16595e2016-09-10 23:36:59 -0700972 for last_value in reversed(last_values):
973 try:
974 return last_value + 1
975 except TypeError:
976 pass
977 else:
Ethan Furmanee47e5c2016-08-31 00:12:15 -0700978 return start
Ethan Furmanc16595e2016-09-10 23:36:59 -0700979
Ethan Furmanee47e5c2016-08-31 00:12:15 -0700980 @classmethod
981 def _missing_(cls, value):
Ethan Furmanc95ad7a2020-09-16 10:26:50 -0700982 return None
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700983
984 def __repr__(self):
Ethan Furmanb7751062021-03-30 21:17:26 -0700985 return "%s.%s" % ( self.__class__.__name__, self._name_)
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700986
987 def __str__(self):
Ethan Furmanb7751062021-03-30 21:17:26 -0700988 return "%s" % (self._name_, )
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700989
Ethan Furman388a3922013-08-12 06:51:41 -0700990 def __dir__(self):
Ethan Furman6d3dfee2020-12-08 12:26:56 -0800991 """
992 Returns all members and all public methods
993 """
Ethan Furman0ae550b2014-10-14 08:58:32 -0700994 added_behavior = [
995 m
996 for cls in self.__class__.mro()
997 for m in cls.__dict__
Ethan Furman354ecf12015-03-11 08:43:12 -0700998 if m[0] != '_' and m not in self._member_map_
Angelin BOOZ68526fe2020-09-21 15:11:06 +0200999 ] + [m for m in self.__dict__ if m[0] != '_']
Ethan Furmanec5f8eb2014-10-21 13:40:35 -07001000 return (['__class__', '__doc__', '__module__'] + added_behavior)
Ethan Furman388a3922013-08-12 06:51:41 -07001001
Ethan Furmanec15a822013-08-31 19:17:41 -07001002 def __format__(self, format_spec):
Ethan Furman6d3dfee2020-12-08 12:26:56 -08001003 """
1004 Returns format using actual value type unless __str__ has been overridden.
1005 """
Ethan Furmanec15a822013-08-31 19:17:41 -07001006 # mixed-in Enums should use the mixed-in type's __format__, otherwise
1007 # we can get strange results with the Enum name showing up instead of
1008 # the value
1009
thatneat2f19e822019-07-04 11:28:37 -07001010 # pure Enum branch, or branch with __str__ explicitly overridden
Ethan Furman37440ee2020-12-08 11:14:10 -08001011 str_overridden = type(self).__str__ not in (Enum.__str__, Flag.__str__)
thatneat2f19e822019-07-04 11:28:37 -07001012 if self._member_type_ is object or str_overridden:
Ethan Furmanec15a822013-08-31 19:17:41 -07001013 cls = str
1014 val = str(self)
1015 # mix-in branch
1016 else:
Ethan Furman6bd92882021-04-27 13:05:08 -07001017 if not format_spec or format_spec in ('{}','{:}'):
1018 import warnings
1019 warnings.warn(
1020 "in 3.12 format() will use the enum member, not the enum member's value;\n"
Pablo Galindo9a42d502021-05-01 20:26:09 +01001021 "use a format specifier, such as :d for an IntEnum member, to maintain "
Ethan Furman6bd92882021-04-27 13:05:08 -07001022 "the current display",
1023 DeprecationWarning,
1024 stacklevel=2,
1025 )
Ethan Furmanec15a822013-08-31 19:17:41 -07001026 cls = self._member_type_
Ethan Furman0081f232014-09-16 17:31:23 -07001027 val = self._value_
Ethan Furmanec15a822013-08-31 19:17:41 -07001028 return cls.__format__(val, format_spec)
1029
Ethan Furman6b3d64a2013-06-14 16:55:46 -07001030 def __hash__(self):
Ethan Furman520ad572013-07-19 19:47:21 -07001031 return hash(self._name_)
Ethan Furman6b3d64a2013-06-14 16:55:46 -07001032
Ethan Furmanca1b7942014-02-08 11:36:27 -08001033 def __reduce_ex__(self, proto):
Ethan Furmandc870522014-02-18 12:37:12 -08001034 return self.__class__, (self._value_, )
Ethan Furmanca1b7942014-02-08 11:36:27 -08001035
Ethan Furmanc314e602021-01-12 23:47:57 -08001036 # enum.property is used to provide access to the `name` and
1037 # `value` attributes of enum members while keeping some measure of
Ethan Furman6b3d64a2013-06-14 16:55:46 -07001038 # protection from modification, while still allowing for an enumeration
1039 # to have members named `name` and `value`. This works because enumeration
Ethan Furmanc314e602021-01-12 23:47:57 -08001040 # members are not set directly on the enum class; they are kept in a
1041 # separate structure, _member_map_, which is where enum.property looks for
1042 # them
Ethan Furman6b3d64a2013-06-14 16:55:46 -07001043
Ethan Furmanc314e602021-01-12 23:47:57 -08001044 @property
Ethan Furman6b3d64a2013-06-14 16:55:46 -07001045 def name(self):
Ethan Furmanc850f342013-09-15 16:59:35 -07001046 """The name of the Enum member."""
Ethan Furman520ad572013-07-19 19:47:21 -07001047 return self._name_
Ethan Furman6b3d64a2013-06-14 16:55:46 -07001048
Ethan Furmanc314e602021-01-12 23:47:57 -08001049 @property
Ethan Furman6b3d64a2013-06-14 16:55:46 -07001050 def value(self):
Ethan Furmanc850f342013-09-15 16:59:35 -07001051 """The value of the Enum member."""
Ethan Furman520ad572013-07-19 19:47:21 -07001052 return self._value_
Ethan Furman6b3d64a2013-06-14 16:55:46 -07001053
1054
1055class IntEnum(int, Enum):
Ethan Furman0063ff42020-09-21 17:23:13 -07001056 """
1057 Enum where members are also (and must be) ints
1058 """
1059
1060
1061class StrEnum(str, Enum):
1062 """
1063 Enum where members are also (and must be) strings
1064 """
1065
1066 def __new__(cls, *values):
1067 if len(values) > 3:
1068 raise TypeError('too many arguments for str(): %r' % (values, ))
1069 if len(values) == 1:
1070 # it must be a string
1071 if not isinstance(values[0], str):
1072 raise TypeError('%r is not a string' % (values[0], ))
Ethan Furman7aaeb2a2021-01-25 14:26:19 -08001073 if len(values) >= 2:
Ethan Furman0063ff42020-09-21 17:23:13 -07001074 # check that encoding argument is a string
1075 if not isinstance(values[1], str):
1076 raise TypeError('encoding must be a string, not %r' % (values[1], ))
Ethan Furman7aaeb2a2021-01-25 14:26:19 -08001077 if len(values) == 3:
1078 # check that errors argument is a string
1079 if not isinstance(values[2], str):
1080 raise TypeError('errors must be a string, not %r' % (values[2]))
Ethan Furman0063ff42020-09-21 17:23:13 -07001081 value = str(*values)
1082 member = str.__new__(cls, value)
1083 member._value_ = value
1084 return member
Ethan Furmanf24bb352013-07-18 17:05:39 -07001085
Ethan Furmand986d162020-09-22 13:00:07 -07001086 __str__ = str.__str__
1087
Ethan Furmanefb13be2020-12-10 12:20:06 -08001088 def _generate_next_value_(name, start, count, last_values):
1089 """
1090 Return the lower-cased version of the member name.
1091 """
1092 return name.lower()
1093
Ethan Furmanf24bb352013-07-18 17:05:39 -07001094
Ethan Furman24e837f2015-03-18 17:27:57 -07001095def _reduce_ex_by_name(self, proto):
1096 return self.name
1097
Ethan Furman7aaeb2a2021-01-25 14:26:19 -08001098class FlagBoundary(StrEnum):
1099 """
1100 control how out of range values are handled
1101 "strict" -> error is raised [default for Flag]
1102 "conform" -> extra bits are discarded
1103 "eject" -> lose flag status [default for IntFlag]
1104 "keep" -> keep flag status and all bits
1105 """
1106 STRICT = auto()
1107 CONFORM = auto()
1108 EJECT = auto()
1109 KEEP = auto()
1110STRICT, CONFORM, EJECT, KEEP = FlagBoundary
1111
1112
1113class Flag(Enum, boundary=STRICT):
Ethan Furman6d3dfee2020-12-08 12:26:56 -08001114 """
1115 Support for flags
1116 """
Ethan Furmanc16595e2016-09-10 23:36:59 -07001117
1118 def _generate_next_value_(name, start, count, last_values):
Ethan Furmanee47e5c2016-08-31 00:12:15 -07001119 """
1120 Generate the next value when not given.
1121
1122 name: the name of the member
HongWeipengbb16fb22019-09-21 13:22:54 +08001123 start: the initial start value or None
Ethan Furmanee47e5c2016-08-31 00:12:15 -07001124 count: the number of existing members
1125 last_value: the last value assigned or None
1126 """
1127 if not count:
1128 return start if start is not None else 1
Ethan Furman7aaeb2a2021-01-25 14:26:19 -08001129 last_value = max(last_values)
1130 try:
1131 high_bit = _high_bit(last_value)
1132 except Exception:
1133 raise TypeError('Invalid Flag value: %r' % last_value) from None
Ethan Furmanee47e5c2016-08-31 00:12:15 -07001134 return 2 ** (high_bit+1)
1135
1136 @classmethod
Ethan Furman7aaeb2a2021-01-25 14:26:19 -08001137 def _iter_member_by_value_(cls, value):
1138 """
1139 Extract all members from the value in definition (i.e. increasing value) order.
1140 """
1141 for val in _iter_bits_lsb(value & cls._flag_mask_):
1142 yield cls._value2member_map_.get(val)
1143
1144 _iter_member_ = _iter_member_by_value_
1145
1146 @classmethod
1147 def _iter_member_by_def_(cls, value):
1148 """
1149 Extract all members from the value in definition order.
1150 """
1151 yield from sorted(
1152 cls._iter_member_by_value_(value),
1153 key=lambda m: m._sort_order_,
1154 )
1155
1156 @classmethod
Ethan Furmanee47e5c2016-08-31 00:12:15 -07001157 def _missing_(cls, value):
Ethan Furman6d3dfee2020-12-08 12:26:56 -08001158 """
Ethan Furman0dca5eb2021-04-15 06:49:54 -07001159 Create a composite member containing all canonical members present in `value`.
1160
1161 If non-member values are present, result depends on `_boundary_` setting.
Ethan Furman3515dcc2016-09-18 13:15:41 -07001162 """
Ethan Furman7aaeb2a2021-01-25 14:26:19 -08001163 if not isinstance(value, int):
1164 raise ValueError(
1165 "%r is not a valid %s" % (value, cls.__qualname__)
1166 )
1167 # check boundaries
1168 # - value must be in range (e.g. -16 <-> +15, i.e. ~15 <-> 15)
1169 # - value must not include any skipped flags (e.g. if bit 2 is not
1170 # defined, then 0d10 is invalid)
1171 flag_mask = cls._flag_mask_
1172 all_bits = cls._all_bits_
1173 neg_value = None
1174 if (
1175 not ~all_bits <= value <= all_bits
1176 or value & (all_bits ^ flag_mask)
1177 ):
1178 if cls._boundary_ is STRICT:
1179 max_bits = max(value.bit_length(), flag_mask.bit_length())
1180 raise ValueError(
1181 "%s: invalid value: %r\n given %s\n allowed %s" % (
1182 cls.__name__, value, bin(value, max_bits), bin(flag_mask, max_bits),
1183 ))
1184 elif cls._boundary_ is CONFORM:
1185 value = value & flag_mask
1186 elif cls._boundary_ is EJECT:
1187 return value
1188 elif cls._boundary_ is KEEP:
1189 if value < 0:
1190 value = (
1191 max(all_bits+1, 2**(value.bit_length()))
1192 + value
1193 )
1194 else:
1195 raise ValueError(
1196 'unknown flag boundary: %r' % (cls._boundary_, )
1197 )
1198 if value < 0:
1199 neg_value = value
1200 value = all_bits + 1 + value
1201 # get members and unknown
1202 unknown = value & ~flag_mask
1203 member_value = value & flag_mask
1204 if unknown and cls._boundary_ is not KEEP:
1205 raise ValueError(
1206 '%s(%r) --> unknown values %r [%s]'
1207 % (cls.__name__, value, unknown, bin(unknown))
1208 )
1209 # normal Flag?
1210 __new__ = getattr(cls, '__new_member__', None)
1211 if cls._member_type_ is object and not __new__:
Ethan Furman3515dcc2016-09-18 13:15:41 -07001212 # construct a singleton enum pseudo-member
Ethan Furmanee47e5c2016-08-31 00:12:15 -07001213 pseudo_member = object.__new__(cls)
Ethan Furman7aaeb2a2021-01-25 14:26:19 -08001214 else:
1215 pseudo_member = (__new__ or cls._member_type_.__new__)(cls, value)
Ethan Furmana02cb472021-04-21 10:20:44 -07001216 if not hasattr(pseudo_member, '_value_'):
Ethan Furmanee47e5c2016-08-31 00:12:15 -07001217 pseudo_member._value_ = value
Ethan Furman7aaeb2a2021-01-25 14:26:19 -08001218 if member_value:
1219 pseudo_member._name_ = '|'.join([
1220 m._name_ for m in cls._iter_member_(member_value)
1221 ])
1222 if unknown:
1223 pseudo_member._name_ += '|0x%x' % unknown
1224 else:
1225 pseudo_member._name_ = None
1226 # use setdefault in case another thread already created a composite
1227 # with this value, but only if all members are known
1228 # note: zero is a special case -- add it
1229 if not unknown:
Ethan Furman28cf6632017-01-24 12:12:06 -08001230 pseudo_member = cls._value2member_map_.setdefault(value, pseudo_member)
Ethan Furman7aaeb2a2021-01-25 14:26:19 -08001231 if neg_value is not None:
1232 cls._value2member_map_[neg_value] = pseudo_member
Ethan Furmanee47e5c2016-08-31 00:12:15 -07001233 return pseudo_member
1234
Ethan Furmanee47e5c2016-08-31 00:12:15 -07001235 def __contains__(self, other):
Ethan Furman6d3dfee2020-12-08 12:26:56 -08001236 """
1237 Returns True if self has at least the same flags set as other.
1238 """
Ethan Furmanee47e5c2016-08-31 00:12:15 -07001239 if not isinstance(other, self.__class__):
Rahul Jha94306522018-09-10 23:51:04 +05301240 raise TypeError(
1241 "unsupported operand type(s) for 'in': '%s' and '%s'" % (
1242 type(other).__qualname__, self.__class__.__qualname__))
Ethan Furman7aaeb2a2021-01-25 14:26:19 -08001243 if other._value_ == 0 or self._value_ == 0:
1244 return False
Ethan Furmanee47e5c2016-08-31 00:12:15 -07001245 return other._value_ & self._value_ == other._value_
1246
Ethan Furman7219e272020-09-16 13:01:00 -07001247 def __iter__(self):
Ethan Furman6d3dfee2020-12-08 12:26:56 -08001248 """
Ethan Furman7aaeb2a2021-01-25 14:26:19 -08001249 Returns flags in definition order.
Ethan Furman6d3dfee2020-12-08 12:26:56 -08001250 """
Ethan Furman7aaeb2a2021-01-25 14:26:19 -08001251 yield from self._iter_member_(self._value_)
1252
1253 def __len__(self):
1254 return self._value_.bit_count()
Ethan Furman7219e272020-09-16 13:01:00 -07001255
Ethan Furmanee47e5c2016-08-31 00:12:15 -07001256 def __repr__(self):
Ethan Furmanb7751062021-03-30 21:17:26 -07001257 cls_name = self.__class__.__name__
1258 if self._name_ is None:
1259 return "0x%x" % (self._value_, )
1260 if _is_single_bit(self._value_):
1261 return '%s.%s' % (cls_name, self._name_)
1262 if self._boundary_ is not FlagBoundary.KEEP:
1263 return '%s.' % cls_name + ('|%s.' % cls_name).join(self.name.split('|'))
Ethan Furman7aaeb2a2021-01-25 14:26:19 -08001264 else:
Ethan Furmanb7751062021-03-30 21:17:26 -07001265 name = []
1266 for n in self._name_.split('|'):
1267 if n.startswith('0'):
1268 name.append(n)
1269 else:
1270 name.append('%s.%s' % (cls_name, n))
1271 return '|'.join(name)
Ethan Furmanee47e5c2016-08-31 00:12:15 -07001272
1273 def __str__(self):
1274 cls = self.__class__
Ethan Furmanb7751062021-03-30 21:17:26 -07001275 if self._name_ is None:
1276 return '%s(%x)' % (cls.__name__, self._value_)
Ethan Furmanee47e5c2016-08-31 00:12:15 -07001277 else:
Ethan Furmanb7751062021-03-30 21:17:26 -07001278 return self._name_
Ethan Furmanee47e5c2016-08-31 00:12:15 -07001279
Ethan Furman25d94bb2016-09-02 16:32:32 -07001280 def __bool__(self):
1281 return bool(self._value_)
1282
Ethan Furmanee47e5c2016-08-31 00:12:15 -07001283 def __or__(self, other):
1284 if not isinstance(other, self.__class__):
1285 return NotImplemented
1286 return self.__class__(self._value_ | other._value_)
1287
1288 def __and__(self, other):
1289 if not isinstance(other, self.__class__):
1290 return NotImplemented
1291 return self.__class__(self._value_ & other._value_)
1292
1293 def __xor__(self, other):
1294 if not isinstance(other, self.__class__):
1295 return NotImplemented
1296 return self.__class__(self._value_ ^ other._value_)
1297
1298 def __invert__(self):
Ethan Furman7aaeb2a2021-01-25 14:26:19 -08001299 if self._inverted_ is None:
1300 if self._boundary_ is KEEP:
1301 # use all bits
1302 self._inverted_ = self.__class__(~self._value_)
1303 else:
1304 # calculate flags not in this member
1305 self._inverted_ = self.__class__(self._flag_mask_ ^ self._value_)
1306 self._inverted_._inverted_ = self
1307 return self._inverted_
Ethan Furmanee47e5c2016-08-31 00:12:15 -07001308
1309
Ethan Furman7aaeb2a2021-01-25 14:26:19 -08001310class IntFlag(int, Flag, boundary=EJECT):
Ethan Furman6d3dfee2020-12-08 12:26:56 -08001311 """
1312 Support for integer-based Flags
1313 """
Ethan Furmanee47e5c2016-08-31 00:12:15 -07001314
Ethan Furmanee47e5c2016-08-31 00:12:15 -07001315 def __or__(self, other):
Ethan Furman7aaeb2a2021-01-25 14:26:19 -08001316 if isinstance(other, self.__class__):
1317 other = other._value_
1318 elif isinstance(other, int):
1319 other = other
1320 else:
Ethan Furmanee47e5c2016-08-31 00:12:15 -07001321 return NotImplemented
Ethan Furman7aaeb2a2021-01-25 14:26:19 -08001322 value = self._value_
1323 return self.__class__(value | other)
Ethan Furmanee47e5c2016-08-31 00:12:15 -07001324
1325 def __and__(self, other):
Ethan Furman7aaeb2a2021-01-25 14:26:19 -08001326 if isinstance(other, self.__class__):
1327 other = other._value_
1328 elif isinstance(other, int):
1329 other = other
1330 else:
Ethan Furmanee47e5c2016-08-31 00:12:15 -07001331 return NotImplemented
Ethan Furman7aaeb2a2021-01-25 14:26:19 -08001332 value = self._value_
1333 return self.__class__(value & other)
Ethan Furmanee47e5c2016-08-31 00:12:15 -07001334
1335 def __xor__(self, other):
Ethan Furman7aaeb2a2021-01-25 14:26:19 -08001336 if isinstance(other, self.__class__):
1337 other = other._value_
1338 elif isinstance(other, int):
1339 other = other
1340 else:
Ethan Furmanee47e5c2016-08-31 00:12:15 -07001341 return NotImplemented
Ethan Furman7aaeb2a2021-01-25 14:26:19 -08001342 value = self._value_
1343 return self.__class__(value ^ other)
Ethan Furmanee47e5c2016-08-31 00:12:15 -07001344
1345 __ror__ = __or__
1346 __rand__ = __and__
1347 __rxor__ = __xor__
Ethan Furman7aaeb2a2021-01-25 14:26:19 -08001348 __invert__ = Flag.__invert__
Ethan Furmanee47e5c2016-08-31 00:12:15 -07001349
1350def _high_bit(value):
Ethan Furman6d3dfee2020-12-08 12:26:56 -08001351 """
1352 returns index of highest bit, or -1 if value is zero or negative
1353 """
Ethan Furman3515dcc2016-09-18 13:15:41 -07001354 return value.bit_length() - 1
Ethan Furmanee47e5c2016-08-31 00:12:15 -07001355
Ethan Furmanf24bb352013-07-18 17:05:39 -07001356def unique(enumeration):
Ethan Furman6d3dfee2020-12-08 12:26:56 -08001357 """
1358 Class decorator for enumerations ensuring unique member values.
1359 """
Ethan Furmanf24bb352013-07-18 17:05:39 -07001360 duplicates = []
1361 for name, member in enumeration.__members__.items():
1362 if name != member.name:
1363 duplicates.append((name, member.name))
1364 if duplicates:
1365 alias_details = ', '.join(
1366 ["%s -> %s" % (alias, name) for (alias, name) in duplicates])
1367 raise ValueError('duplicate values found in %r: %s' %
1368 (enumeration, alias_details))
1369 return enumeration
Ethan Furman3515dcc2016-09-18 13:15:41 -07001370
Ethan Furman7aaeb2a2021-01-25 14:26:19 -08001371def _power_of_two(value):
1372 if value < 1:
1373 return False
1374 return value == 2 ** _high_bit(value)
Ethan Furmanb7751062021-03-30 21:17:26 -07001375
1376def global_enum_repr(self):
1377 return '%s.%s' % (self.__class__.__module__, self._name_)
1378
1379def global_flag_repr(self):
1380 module = self.__class__.__module__
1381 cls_name = self.__class__.__name__
1382 if self._name_ is None:
1383 return "%x" % (module, cls_name, self._value_)
1384 if _is_single_bit(self):
1385 return '%s.%s' % (module, self._name_)
1386 if self._boundary_ is not FlagBoundary.KEEP:
1387 return module + module.join(self.name.split('|'))
1388 else:
1389 name = []
1390 for n in self._name_.split('|'):
1391 if n.startswith('0'):
1392 name.append(n)
1393 else:
1394 name.append('%s.%s' % (module, n))
1395 return '|'.join(name)
1396
1397
1398def global_enum(cls):
1399 """
1400 decorator that makes the repr() of an enum member reference its module
1401 instead of its class; also exports all members to the enum's module's
1402 global namespace
1403 """
1404 if issubclass(cls, Flag):
1405 cls.__repr__ = global_flag_repr
1406 else:
1407 cls.__repr__ = global_enum_repr
1408 sys.modules[cls.__module__].__dict__.update(cls.__members__)
1409 return cls
Ethan Furmana02cb472021-04-21 10:20:44 -07001410
1411def _simple_enum(etype=Enum, *, boundary=None, use_args=None):
1412 """
1413 Class decorator that converts a normal class into an :class:`Enum`. No
1414 safety checks are done, and some advanced behavior (such as
1415 :func:`__init_subclass__`) is not available. Enum creation can be faster
1416 using :func:`simple_enum`.
1417
1418 >>> from enum import Enum, _simple_enum
1419 >>> @_simple_enum(Enum)
1420 ... class Color:
1421 ... RED = auto()
1422 ... GREEN = auto()
1423 ... BLUE = auto()
1424 >>> Color
1425 <enum 'Color'>
1426 """
1427 def convert_class(cls):
1428 nonlocal use_args
1429 cls_name = cls.__name__
1430 if use_args is None:
1431 use_args = etype._use_args_
1432 __new__ = cls.__dict__.get('__new__')
1433 if __new__ is not None:
1434 new_member = __new__.__func__
1435 else:
1436 new_member = etype._member_type_.__new__
1437 attrs = {}
1438 body = {}
1439 if __new__ is not None:
1440 body['__new_member__'] = new_member
1441 body['_new_member_'] = new_member
1442 body['_use_args_'] = use_args
1443 body['_generate_next_value_'] = gnv = etype._generate_next_value_
1444 body['_member_names_'] = member_names = []
1445 body['_member_map_'] = member_map = {}
1446 body['_value2member_map_'] = value2member_map = {}
Ethan Furman6bd92882021-04-27 13:05:08 -07001447 body['_unhashable_values_'] = []
Ethan Furmana02cb472021-04-21 10:20:44 -07001448 body['_member_type_'] = member_type = etype._member_type_
1449 if issubclass(etype, Flag):
1450 body['_boundary_'] = boundary or etype._boundary_
1451 body['_flag_mask_'] = None
1452 body['_all_bits_'] = None
1453 body['_inverted_'] = None
1454 for name, obj in cls.__dict__.items():
1455 if name in ('__dict__', '__weakref__'):
1456 continue
1457 if _is_dunder(name) or _is_private(cls_name, name) or _is_sunder(name) or _is_descriptor(obj):
1458 body[name] = obj
1459 else:
1460 attrs[name] = obj
1461 if cls.__dict__.get('__doc__') is None:
1462 body['__doc__'] = 'An enumeration.'
1463 #
1464 # double check that repr and friends are not the mixin's or various
1465 # things break (such as pickle)
1466 # however, if the method is defined in the Enum itself, don't replace
1467 # it
1468 enum_class = type(cls_name, (etype, ), body, boundary=boundary, _simple=True)
1469 for name in ('__repr__', '__str__', '__format__', '__reduce_ex__'):
1470 if name in body:
1471 continue
1472 class_method = getattr(enum_class, name)
1473 obj_method = getattr(member_type, name, None)
1474 enum_method = getattr(etype, name, None)
1475 if obj_method is not None and obj_method is class_method:
1476 setattr(enum_class, name, enum_method)
1477 gnv_last_values = []
1478 if issubclass(enum_class, Flag):
1479 # Flag / IntFlag
1480 single_bits = multi_bits = 0
1481 for name, value in attrs.items():
1482 if isinstance(value, auto) and auto.value is _auto_null:
1483 value = gnv(name, 1, len(member_names), gnv_last_values)
1484 if value in value2member_map:
1485 # an alias to an existing member
1486 redirect = property()
1487 redirect.__set_name__(enum_class, name)
1488 setattr(enum_class, name, redirect)
1489 member_map[name] = value2member_map[value]
1490 else:
1491 # create the member
1492 if use_args:
1493 if not isinstance(value, tuple):
1494 value = (value, )
1495 member = new_member(enum_class, *value)
1496 value = value[0]
1497 else:
1498 member = new_member(enum_class)
1499 if __new__ is None:
1500 member._value_ = value
1501 member._name_ = name
1502 member.__objclass__ = enum_class
1503 member.__init__(value)
1504 redirect = property()
1505 redirect.__set_name__(enum_class, name)
1506 setattr(enum_class, name, redirect)
1507 member_map[name] = member
1508 member._sort_order_ = len(member_names)
1509 value2member_map[value] = member
1510 if _is_single_bit(value):
1511 # not a multi-bit alias, record in _member_names_ and _flag_mask_
1512 member_names.append(name)
1513 single_bits |= value
1514 else:
1515 multi_bits |= value
1516 gnv_last_values.append(value)
1517 enum_class._flag_mask_ = single_bits
1518 enum_class._all_bits_ = 2 ** ((single_bits|multi_bits).bit_length()) - 1
1519 # set correct __iter__
1520 member_list = [m._value_ for m in enum_class]
1521 if member_list != sorted(member_list):
1522 enum_class._iter_member_ = enum_class._iter_member_by_def_
1523 else:
1524 # Enum / IntEnum / StrEnum
1525 for name, value in attrs.items():
1526 if isinstance(value, auto):
1527 if value.value is _auto_null:
1528 value.value = gnv(name, 1, len(member_names), gnv_last_values)
1529 value = value.value
1530 if value in value2member_map:
1531 # an alias to an existing member
1532 redirect = property()
1533 redirect.__set_name__(enum_class, name)
1534 setattr(enum_class, name, redirect)
1535 member_map[name] = value2member_map[value]
1536 else:
1537 # create the member
1538 if use_args:
1539 if not isinstance(value, tuple):
1540 value = (value, )
1541 member = new_member(enum_class, *value)
1542 value = value[0]
1543 else:
1544 member = new_member(enum_class)
1545 if __new__ is None:
1546 member._value_ = value
1547 member._name_ = name
1548 member.__objclass__ = enum_class
1549 member.__init__(value)
1550 member._sort_order_ = len(member_names)
1551 redirect = property()
1552 redirect.__set_name__(enum_class, name)
1553 setattr(enum_class, name, redirect)
1554 member_map[name] = member
1555 value2member_map[value] = member
1556 member_names.append(name)
1557 gnv_last_values.append(value)
1558 if '__new__' in body:
1559 enum_class.__new_member__ = enum_class.__new__
1560 enum_class.__new__ = Enum.__new__
1561 return enum_class
1562 return convert_class
1563
1564def _test_simple_enum(checked_enum, simple_enum):
1565 """
1566 A function that can be used to test an enum created with :func:`_simple_enum`
1567 against the version created by subclassing :class:`Enum`::
1568
1569 >>> from enum import Enum, _simple_enum, _test_simple_enum
1570 >>> @_simple_enum(Enum)
1571 ... class Color:
1572 ... RED = auto()
1573 ... GREEN = auto()
1574 ... BLUE = auto()
1575 >>> class CheckedColor(Enum):
1576 ... RED = auto()
1577 ... GREEN = auto()
1578 ... BLUE = auto()
1579 >>> _test_simple_enum(CheckedColor, Color)
1580
1581 If differences are found, a :exc:`TypeError` is raised.
1582 """
1583 failed = []
1584 if checked_enum.__dict__ != simple_enum.__dict__:
1585 checked_dict = checked_enum.__dict__
1586 checked_keys = list(checked_dict.keys())
1587 simple_dict = simple_enum.__dict__
1588 simple_keys = list(simple_dict.keys())
1589 member_names = set(
1590 list(checked_enum._member_map_.keys())
1591 + list(simple_enum._member_map_.keys())
1592 )
1593 for key in set(checked_keys + simple_keys):
1594 if key in ('__module__', '_member_map_', '_value2member_map_'):
1595 # keys known to be different
1596 continue
1597 elif key in member_names:
1598 # members are checked below
1599 continue
1600 elif key not in simple_keys:
1601 failed.append("missing key: %r" % (key, ))
1602 elif key not in checked_keys:
1603 failed.append("extra key: %r" % (key, ))
1604 else:
1605 checked_value = checked_dict[key]
1606 simple_value = simple_dict[key]
1607 if callable(checked_value):
1608 continue
1609 if key == '__doc__':
1610 # remove all spaces/tabs
1611 compressed_checked_value = checked_value.replace(' ','').replace('\t','')
1612 compressed_simple_value = simple_value.replace(' ','').replace('\t','')
1613 if compressed_checked_value != compressed_simple_value:
1614 failed.append("%r:\n %s\n %s" % (
1615 key,
1616 "checked -> %r" % (checked_value, ),
1617 "simple -> %r" % (simple_value, ),
1618 ))
1619 elif checked_value != simple_value:
1620 failed.append("%r:\n %s\n %s" % (
1621 key,
1622 "checked -> %r" % (checked_value, ),
1623 "simple -> %r" % (simple_value, ),
1624 ))
1625 failed.sort()
1626 for name in member_names:
1627 failed_member = []
1628 if name not in simple_keys:
1629 failed.append('missing member from simple enum: %r' % name)
1630 elif name not in checked_keys:
1631 failed.append('extra member in simple enum: %r' % name)
1632 else:
1633 checked_member_dict = checked_enum[name].__dict__
1634 checked_member_keys = list(checked_member_dict.keys())
1635 simple_member_dict = simple_enum[name].__dict__
1636 simple_member_keys = list(simple_member_dict.keys())
1637 for key in set(checked_member_keys + simple_member_keys):
Ethan Furman6c681e12021-04-23 19:08:22 -07001638 if key in ('__module__', '__objclass__', '_inverted_'):
1639 # keys known to be different or absent
Ethan Furmana02cb472021-04-21 10:20:44 -07001640 continue
1641 elif key not in simple_member_keys:
1642 failed_member.append("missing key %r not in the simple enum member %r" % (key, name))
1643 elif key not in checked_member_keys:
1644 failed_member.append("extra key %r in simple enum member %r" % (key, name))
1645 else:
1646 checked_value = checked_member_dict[key]
1647 simple_value = simple_member_dict[key]
1648 if checked_value != simple_value:
1649 failed_member.append("%r:\n %s\n %s" % (
1650 key,
1651 "checked member -> %r" % (checked_value, ),
1652 "simple member -> %r" % (simple_value, ),
1653 ))
1654 if failed_member:
1655 failed.append('%r member mismatch:\n %s' % (
1656 name, '\n '.join(failed_member),
1657 ))
1658 for method in (
1659 '__str__', '__repr__', '__reduce_ex__', '__format__',
1660 '__getnewargs_ex__', '__getnewargs__', '__reduce_ex__', '__reduce__'
1661 ):
1662 if method in simple_keys and method in checked_keys:
1663 # cannot compare functions, and it exists in both, so we're good
1664 continue
1665 elif method not in simple_keys and method not in checked_keys:
1666 # method is inherited -- check it out
1667 checked_method = getattr(checked_enum, method, None)
1668 simple_method = getattr(simple_enum, method, None)
1669 if hasattr(checked_method, '__func__'):
1670 checked_method = checked_method.__func__
1671 simple_method = simple_method.__func__
1672 if checked_method != simple_method:
1673 failed.append("%r: %-30s %s" % (
1674 method,
1675 "checked -> %r" % (checked_method, ),
1676 "simple -> %r" % (simple_method, ),
1677 ))
1678 else:
1679 # if the method existed in only one of the enums, it will have been caught
1680 # in the first checks above
1681 pass
1682 if failed:
1683 raise TypeError('enum mismatch:\n %s' % '\n '.join(failed))
1684
1685def _old_convert_(etype, name, module, filter, source=None, *, boundary=None):
1686 """
1687 Create a new Enum subclass that replaces a collection of global constants
1688 """
1689 # convert all constants from source (or module) that pass filter() to
1690 # a new Enum called name, and export the enum and its members back to
1691 # module;
1692 # also, replace the __reduce_ex__ method so unpickling works in
1693 # previous Python versions
1694 module_globals = sys.modules[module].__dict__
1695 if source:
1696 source = source.__dict__
1697 else:
1698 source = module_globals
1699 # _value2member_map_ is populated in the same order every time
1700 # for a consistent reverse mapping of number to name when there
1701 # are multiple names for the same number.
1702 members = [
1703 (name, value)
1704 for name, value in source.items()
1705 if filter(name)]
1706 try:
1707 # sort by value
1708 members.sort(key=lambda t: (t[1], t[0]))
1709 except TypeError:
1710 # unless some values aren't comparable, in which case sort by name
1711 members.sort(key=lambda t: t[0])
1712 cls = etype(name, members, module=module, boundary=boundary or KEEP)
1713 cls.__reduce_ex__ = _reduce_ex_by_name
1714 cls.__repr__ = global_enum_repr
1715 return cls