blob: ed0c9ce72d01c46e913c382b6fe7eed742ff01e8 [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 Furman6b3d64a2013-06-14 16:55:46 -07003
Ethan Furmane5754ab2015-09-17 22:03:52 -07004
Ethan Furmanc16595e2016-09-10 23:36:59 -07005__all__ = [
6 'EnumMeta',
Ethan Furman0063ff42020-09-21 17:23:13 -07007 'Enum', 'IntEnum', 'StrEnum', 'Flag', 'IntFlag',
Ethan Furmanc16595e2016-09-10 23:36:59 -07008 'auto', 'unique',
9 ]
Ethan Furman6b3d64a2013-06-14 16:55:46 -070010
11
Ethan Furman6bd94de2020-12-09 16:41:22 -080012class _NoInitSubclass:
13 """
14 temporary base class to suppress calling __init_subclass__
15 """
16 @classmethod
17 def __init_subclass__(cls, **kwds):
18 pass
19
Ethan Furman101e0742013-09-15 12:34:36 -070020def _is_descriptor(obj):
Ethan Furman6d3dfee2020-12-08 12:26:56 -080021 """
22 Returns True if obj is a descriptor, False otherwise.
23 """
Ethan Furman101e0742013-09-15 12:34:36 -070024 return (
25 hasattr(obj, '__get__') or
26 hasattr(obj, '__set__') or
Ethan Furman6d3dfee2020-12-08 12:26:56 -080027 hasattr(obj, '__delete__')
28 )
Ethan Furman101e0742013-09-15 12:34:36 -070029
Ethan Furman6b3d64a2013-06-14 16:55:46 -070030def _is_dunder(name):
Ethan Furman6d3dfee2020-12-08 12:26:56 -080031 """
32 Returns True if a __dunder__ name, False otherwise.
33 """
34 return (
35 len(name) > 4 and
Brennan D Baraban8b914d22019-03-03 14:09:11 -080036 name[:2] == name[-2:] == '__' and
37 name[2] != '_' and
Ethan Furman6d3dfee2020-12-08 12:26:56 -080038 name[-3] != '_'
39 )
Ethan Furman6b3d64a2013-06-14 16:55:46 -070040
41def _is_sunder(name):
Ethan Furman6d3dfee2020-12-08 12:26:56 -080042 """
43 Returns True if a _sunder_ name, False otherwise.
44 """
45 return (
46 len(name) > 2 and
Brennan D Baraban8b914d22019-03-03 14:09:11 -080047 name[0] == name[-1] == '_' and
Ethan Furman6b3d64a2013-06-14 16:55:46 -070048 name[1:2] != '_' and
Ethan Furman6d3dfee2020-12-08 12:26:56 -080049 name[-2:-1] != '_'
50 )
Ethan Furman6b3d64a2013-06-14 16:55:46 -070051
Ethan Furman7cf0aad2020-12-09 17:12:11 -080052def _is_private(cls_name, name):
53 # do not use `re` as `re` imports `enum`
54 pattern = '_%s__' % (cls_name, )
55 if (
56 len(name) >= 5
57 and name.startswith(pattern)
58 and name[len(pattern)] != '_'
59 and (name[-1] != '_' or name[-2] != '_')
60 ):
61 return True
62 else:
63 return False
64
Ethan Furman6b3d64a2013-06-14 16:55:46 -070065def _make_class_unpicklable(cls):
Ethan Furman6d3dfee2020-12-08 12:26:56 -080066 """
67 Make the given class un-picklable.
68 """
Ethan Furmanca1b7942014-02-08 11:36:27 -080069 def _break_on_call_reduce(self, proto):
Ethan Furman6b3d64a2013-06-14 16:55:46 -070070 raise TypeError('%r cannot be pickled' % self)
Ethan Furmanca1b7942014-02-08 11:36:27 -080071 cls.__reduce_ex__ = _break_on_call_reduce
Ethan Furman6b3d64a2013-06-14 16:55:46 -070072 cls.__module__ = '<unknown>'
73
Ethan Furman3515dcc2016-09-18 13:15:41 -070074_auto_null = object()
Ethan Furmanc16595e2016-09-10 23:36:59 -070075class auto:
76 """
77 Instances are replaced with an appropriate value in Enum class suites.
78 """
Ethan Furman3515dcc2016-09-18 13:15:41 -070079 value = _auto_null
Ethan Furmanc16595e2016-09-10 23:36:59 -070080
Ethan Furman101e0742013-09-15 12:34:36 -070081
Ethan Furman6b3d64a2013-06-14 16:55:46 -070082class _EnumDict(dict):
Ethan Furman6d3dfee2020-12-08 12:26:56 -080083 """
84 Track enum member order and ensure member names are not reused.
Ethan Furman6b3d64a2013-06-14 16:55:46 -070085
86 EnumMeta will use the names found in self._member_names as the
87 enumeration member names.
Ethan Furman6b3d64a2013-06-14 16:55:46 -070088 """
89 def __init__(self):
90 super().__init__()
91 self._member_names = []
Ethan Furmanc16595e2016-09-10 23:36:59 -070092 self._last_values = []
Ethan Furmana4b1bb42018-01-22 07:56:37 -080093 self._ignore = []
Ethan Onstottd9a43e22020-04-28 13:20:55 -040094 self._auto_called = False
Ethan Furman6b3d64a2013-06-14 16:55:46 -070095
96 def __setitem__(self, key, value):
Ethan Furman6d3dfee2020-12-08 12:26:56 -080097 """
98 Changes anything not dundered or not a descriptor.
Ethan Furman6b3d64a2013-06-14 16:55:46 -070099
100 If an enum member name is used twice, an error is raised; duplicate
101 values are not checked for.
102
103 Single underscore (sunder) names are reserved.
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700104 """
Ethan Furman7cf0aad2020-12-09 17:12:11 -0800105 if _is_private(self._cls_name, key):
106 # do nothing, name will be a normal attribute
107 pass
108 elif _is_sunder(key):
Ethan Furmanee47e5c2016-08-31 00:12:15 -0700109 if key not in (
Ethan Furman3515dcc2016-09-18 13:15:41 -0700110 '_order_', '_create_pseudo_member_',
Ethan Furmana4b1bb42018-01-22 07:56:37 -0800111 '_generate_next_value_', '_missing_', '_ignore_',
Ethan Furmanee47e5c2016-08-31 00:12:15 -0700112 ):
Ethan Furman6d3dfee2020-12-08 12:26:56 -0800113 raise ValueError(
114 '_sunder_ names, such as %r, are reserved for future Enum use'
115 % (key, )
116 )
Ethan Furmanc16595e2016-09-10 23:36:59 -0700117 if key == '_generate_next_value_':
Ethan Onstottd9a43e22020-04-28 13:20:55 -0400118 # check if members already defined as auto()
119 if self._auto_called:
120 raise TypeError("_generate_next_value_ must be defined before members")
Ethan Furmanc16595e2016-09-10 23:36:59 -0700121 setattr(self, '_generate_next_value', value)
Ethan Furmana4b1bb42018-01-22 07:56:37 -0800122 elif key == '_ignore_':
123 if isinstance(value, str):
124 value = value.replace(',',' ').split()
125 else:
126 value = list(value)
127 self._ignore = value
128 already = set(value) & set(self._member_names)
129 if already:
Ethan Furman6d3dfee2020-12-08 12:26:56 -0800130 raise ValueError(
131 '_ignore_ cannot specify already set names: %r'
132 % (already, )
133 )
Ethan Furman101e0742013-09-15 12:34:36 -0700134 elif _is_dunder(key):
Ethan Furmane8e61272016-08-20 07:19:31 -0700135 if key == '__order__':
136 key = '_order_'
Ethan Furman101e0742013-09-15 12:34:36 -0700137 elif key in self._member_names:
138 # descriptor overwriting an enum?
139 raise TypeError('Attempted to reuse key: %r' % key)
Ethan Furmana4b1bb42018-01-22 07:56:37 -0800140 elif key in self._ignore:
141 pass
Ethan Furman101e0742013-09-15 12:34:36 -0700142 elif not _is_descriptor(value):
143 if key in self:
144 # enum overwriting a descriptor?
Ethan Furmanee47e5c2016-08-31 00:12:15 -0700145 raise TypeError('%r already defined as: %r' % (key, self[key]))
Ethan Furmanc16595e2016-09-10 23:36:59 -0700146 if isinstance(value, auto):
Ethan Furman3515dcc2016-09-18 13:15:41 -0700147 if value.value == _auto_null:
Ethan Furman6d3dfee2020-12-08 12:26:56 -0800148 value.value = self._generate_next_value(
149 key,
150 1,
151 len(self._member_names),
152 self._last_values[:],
153 )
Ethan Furmanfc23a942020-09-16 12:37:54 -0700154 self._auto_called = True
Ethan Furman3515dcc2016-09-18 13:15:41 -0700155 value = value.value
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700156 self._member_names.append(key)
Ethan Furmanc16595e2016-09-10 23:36:59 -0700157 self._last_values.append(value)
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700158 super().__setitem__(key, value)
159
160
Ezio Melotti9a3777e2013-08-17 15:53:55 +0300161# Dummy value for Enum as EnumMeta explicitly checks for it, but of course
162# until EnumMeta finishes running the first time the Enum class doesn't exist.
163# This is also why there are checks in EnumMeta like `if Enum is not None`
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700164Enum = None
165
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700166class EnumMeta(type):
Ethan Furman6d3dfee2020-12-08 12:26:56 -0800167 """
168 Metaclass for Enum
169 """
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700170 @classmethod
Ethan Furman332dbc72016-08-20 00:00:52 -0700171 def __prepare__(metacls, cls, bases):
Ethan Furman3064dbf2020-09-16 07:11:57 -0700172 # check that previous enum members do not exist
173 metacls._check_for_existing_members(cls, bases)
Ethan Furmanee47e5c2016-08-31 00:12:15 -0700174 # create the namespace dict
175 enum_dict = _EnumDict()
Ethan Furman7cf0aad2020-12-09 17:12:11 -0800176 enum_dict._cls_name = cls
Ethan Furmanee47e5c2016-08-31 00:12:15 -0700177 # inherit previous flags and _generate_next_value_ function
Ethan Furman3064dbf2020-09-16 07:11:57 -0700178 member_type, first_enum = metacls._get_mixins_(cls, bases)
Ethan Furmanee47e5c2016-08-31 00:12:15 -0700179 if first_enum is not None:
Ethan Furman6d3dfee2020-12-08 12:26:56 -0800180 enum_dict['_generate_next_value_'] = getattr(
181 first_enum, '_generate_next_value_', None,
182 )
Ethan Furmanee47e5c2016-08-31 00:12:15 -0700183 return enum_dict
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700184
Ethan Furman6bd94de2020-12-09 16:41:22 -0800185 def __new__(metacls, cls, bases, classdict, **kwds):
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700186 # an Enum class is final once enumeration items have been defined; it
187 # cannot be mixed with other types (int, float, etc.) if it has an
188 # inherited __new__ unless a new __new__ is defined (or the resulting
189 # class will fail).
Ethan Furmana4b1bb42018-01-22 07:56:37 -0800190 #
191 # remove any keys listed in _ignore_
192 classdict.setdefault('_ignore_', []).append('_ignore_')
193 ignore = classdict['_ignore_']
194 for key in ignore:
195 classdict.pop(key, None)
Ethan Furman3064dbf2020-09-16 07:11:57 -0700196 member_type, first_enum = metacls._get_mixins_(cls, bases)
Ethan Furmanc2667362020-12-07 00:17:31 -0800197 __new__, save_new, use_args = metacls._find_new_(
198 classdict, member_type, first_enum,
199 )
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700200
201 # save enum items into separate mapping so they don't get baked into
202 # the new class
Ethan Furmanee47e5c2016-08-31 00:12:15 -0700203 enum_members = {k: classdict[k] for k in classdict._member_names}
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700204 for name in classdict._member_names:
205 del classdict[name]
206
Ethan Furmane8e61272016-08-20 07:19:31 -0700207 # adjust the sunders
208 _order_ = classdict.pop('_order_', None)
209
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700210 # check for illegal enum names (any others?)
Brennan D Baraban8b914d22019-03-03 14:09:11 -0800211 invalid_names = set(enum_members) & {'mro', ''}
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700212 if invalid_names:
213 raise ValueError('Invalid enum member name: {0}'.format(
214 ','.join(invalid_names)))
215
Ethan Furman48a724f2015-04-11 23:23:06 -0700216 # create a default docstring if one has not been provided
217 if '__doc__' not in classdict:
218 classdict['__doc__'] = 'An enumeration.'
219
Ethan Furman6bd94de2020-12-09 16:41:22 -0800220 # postpone calling __init_subclass__
221 if '__init_subclass__' in classdict and classdict['__init_subclass__'] is None:
222 raise TypeError('%s.__init_subclass__ cannot be None')
223 # remove current __init_subclass__ so previous one can be found with getattr
224 new_init_subclass = classdict.pop('__init_subclass__', None)
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700225 # create our new Enum type
Ethan Furman6bd94de2020-12-09 16:41:22 -0800226 if bases:
227 bases = (_NoInitSubclass, ) + bases
228 enum_class = type.__new__(metacls, cls, bases, classdict)
229 enum_class.__bases__ = enum_class.__bases__[1:] #or (object, )
230 else:
231 enum_class = type.__new__(metacls, cls, bases, classdict)
232 old_init_subclass = getattr(enum_class, '__init_subclass__', None)
233 # and restore the new one (if there was one)
234 if new_init_subclass is not None:
235 enum_class.__init_subclass__ = classmethod(new_init_subclass)
Ethan Furman520ad572013-07-19 19:47:21 -0700236 enum_class._member_names_ = [] # names in definition order
INADA Naokie57f91a2018-06-19 01:14:26 +0900237 enum_class._member_map_ = {} # name->value map
Ethan Furman5e5a8232013-08-04 08:42:23 -0700238 enum_class._member_type_ = member_type
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700239
orlnub1230fb9fad2018-09-12 20:28:53 +0300240 # save DynamicClassAttribute attributes from super classes so we know
241 # if we can take the shortcut of storing members in the class dict
Ethan Furman6d3dfee2020-12-08 12:26:56 -0800242 dynamic_attributes = {
243 k for c in enum_class.mro()
244 for k, v in c.__dict__.items()
245 if isinstance(v, DynamicClassAttribute)
246 }
Ethan Furman354ecf12015-03-11 08:43:12 -0700247
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700248 # Reverse value->name map for hashable values.
Ethan Furman520ad572013-07-19 19:47:21 -0700249 enum_class._value2member_map_ = {}
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700250
Ethan Furman2da95042014-03-03 12:42:52 -0800251 # If a custom type is mixed into the Enum, and it does not know how
252 # to pickle itself, pickle.dumps will succeed but pickle.loads will
253 # fail. Rather than have the error show up later and possibly far
254 # from the source, sabotage the pickle protocol for this class so
255 # that pickle.dumps also fails.
256 #
257 # However, if the new class implements its own __reduce_ex__, do not
258 # sabotage -- it's on them to make sure it works correctly. We use
259 # __reduce_ex__ instead of any of the others as it is preferred by
260 # pickle over __reduce__, and it handles all pickle protocols.
261 if '__reduce_ex__' not in classdict:
Ethan Furmandc870522014-02-18 12:37:12 -0800262 if member_type is not object:
263 methods = ('__getnewargs_ex__', '__getnewargs__',
264 '__reduce_ex__', '__reduce__')
Ethan Furman2da95042014-03-03 12:42:52 -0800265 if not any(m in member_type.__dict__ for m in methods):
Ethan Furmandc870522014-02-18 12:37:12 -0800266 _make_class_unpicklable(enum_class)
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700267
268 # instantiate them, checking for duplicates as we go
269 # we instantiate first instead of checking for duplicates first in case
270 # a custom __new__ is doing something funky with the values -- such as
271 # auto-numbering ;)
272 for member_name in classdict._member_names:
Ethan Furmanee47e5c2016-08-31 00:12:15 -0700273 value = enum_members[member_name]
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700274 if not isinstance(value, tuple):
275 args = (value, )
276 else:
277 args = value
278 if member_type is tuple: # special case for tuple enums
279 args = (args, ) # wrap it one more time
280 if not use_args:
281 enum_member = __new__(enum_class)
Ethan Furmanb41803e2013-07-25 13:50:45 -0700282 if not hasattr(enum_member, '_value_'):
283 enum_member._value_ = value
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700284 else:
285 enum_member = __new__(enum_class, *args)
Ethan Furmanb41803e2013-07-25 13:50:45 -0700286 if not hasattr(enum_member, '_value_'):
Ethan Furmanee47e5c2016-08-31 00:12:15 -0700287 if member_type is object:
288 enum_member._value_ = value
289 else:
290 enum_member._value_ = member_type(*args)
Ethan Furman520ad572013-07-19 19:47:21 -0700291 value = enum_member._value_
Ethan Furman520ad572013-07-19 19:47:21 -0700292 enum_member._name_ = member_name
Ethan Furmanc850f342013-09-15 16:59:35 -0700293 enum_member.__objclass__ = enum_class
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700294 enum_member.__init__(*args)
295 # If another member with the same value was already defined, the
296 # new member becomes an alias to the existing one.
Ethan Furman520ad572013-07-19 19:47:21 -0700297 for name, canonical_member in enum_class._member_map_.items():
Ethan Furman0081f232014-09-16 17:31:23 -0700298 if canonical_member._value_ == enum_member._value_:
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700299 enum_member = canonical_member
300 break
301 else:
302 # Aliases don't appear in member names (only in __members__).
Ethan Furman520ad572013-07-19 19:47:21 -0700303 enum_class._member_names_.append(member_name)
Ethan Furman354ecf12015-03-11 08:43:12 -0700304 # performance boost for any member that would not shadow
305 # a DynamicClassAttribute
orlnub1230fb9fad2018-09-12 20:28:53 +0300306 if member_name not in dynamic_attributes:
Ethan Furman354ecf12015-03-11 08:43:12 -0700307 setattr(enum_class, member_name, enum_member)
308 # now add to _member_map_
Ethan Furman520ad572013-07-19 19:47:21 -0700309 enum_class._member_map_[member_name] = enum_member
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700310 try:
311 # This may fail if value is not hashable. We can't add the value
312 # to the map, and by-value lookups for this value will be
313 # linear.
Ethan Furman520ad572013-07-19 19:47:21 -0700314 enum_class._value2member_map_[value] = enum_member
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700315 except TypeError:
316 pass
317
318 # double check that repr and friends are not the mixin's or various
319 # things break (such as pickle)
Ethan Furman22415ad2020-09-15 16:28:25 -0700320 # however, if the method is defined in the Enum itself, don't replace
321 # it
Ethan Furmandc870522014-02-18 12:37:12 -0800322 for name in ('__repr__', '__str__', '__format__', '__reduce_ex__'):
Ethan Furman22415ad2020-09-15 16:28:25 -0700323 if name in classdict:
324 continue
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700325 class_method = getattr(enum_class, name)
326 obj_method = getattr(member_type, name, None)
327 enum_method = getattr(first_enum, name, None)
328 if obj_method is not None and obj_method is class_method:
329 setattr(enum_class, name, enum_method)
330
331 # replace any other __new__ with our own (as long as Enum is not None,
332 # anyway) -- again, this is to support pickle
333 if Enum is not None:
334 # if the user defined their own __new__, save it before it gets
335 # clobbered in case they subclass later
336 if save_new:
337 enum_class.__new_member__ = __new__
338 enum_class.__new__ = Enum.__new__
Ethan Furmane8e61272016-08-20 07:19:31 -0700339
340 # py3 support for definition order (helps keep py2/py3 code in sync)
341 if _order_ is not None:
342 if isinstance(_order_, str):
343 _order_ = _order_.replace(',', ' ').split()
344 if _order_ != enum_class._member_names_:
345 raise TypeError('member order does not match _order_')
346
Ethan Furman6bd94de2020-12-09 16:41:22 -0800347 # finally, call parents' __init_subclass__
348 if Enum is not None and old_init_subclass is not None:
349 old_init_subclass(**kwds)
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700350 return enum_class
351
Ethan Furman5de67b12016-04-13 23:52:09 -0700352 def __bool__(self):
353 """
354 classes/types should always be True.
355 """
356 return True
357
Ethan Furmand9925a12014-09-16 20:35:55 -0700358 def __call__(cls, value, names=None, *, module=None, qualname=None, type=None, start=1):
Ethan Furman6d3dfee2020-12-08 12:26:56 -0800359 """
360 Either returns an existing member, or creates a new enum class.
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700361
362 This method is used both when an enum class is given a value to match
363 to an enumeration member (i.e. Color(3)) and for the functional API
Ethan Furman23bb6f42016-11-21 09:22:05 -0800364 (i.e. Color = Enum('Color', names='RED GREEN BLUE')).
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700365
Ethan Furman2da95042014-03-03 12:42:52 -0800366 When used for the functional API:
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700367
Ethan Furman2da95042014-03-03 12:42:52 -0800368 `value` will be the name of the new class.
369
370 `names` should be either a string of white-space/comma delimited names
Ethan Furmand9925a12014-09-16 20:35:55 -0700371 (values will start at `start`), or an iterator/mapping of name, value pairs.
Ethan Furman2da95042014-03-03 12:42:52 -0800372
373 `module` should be set to the module this class is being created in;
374 if it is not set, an attempt to find that module will be made, but if
375 it fails the class will not be picklable.
376
377 `qualname` should be set to the actual location this class can be found
378 at in its module; by default it is set to the global scope. If this is
379 not correct, unpickling will fail in some circumstances.
380
381 `type`, if set, will be mixed in as the first base class.
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700382 """
383 if names is None: # simple value lookup
384 return cls.__new__(cls, value)
385 # otherwise, functional API: we're creating a new Enum type
Ethan Furman6d3dfee2020-12-08 12:26:56 -0800386 return cls._create_(
387 value,
388 names,
389 module=module,
390 qualname=qualname,
391 type=type,
392 start=start,
393 )
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700394
395 def __contains__(cls, member):
Rahul Jha94306522018-09-10 23:51:04 +0530396 if not isinstance(member, Enum):
397 raise TypeError(
398 "unsupported operand type(s) for 'in': '%s' and '%s'" % (
399 type(member).__qualname__, cls.__class__.__qualname__))
Ethan Furman0081f232014-09-16 17:31:23 -0700400 return isinstance(member, cls) and member._name_ in cls._member_map_
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700401
Ethan Furman64a99722013-09-22 16:18:19 -0700402 def __delattr__(cls, attr):
403 # nicer error message when someone tries to delete an attribute
404 # (see issue19025).
405 if attr in cls._member_map_:
Ethan Furman6d3dfee2020-12-08 12:26:56 -0800406 raise AttributeError("%s: cannot delete Enum member %r." % (cls.__name__, attr))
Ethan Furman64a99722013-09-22 16:18:19 -0700407 super().__delattr__(attr)
408
Ethan Furman388a3922013-08-12 06:51:41 -0700409 def __dir__(self):
Ethan Furman6d3dfee2020-12-08 12:26:56 -0800410 return (
411 ['__class__', '__doc__', '__members__', '__module__']
412 + self._member_names_
413 )
Ethan Furman388a3922013-08-12 06:51:41 -0700414
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700415 def __getattr__(cls, name):
Ethan Furman6d3dfee2020-12-08 12:26:56 -0800416 """
417 Return the enum member matching `name`
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700418
419 We use __getattr__ instead of descriptors or inserting into the enum
420 class' __dict__ in order to support `name` and `value` being both
421 properties for enum members (which live in the class' __dict__) and
422 enum members themselves.
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700423 """
424 if _is_dunder(name):
425 raise AttributeError(name)
426 try:
Ethan Furman520ad572013-07-19 19:47:21 -0700427 return cls._member_map_[name]
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700428 except KeyError:
429 raise AttributeError(name) from None
430
431 def __getitem__(cls, name):
Ethan Furman520ad572013-07-19 19:47:21 -0700432 return cls._member_map_[name]
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700433
434 def __iter__(cls):
Ethan Furman6d3dfee2020-12-08 12:26:56 -0800435 """
436 Returns members in definition order.
437 """
Ethan Furman520ad572013-07-19 19:47:21 -0700438 return (cls._member_map_[name] for name in cls._member_names_)
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700439
440 def __len__(cls):
Ethan Furman520ad572013-07-19 19:47:21 -0700441 return len(cls._member_names_)
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700442
Ethan Furman2131a4a2013-09-14 18:11:24 -0700443 @property
444 def __members__(cls):
Ethan Furman6d3dfee2020-12-08 12:26:56 -0800445 """
446 Returns a mapping of member name->value.
Ethan Furman2131a4a2013-09-14 18:11:24 -0700447
448 This mapping lists all enum members, including aliases. Note that this
449 is a read-only view of the internal mapping.
Ethan Furman2131a4a2013-09-14 18:11:24 -0700450 """
451 return MappingProxyType(cls._member_map_)
452
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700453 def __repr__(cls):
454 return "<enum %r>" % cls.__name__
455
Ethan Furman2131a4a2013-09-14 18:11:24 -0700456 def __reversed__(cls):
Ethan Furman6d3dfee2020-12-08 12:26:56 -0800457 """
458 Returns members in reverse definition order.
459 """
Ethan Furman2131a4a2013-09-14 18:11:24 -0700460 return (cls._member_map_[name] for name in reversed(cls._member_names_))
461
Ethan Furmanf203f2d2013-09-06 07:16:48 -0700462 def __setattr__(cls, name, value):
Ethan Furman6d3dfee2020-12-08 12:26:56 -0800463 """
464 Block attempts to reassign Enum members.
Ethan Furmanf203f2d2013-09-06 07:16:48 -0700465
466 A simple assignment to the class namespace only changes one of the
467 several possible ways to get an Enum member from the Enum class,
468 resulting in an inconsistent Enumeration.
Ethan Furmanf203f2d2013-09-06 07:16:48 -0700469 """
470 member_map = cls.__dict__.get('_member_map_', {})
471 if name in member_map:
472 raise AttributeError('Cannot reassign members.')
473 super().__setattr__(name, value)
474
anentropicb8e21f12018-04-16 04:40:35 +0100475 def _create_(cls, class_name, names, *, module=None, qualname=None, type=None, start=1):
Ethan Furman6d3dfee2020-12-08 12:26:56 -0800476 """
477 Convenience method to create a new Enum class.
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700478
479 `names` can be:
480
481 * A string containing member names, separated either with spaces or
Ethan Furmand9925a12014-09-16 20:35:55 -0700482 commas. Values are incremented by 1 from `start`.
483 * An iterable of member names. Values are incremented by 1 from `start`.
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700484 * An iterable of (member name, value) pairs.
Ethan Furmand9925a12014-09-16 20:35:55 -0700485 * A mapping of member name -> value pairs.
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700486 """
487 metacls = cls.__class__
488 bases = (cls, ) if type is None else (type, cls)
Ethan Furman3064dbf2020-09-16 07:11:57 -0700489 _, first_enum = cls._get_mixins_(cls, bases)
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700490 classdict = metacls.__prepare__(class_name, bases)
491
492 # special processing needed for names?
493 if isinstance(names, str):
494 names = names.replace(',', ' ').split()
Dong-hee Nadcc8ce42017-06-22 01:52:32 +0900495 if isinstance(names, (tuple, list)) and names and isinstance(names[0], str):
Ethan Furmanee47e5c2016-08-31 00:12:15 -0700496 original_names, names = names, []
Ethan Furmanc16595e2016-09-10 23:36:59 -0700497 last_values = []
Ethan Furmanee47e5c2016-08-31 00:12:15 -0700498 for count, name in enumerate(original_names):
Ethan Furmanc16595e2016-09-10 23:36:59 -0700499 value = first_enum._generate_next_value_(name, start, count, last_values[:])
500 last_values.append(value)
501 names.append((name, value))
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700502
503 # Here, names is either an iterable of (name, value) or a mapping.
504 for item in names:
505 if isinstance(item, str):
506 member_name, member_value = item, names[item]
507 else:
508 member_name, member_value = item
509 classdict[member_name] = member_value
510 enum_class = metacls.__new__(metacls, class_name, bases, classdict)
511
512 # TODO: replace the frame hack if a blessed way to know the calling
513 # module is ever developed
514 if module is None:
515 try:
516 module = sys._getframe(2).f_globals['__name__']
Pablo Galindo293dd232019-11-19 21:34:03 +0000517 except (AttributeError, ValueError, KeyError):
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700518 pass
519 if module is None:
520 _make_class_unpicklable(enum_class)
521 else:
522 enum_class.__module__ = module
Ethan Furmanca1b7942014-02-08 11:36:27 -0800523 if qualname is not None:
524 enum_class.__qualname__ = qualname
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700525
526 return enum_class
527
orlnub1230fb9fad2018-09-12 20:28:53 +0300528 def _convert_(cls, name, module, filter, source=None):
529 """
530 Create a new Enum subclass that replaces a collection of global constants
531 """
532 # convert all constants from source (or module) that pass filter() to
533 # a new Enum called name, and export the enum and its members back to
534 # module;
535 # also, replace the __reduce_ex__ method so unpickling works in
536 # previous Python versions
537 module_globals = vars(sys.modules[module])
538 if source:
539 source = vars(source)
540 else:
541 source = module_globals
542 # _value2member_map_ is populated in the same order every time
543 # for a consistent reverse mapping of number to name when there
544 # are multiple names for the same number.
545 members = [
546 (name, value)
547 for name, value in source.items()
548 if filter(name)]
549 try:
550 # sort by value
551 members.sort(key=lambda t: (t[1], t[0]))
552 except TypeError:
553 # unless some values aren't comparable, in which case sort by name
554 members.sort(key=lambda t: t[0])
555 cls = cls(name, members, module=module)
556 cls.__reduce_ex__ = _reduce_ex_by_name
557 module_globals.update(cls.__members__)
558 module_globals[name] = cls
559 return cls
560
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700561 @staticmethod
Ethan Furman3064dbf2020-09-16 07:11:57 -0700562 def _check_for_existing_members(class_name, bases):
563 for chain in bases:
564 for base in chain.__mro__:
565 if issubclass(base, Enum) and base._member_names_:
Ethan Furman6d3dfee2020-12-08 12:26:56 -0800566 raise TypeError(
567 "%s: cannot extend enumeration %r"
568 % (class_name, base.__name__)
569 )
Ethan Furman3064dbf2020-09-16 07:11:57 -0700570
571 @staticmethod
572 def _get_mixins_(class_name, bases):
Ethan Furman6d3dfee2020-12-08 12:26:56 -0800573 """
574 Returns the type for creating enum members, and the first inherited
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700575 enum class.
576
577 bases: the tuple of bases that was given to __new__
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700578 """
579 if not bases:
580 return object, Enum
581
Ethan Furman5bdab642018-09-21 19:03:09 -0700582 def _find_data_type(bases):
Ethan Furmanbff01f32020-09-15 15:56:26 -0700583 data_types = []
Ethan Furman5bdab642018-09-21 19:03:09 -0700584 for chain in bases:
Ethan Furmanbff01f32020-09-15 15:56:26 -0700585 candidate = None
Ethan Furman5bdab642018-09-21 19:03:09 -0700586 for base in chain.__mro__:
587 if base is object:
588 continue
Ethan Furmanc2667362020-12-07 00:17:31 -0800589 elif issubclass(base, Enum):
590 if base._member_type_ is not object:
591 data_types.append(base._member_type_)
592 break
Ethan Furman5bdab642018-09-21 19:03:09 -0700593 elif '__new__' in base.__dict__:
Ethan Furmancd453852018-10-05 23:29:36 -0700594 if issubclass(base, Enum):
Ethan Furman5bdab642018-09-21 19:03:09 -0700595 continue
Ethan Furmanbff01f32020-09-15 15:56:26 -0700596 data_types.append(candidate or base)
597 break
Ethan Furmanc2667362020-12-07 00:17:31 -0800598 else:
Ethan Furmanbff01f32020-09-15 15:56:26 -0700599 candidate = base
600 if len(data_types) > 1:
Ethan Furman3064dbf2020-09-16 07:11:57 -0700601 raise TypeError('%r: too many data types: %r' % (class_name, data_types))
Ethan Furmanbff01f32020-09-15 15:56:26 -0700602 elif data_types:
603 return data_types[0]
604 else:
605 return None
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700606
Ethan Furman5bdab642018-09-21 19:03:09 -0700607 # ensure final parent class is an Enum derivative, find any concrete
608 # data type, and check that Enum has no members
609 first_enum = bases[-1]
610 if not issubclass(first_enum, Enum):
611 raise TypeError("new enumerations should be created as "
612 "`EnumName([mixin_type, ...] [data_type,] enum_type)`")
613 member_type = _find_data_type(bases) or object
614 if first_enum._member_names_:
615 raise TypeError("Cannot extend enumerations")
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700616 return member_type, first_enum
617
618 @staticmethod
619 def _find_new_(classdict, member_type, first_enum):
Ethan Furman6d3dfee2020-12-08 12:26:56 -0800620 """
621 Returns the __new__ to be used for creating the enum members.
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700622
623 classdict: the class dictionary given to __new__
624 member_type: the data type whose __new__ will be used by default
625 first_enum: enumeration to check for an overriding __new__
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700626 """
627 # now find the correct __new__, checking to see of one was defined
628 # by the user; also check earlier enum classes in case a __new__ was
629 # saved as __new_member__
630 __new__ = classdict.get('__new__', None)
631
632 # should __new__ be saved as __new_member__ later?
633 save_new = __new__ is not None
634
635 if __new__ is None:
636 # check all possibles for __new_member__ before falling back to
637 # __new__
638 for method in ('__new_member__', '__new__'):
639 for possible in (member_type, first_enum):
640 target = getattr(possible, method, None)
641 if target not in {
642 None,
643 None.__new__,
644 object.__new__,
645 Enum.__new__,
646 }:
647 __new__ = target
648 break
649 if __new__ is not None:
650 break
651 else:
652 __new__ = object.__new__
653
654 # if a non-object.__new__ is used then whatever value/tuple was
655 # assigned to the enum member name will be passed to __new__ and to the
656 # new enum member's __init__
657 if __new__ is object.__new__:
658 use_args = False
659 else:
660 use_args = True
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700661 return __new__, save_new, use_args
662
663
664class Enum(metaclass=EnumMeta):
Ethan Furman6d3dfee2020-12-08 12:26:56 -0800665 """
666 Generic enumeration.
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700667
668 Derive from this class to define new enumerations.
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700669 """
670 def __new__(cls, value):
671 # all enum instances are actually created during class construction
672 # without calling this method; this method is called by the metaclass'
673 # __call__ (i.e. Color(3) ), and by pickle
674 if type(value) is cls:
Ethan Furman23bb6f42016-11-21 09:22:05 -0800675 # For lookups like Color(Color.RED)
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700676 return value
677 # by-value search for a matching enum member
678 # see if it's in the reverse mapping (for hashable values)
Ethan Furman2aa27322013-07-19 19:35:56 -0700679 try:
Andrew Svetlov34ae04f2018-12-26 20:45:33 +0200680 return cls._value2member_map_[value]
681 except KeyError:
682 # Not found, no need to do long O(n) search
683 pass
Ethan Furman2aa27322013-07-19 19:35:56 -0700684 except TypeError:
685 # not there, now do long search -- O(n) behavior
Ethan Furman520ad572013-07-19 19:47:21 -0700686 for member in cls._member_map_.values():
Ethan Furman0081f232014-09-16 17:31:23 -0700687 if member._value_ == value:
Ethan Furman2aa27322013-07-19 19:35:56 -0700688 return member
Ethan Furmanee47e5c2016-08-31 00:12:15 -0700689 # still not found -- try _missing_ hook
Ethan Furman019f0a02018-09-12 11:43:34 -0700690 try:
691 exc = None
692 result = cls._missing_(value)
693 except Exception as e:
694 exc = e
695 result = None
696 if isinstance(result, cls):
697 return result
698 else:
Walter Dörwald323842c2019-07-18 20:37:13 +0200699 ve_exc = ValueError("%r is not a valid %s" % (value, cls.__qualname__))
Ethan Furman019f0a02018-09-12 11:43:34 -0700700 if result is None and exc is None:
701 raise ve_exc
702 elif exc is None:
703 exc = TypeError(
704 'error in %s._missing_: returned %r instead of None or a valid member'
705 % (cls.__name__, result)
706 )
707 exc.__context__ = ve_exc
708 raise exc
Ethan Furmanee47e5c2016-08-31 00:12:15 -0700709
Ethan Furmanc16595e2016-09-10 23:36:59 -0700710 def _generate_next_value_(name, start, count, last_values):
Ethan Furman6d3dfee2020-12-08 12:26:56 -0800711 """
712 Generate the next value when not given.
713
714 name: the name of the member
715 start: the initial start value or None
716 count: the number of existing members
717 last_value: the last value assigned or None
718 """
Ethan Furmanc16595e2016-09-10 23:36:59 -0700719 for last_value in reversed(last_values):
720 try:
721 return last_value + 1
722 except TypeError:
723 pass
724 else:
Ethan Furmanee47e5c2016-08-31 00:12:15 -0700725 return start
Ethan Furmanc16595e2016-09-10 23:36:59 -0700726
Ethan Furman6bd94de2020-12-09 16:41:22 -0800727 def __init_subclass__(cls, **kwds):
728 super().__init_subclass__(**kwds)
729
Ethan Furmanee47e5c2016-08-31 00:12:15 -0700730 @classmethod
731 def _missing_(cls, value):
Ethan Furmanc95ad7a2020-09-16 10:26:50 -0700732 return None
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700733
734 def __repr__(self):
735 return "<%s.%s: %r>" % (
Ethan Furman520ad572013-07-19 19:47:21 -0700736 self.__class__.__name__, self._name_, self._value_)
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700737
738 def __str__(self):
Ethan Furman520ad572013-07-19 19:47:21 -0700739 return "%s.%s" % (self.__class__.__name__, self._name_)
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700740
Ethan Furman388a3922013-08-12 06:51:41 -0700741 def __dir__(self):
Ethan Furman6d3dfee2020-12-08 12:26:56 -0800742 """
743 Returns all members and all public methods
744 """
Ethan Furman0ae550b2014-10-14 08:58:32 -0700745 added_behavior = [
746 m
747 for cls in self.__class__.mro()
748 for m in cls.__dict__
Ethan Furman354ecf12015-03-11 08:43:12 -0700749 if m[0] != '_' and m not in self._member_map_
Angelin BOOZ68526fe2020-09-21 15:11:06 +0200750 ] + [m for m in self.__dict__ if m[0] != '_']
Ethan Furmanec5f8eb2014-10-21 13:40:35 -0700751 return (['__class__', '__doc__', '__module__'] + added_behavior)
Ethan Furman388a3922013-08-12 06:51:41 -0700752
Ethan Furmanec15a822013-08-31 19:17:41 -0700753 def __format__(self, format_spec):
Ethan Furman6d3dfee2020-12-08 12:26:56 -0800754 """
755 Returns format using actual value type unless __str__ has been overridden.
756 """
Ethan Furmanec15a822013-08-31 19:17:41 -0700757 # mixed-in Enums should use the mixed-in type's __format__, otherwise
758 # we can get strange results with the Enum name showing up instead of
759 # the value
760
thatneat2f19e822019-07-04 11:28:37 -0700761 # pure Enum branch, or branch with __str__ explicitly overridden
Ethan Furman37440ee2020-12-08 11:14:10 -0800762 str_overridden = type(self).__str__ not in (Enum.__str__, Flag.__str__)
thatneat2f19e822019-07-04 11:28:37 -0700763 if self._member_type_ is object or str_overridden:
Ethan Furmanec15a822013-08-31 19:17:41 -0700764 cls = str
765 val = str(self)
766 # mix-in branch
767 else:
768 cls = self._member_type_
Ethan Furman0081f232014-09-16 17:31:23 -0700769 val = self._value_
Ethan Furmanec15a822013-08-31 19:17:41 -0700770 return cls.__format__(val, format_spec)
771
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700772 def __hash__(self):
Ethan Furman520ad572013-07-19 19:47:21 -0700773 return hash(self._name_)
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700774
Ethan Furmanca1b7942014-02-08 11:36:27 -0800775 def __reduce_ex__(self, proto):
Ethan Furmandc870522014-02-18 12:37:12 -0800776 return self.__class__, (self._value_, )
Ethan Furmanca1b7942014-02-08 11:36:27 -0800777
Ethan Furman33918c12013-09-27 23:02:02 -0700778 # DynamicClassAttribute is used to provide access to the `name` and
779 # `value` properties of enum members while keeping some measure of
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700780 # protection from modification, while still allowing for an enumeration
781 # to have members named `name` and `value`. This works because enumeration
782 # members are not set directly on the enum class -- __getattr__ is
783 # used to look them up.
784
Ethan Furmane03ea372013-09-25 07:14:41 -0700785 @DynamicClassAttribute
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700786 def name(self):
Ethan Furmanc850f342013-09-15 16:59:35 -0700787 """The name of the Enum member."""
Ethan Furman520ad572013-07-19 19:47:21 -0700788 return self._name_
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700789
Ethan Furmane03ea372013-09-25 07:14:41 -0700790 @DynamicClassAttribute
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700791 def value(self):
Ethan Furmanc850f342013-09-15 16:59:35 -0700792 """The value of the Enum member."""
Ethan Furman520ad572013-07-19 19:47:21 -0700793 return self._value_
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700794
795
796class IntEnum(int, Enum):
Ethan Furman0063ff42020-09-21 17:23:13 -0700797 """
798 Enum where members are also (and must be) ints
799 """
800
801
802class StrEnum(str, Enum):
803 """
804 Enum where members are also (and must be) strings
805 """
806
807 def __new__(cls, *values):
808 if len(values) > 3:
809 raise TypeError('too many arguments for str(): %r' % (values, ))
810 if len(values) == 1:
811 # it must be a string
812 if not isinstance(values[0], str):
813 raise TypeError('%r is not a string' % (values[0], ))
814 if len(values) > 1:
815 # check that encoding argument is a string
816 if not isinstance(values[1], str):
817 raise TypeError('encoding must be a string, not %r' % (values[1], ))
818 if len(values) > 2:
819 # check that errors argument is a string
820 if not isinstance(values[2], str):
821 raise TypeError('errors must be a string, not %r' % (values[2], ))
822 value = str(*values)
823 member = str.__new__(cls, value)
824 member._value_ = value
825 return member
Ethan Furmanf24bb352013-07-18 17:05:39 -0700826
Ethan Furmand986d162020-09-22 13:00:07 -0700827 __str__ = str.__str__
828
Ethan Furmanefb13be2020-12-10 12:20:06 -0800829 def _generate_next_value_(name, start, count, last_values):
830 """
831 Return the lower-cased version of the member name.
832 """
833 return name.lower()
834
Ethan Furmanf24bb352013-07-18 17:05:39 -0700835
Ethan Furman24e837f2015-03-18 17:27:57 -0700836def _reduce_ex_by_name(self, proto):
837 return self.name
838
Ethan Furman65a5a472016-09-01 23:55:19 -0700839class Flag(Enum):
Ethan Furman6d3dfee2020-12-08 12:26:56 -0800840 """
841 Support for flags
842 """
Ethan Furmanc16595e2016-09-10 23:36:59 -0700843
844 def _generate_next_value_(name, start, count, last_values):
Ethan Furmanee47e5c2016-08-31 00:12:15 -0700845 """
846 Generate the next value when not given.
847
848 name: the name of the member
HongWeipengbb16fb22019-09-21 13:22:54 +0800849 start: the initial start value or None
Ethan Furmanee47e5c2016-08-31 00:12:15 -0700850 count: the number of existing members
851 last_value: the last value assigned or None
852 """
853 if not count:
854 return start if start is not None else 1
Ethan Furmanc16595e2016-09-10 23:36:59 -0700855 for last_value in reversed(last_values):
856 try:
857 high_bit = _high_bit(last_value)
858 break
Ethan Furman3515dcc2016-09-18 13:15:41 -0700859 except Exception:
Ethan Furmanc16595e2016-09-10 23:36:59 -0700860 raise TypeError('Invalid Flag value: %r' % last_value) from None
Ethan Furmanee47e5c2016-08-31 00:12:15 -0700861 return 2 ** (high_bit+1)
862
863 @classmethod
864 def _missing_(cls, value):
Ethan Furman6d3dfee2020-12-08 12:26:56 -0800865 """
866 Returns member (possibly creating it) if one can be found for value.
867 """
Ethan Furmanee47e5c2016-08-31 00:12:15 -0700868 original_value = value
869 if value < 0:
870 value = ~value
871 possible_member = cls._create_pseudo_member_(value)
Ethan Furmanee47e5c2016-08-31 00:12:15 -0700872 if original_value < 0:
873 possible_member = ~possible_member
874 return possible_member
875
876 @classmethod
877 def _create_pseudo_member_(cls, value):
Ethan Furman3515dcc2016-09-18 13:15:41 -0700878 """
879 Create a composite member iff value contains only members.
880 """
Ethan Furmanee47e5c2016-08-31 00:12:15 -0700881 pseudo_member = cls._value2member_map_.get(value, None)
882 if pseudo_member is None:
Ethan Furman3515dcc2016-09-18 13:15:41 -0700883 # verify all bits are accounted for
884 _, extra_flags = _decompose(cls, value)
885 if extra_flags:
Walter Dörwald323842c2019-07-18 20:37:13 +0200886 raise ValueError("%r is not a valid %s" % (value, cls.__qualname__))
Ethan Furman3515dcc2016-09-18 13:15:41 -0700887 # construct a singleton enum pseudo-member
Ethan Furmanee47e5c2016-08-31 00:12:15 -0700888 pseudo_member = object.__new__(cls)
889 pseudo_member._name_ = None
890 pseudo_member._value_ = value
Ethan Furman28cf6632017-01-24 12:12:06 -0800891 # use setdefault in case another thread already created a composite
892 # with this value
893 pseudo_member = cls._value2member_map_.setdefault(value, pseudo_member)
Ethan Furmanee47e5c2016-08-31 00:12:15 -0700894 return pseudo_member
895
Ethan Furmanee47e5c2016-08-31 00:12:15 -0700896 def __contains__(self, other):
Ethan Furman6d3dfee2020-12-08 12:26:56 -0800897 """
898 Returns True if self has at least the same flags set as other.
899 """
Ethan Furmanee47e5c2016-08-31 00:12:15 -0700900 if not isinstance(other, self.__class__):
Rahul Jha94306522018-09-10 23:51:04 +0530901 raise TypeError(
902 "unsupported operand type(s) for 'in': '%s' and '%s'" % (
903 type(other).__qualname__, self.__class__.__qualname__))
Ethan Furmanee47e5c2016-08-31 00:12:15 -0700904 return other._value_ & self._value_ == other._value_
905
Ethan Furman7219e272020-09-16 13:01:00 -0700906 def __iter__(self):
Ethan Furman6d3dfee2020-12-08 12:26:56 -0800907 """
908 Returns flags in decreasing value order.
909 """
Ethan Furman7219e272020-09-16 13:01:00 -0700910 members, extra_flags = _decompose(self.__class__, self.value)
911 return (m for m in members if m._value_ != 0)
912
Ethan Furmanee47e5c2016-08-31 00:12:15 -0700913 def __repr__(self):
914 cls = self.__class__
915 if self._name_ is not None:
916 return '<%s.%s: %r>' % (cls.__name__, self._name_, self._value_)
Ethan Furman3515dcc2016-09-18 13:15:41 -0700917 members, uncovered = _decompose(cls, self._value_)
Ethan Furman27682d22016-09-04 11:39:01 -0700918 return '<%s.%s: %r>' % (
919 cls.__name__,
920 '|'.join([str(m._name_ or m._value_) for m in members]),
921 self._value_,
922 )
Ethan Furmanee47e5c2016-08-31 00:12:15 -0700923
924 def __str__(self):
925 cls = self.__class__
926 if self._name_ is not None:
927 return '%s.%s' % (cls.__name__, self._name_)
Ethan Furman3515dcc2016-09-18 13:15:41 -0700928 members, uncovered = _decompose(cls, self._value_)
Ethan Furmanee47e5c2016-08-31 00:12:15 -0700929 if len(members) == 1 and members[0]._name_ is None:
930 return '%s.%r' % (cls.__name__, members[0]._value_)
931 else:
932 return '%s.%s' % (
933 cls.__name__,
934 '|'.join([str(m._name_ or m._value_) for m in members]),
935 )
936
Ethan Furman25d94bb2016-09-02 16:32:32 -0700937 def __bool__(self):
938 return bool(self._value_)
939
Ethan Furmanee47e5c2016-08-31 00:12:15 -0700940 def __or__(self, other):
941 if not isinstance(other, self.__class__):
942 return NotImplemented
943 return self.__class__(self._value_ | other._value_)
944
945 def __and__(self, other):
946 if not isinstance(other, self.__class__):
947 return NotImplemented
948 return self.__class__(self._value_ & other._value_)
949
950 def __xor__(self, other):
951 if not isinstance(other, self.__class__):
952 return NotImplemented
953 return self.__class__(self._value_ ^ other._value_)
954
955 def __invert__(self):
Ethan Furman3515dcc2016-09-18 13:15:41 -0700956 members, uncovered = _decompose(self.__class__, self._value_)
Serhiy Storchaka81108372017-09-26 00:55:55 +0300957 inverted = self.__class__(0)
958 for m in self.__class__:
959 if m not in members and not (m._value_ & self._value_):
960 inverted = inverted | m
Ethan Furmanee47e5c2016-08-31 00:12:15 -0700961 return self.__class__(inverted)
962
963
Ethan Furman65a5a472016-09-01 23:55:19 -0700964class IntFlag(int, Flag):
Ethan Furman6d3dfee2020-12-08 12:26:56 -0800965 """
966 Support for integer-based Flags
967 """
Ethan Furmanee47e5c2016-08-31 00:12:15 -0700968
969 @classmethod
Ethan Furman3515dcc2016-09-18 13:15:41 -0700970 def _missing_(cls, value):
Ethan Furman6d3dfee2020-12-08 12:26:56 -0800971 """
972 Returns member (possibly creating it) if one can be found for value.
973 """
Ethan Furman3515dcc2016-09-18 13:15:41 -0700974 if not isinstance(value, int):
Walter Dörwald323842c2019-07-18 20:37:13 +0200975 raise ValueError("%r is not a valid %s" % (value, cls.__qualname__))
Ethan Furman3515dcc2016-09-18 13:15:41 -0700976 new_member = cls._create_pseudo_member_(value)
977 return new_member
978
979 @classmethod
Ethan Furmanee47e5c2016-08-31 00:12:15 -0700980 def _create_pseudo_member_(cls, value):
Ethan Furman6d3dfee2020-12-08 12:26:56 -0800981 """
982 Create a composite member iff value contains only members.
983 """
Ethan Furmanee47e5c2016-08-31 00:12:15 -0700984 pseudo_member = cls._value2member_map_.get(value, None)
985 if pseudo_member is None:
Ethan Furman3515dcc2016-09-18 13:15:41 -0700986 need_to_create = [value]
987 # get unaccounted for bits
988 _, extra_flags = _decompose(cls, value)
989 # timer = 10
990 while extra_flags:
991 # timer -= 1
992 bit = _high_bit(extra_flags)
993 flag_value = 2 ** bit
994 if (flag_value not in cls._value2member_map_ and
995 flag_value not in need_to_create
996 ):
997 need_to_create.append(flag_value)
998 if extra_flags == -flag_value:
999 extra_flags = 0
1000 else:
1001 extra_flags ^= flag_value
1002 for value in reversed(need_to_create):
1003 # construct singleton pseudo-members
1004 pseudo_member = int.__new__(cls, value)
1005 pseudo_member._name_ = None
1006 pseudo_member._value_ = value
Ethan Furman28cf6632017-01-24 12:12:06 -08001007 # use setdefault in case another thread already created a composite
1008 # with this value
1009 pseudo_member = cls._value2member_map_.setdefault(value, pseudo_member)
Ethan Furmanee47e5c2016-08-31 00:12:15 -07001010 return pseudo_member
1011
Ethan Furmanee47e5c2016-08-31 00:12:15 -07001012 def __or__(self, other):
1013 if not isinstance(other, (self.__class__, int)):
1014 return NotImplemented
Ethan Furman3515dcc2016-09-18 13:15:41 -07001015 result = self.__class__(self._value_ | self.__class__(other)._value_)
1016 return result
Ethan Furmanee47e5c2016-08-31 00:12:15 -07001017
1018 def __and__(self, other):
1019 if not isinstance(other, (self.__class__, int)):
1020 return NotImplemented
1021 return self.__class__(self._value_ & self.__class__(other)._value_)
1022
1023 def __xor__(self, other):
1024 if not isinstance(other, (self.__class__, int)):
1025 return NotImplemented
1026 return self.__class__(self._value_ ^ self.__class__(other)._value_)
1027
1028 __ror__ = __or__
1029 __rand__ = __and__
1030 __rxor__ = __xor__
1031
1032 def __invert__(self):
Ethan Furman3515dcc2016-09-18 13:15:41 -07001033 result = self.__class__(~self._value_)
1034 return result
Ethan Furmanee47e5c2016-08-31 00:12:15 -07001035
1036
1037def _high_bit(value):
Ethan Furman6d3dfee2020-12-08 12:26:56 -08001038 """
1039 returns index of highest bit, or -1 if value is zero or negative
1040 """
Ethan Furman3515dcc2016-09-18 13:15:41 -07001041 return value.bit_length() - 1
Ethan Furmanee47e5c2016-08-31 00:12:15 -07001042
Ethan Furmanf24bb352013-07-18 17:05:39 -07001043def unique(enumeration):
Ethan Furman6d3dfee2020-12-08 12:26:56 -08001044 """
1045 Class decorator for enumerations ensuring unique member values.
1046 """
Ethan Furmanf24bb352013-07-18 17:05:39 -07001047 duplicates = []
1048 for name, member in enumeration.__members__.items():
1049 if name != member.name:
1050 duplicates.append((name, member.name))
1051 if duplicates:
1052 alias_details = ', '.join(
1053 ["%s -> %s" % (alias, name) for (alias, name) in duplicates])
1054 raise ValueError('duplicate values found in %r: %s' %
1055 (enumeration, alias_details))
1056 return enumeration
Ethan Furman3515dcc2016-09-18 13:15:41 -07001057
1058def _decompose(flag, value):
Ethan Furman6d3dfee2020-12-08 12:26:56 -08001059 """
1060 Extract all members from the value.
1061 """
Ethan Furman3515dcc2016-09-18 13:15:41 -07001062 # _decompose is only called if the value is not named
1063 not_covered = value
1064 negative = value < 0
Ethan Furman3515dcc2016-09-18 13:15:41 -07001065 members = []
HongWeipeng0b41a922019-11-27 06:36:02 +08001066 for member in flag:
1067 member_value = member.value
Ethan Furman3515dcc2016-09-18 13:15:41 -07001068 if member_value and member_value & value == member_value:
1069 members.append(member)
1070 not_covered &= ~member_value
HongWeipeng0b41a922019-11-27 06:36:02 +08001071 if not negative:
1072 tmp = not_covered
1073 while tmp:
1074 flag_value = 2 ** _high_bit(tmp)
1075 if flag_value in flag._value2member_map_:
1076 members.append(flag._value2member_map_[flag_value])
1077 not_covered &= ~flag_value
1078 tmp &= ~flag_value
Ethan Furman3515dcc2016-09-18 13:15:41 -07001079 if not members and value in flag._value2member_map_:
1080 members.append(flag._value2member_map_[value])
1081 members.sort(key=lambda m: m._value_, reverse=True)
1082 if len(members) > 1 and members[0].value == value:
1083 # we have the breakdown, don't need the value member itself
1084 members.pop(0)
1085 return members, not_covered