blob: 75249bfdc1be17534008af0487bf47aabda5275e [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?
Ethan Furmana6582872020-12-10 13:07:00 -0800139 raise TypeError('%r already defined as: %r' % (key, self[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
Ethan Furmana6582872020-12-10 13:07:00 -0800160 def update(self, members, **more_members):
161 try:
162 for name in members.keys():
163 self[name] = members[name]
164 except AttributeError:
165 for name, value in members:
166 self[name] = value
167 for name, value in more_members.items():
168 self[name] = value
169
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700170
Ezio Melotti9a3777e2013-08-17 15:53:55 +0300171# Dummy value for Enum as EnumMeta explicitly checks for it, but of course
172# until EnumMeta finishes running the first time the Enum class doesn't exist.
173# This is also why there are checks in EnumMeta like `if Enum is not None`
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700174Enum = None
175
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700176class EnumMeta(type):
Ethan Furman6d3dfee2020-12-08 12:26:56 -0800177 """
178 Metaclass for Enum
179 """
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700180 @classmethod
Ethan Furman6ec0ade2020-12-24 10:05:02 -0800181 def __prepare__(metacls, cls, bases, **kwds):
Ethan Furman3064dbf2020-09-16 07:11:57 -0700182 # check that previous enum members do not exist
183 metacls._check_for_existing_members(cls, bases)
Ethan Furmanee47e5c2016-08-31 00:12:15 -0700184 # create the namespace dict
185 enum_dict = _EnumDict()
Ethan Furman7cf0aad2020-12-09 17:12:11 -0800186 enum_dict._cls_name = cls
Ethan Furmanee47e5c2016-08-31 00:12:15 -0700187 # inherit previous flags and _generate_next_value_ function
Ethan Furman3064dbf2020-09-16 07:11:57 -0700188 member_type, first_enum = metacls._get_mixins_(cls, bases)
Ethan Furmanee47e5c2016-08-31 00:12:15 -0700189 if first_enum is not None:
Ethan Furman6d3dfee2020-12-08 12:26:56 -0800190 enum_dict['_generate_next_value_'] = getattr(
191 first_enum, '_generate_next_value_', None,
192 )
Ethan Furmanee47e5c2016-08-31 00:12:15 -0700193 return enum_dict
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700194
Ethan Furman6bd94de2020-12-09 16:41:22 -0800195 def __new__(metacls, cls, bases, classdict, **kwds):
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700196 # an Enum class is final once enumeration items have been defined; it
197 # cannot be mixed with other types (int, float, etc.) if it has an
198 # inherited __new__ unless a new __new__ is defined (or the resulting
199 # class will fail).
Ethan Furmana4b1bb42018-01-22 07:56:37 -0800200 #
201 # remove any keys listed in _ignore_
202 classdict.setdefault('_ignore_', []).append('_ignore_')
203 ignore = classdict['_ignore_']
204 for key in ignore:
205 classdict.pop(key, None)
Ethan Furman3064dbf2020-09-16 07:11:57 -0700206 member_type, first_enum = metacls._get_mixins_(cls, bases)
Ethan Furmanc2667362020-12-07 00:17:31 -0800207 __new__, save_new, use_args = metacls._find_new_(
208 classdict, member_type, first_enum,
209 )
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700210
211 # save enum items into separate mapping so they don't get baked into
212 # the new class
Ethan Furmanee47e5c2016-08-31 00:12:15 -0700213 enum_members = {k: classdict[k] for k in classdict._member_names}
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700214 for name in classdict._member_names:
215 del classdict[name]
216
Ethan Furmane8e61272016-08-20 07:19:31 -0700217 # adjust the sunders
218 _order_ = classdict.pop('_order_', None)
219
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700220 # check for illegal enum names (any others?)
Brennan D Baraban8b914d22019-03-03 14:09:11 -0800221 invalid_names = set(enum_members) & {'mro', ''}
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700222 if invalid_names:
223 raise ValueError('Invalid enum member name: {0}'.format(
224 ','.join(invalid_names)))
225
Ethan Furman48a724f2015-04-11 23:23:06 -0700226 # create a default docstring if one has not been provided
227 if '__doc__' not in classdict:
228 classdict['__doc__'] = 'An enumeration.'
229
Ethan Furman6bd94de2020-12-09 16:41:22 -0800230 # postpone calling __init_subclass__
231 if '__init_subclass__' in classdict and classdict['__init_subclass__'] is None:
232 raise TypeError('%s.__init_subclass__ cannot be None')
233 # remove current __init_subclass__ so previous one can be found with getattr
234 new_init_subclass = classdict.pop('__init_subclass__', None)
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700235 # create our new Enum type
Ethan Furman6bd94de2020-12-09 16:41:22 -0800236 if bases:
237 bases = (_NoInitSubclass, ) + bases
Ethan Furman786d97a2020-12-24 19:31:10 -0800238 enum_class = super().__new__(metacls, cls, bases, classdict, **kwds)
Ethan Furman6bd94de2020-12-09 16:41:22 -0800239 enum_class.__bases__ = enum_class.__bases__[1:] #or (object, )
240 else:
Ethan Furman786d97a2020-12-24 19:31:10 -0800241 enum_class = super().__new__(metacls, cls, bases, classdict, **kwds)
Ethan Furman6bd94de2020-12-09 16:41:22 -0800242 old_init_subclass = getattr(enum_class, '__init_subclass__', None)
243 # and restore the new one (if there was one)
244 if new_init_subclass is not None:
245 enum_class.__init_subclass__ = classmethod(new_init_subclass)
Ethan Furman520ad572013-07-19 19:47:21 -0700246 enum_class._member_names_ = [] # names in definition order
INADA Naokie57f91a2018-06-19 01:14:26 +0900247 enum_class._member_map_ = {} # name->value map
Ethan Furman5e5a8232013-08-04 08:42:23 -0700248 enum_class._member_type_ = member_type
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700249
orlnub1230fb9fad2018-09-12 20:28:53 +0300250 # save DynamicClassAttribute attributes from super classes so we know
251 # if we can take the shortcut of storing members in the class dict
Ethan Furman6d3dfee2020-12-08 12:26:56 -0800252 dynamic_attributes = {
253 k for c in enum_class.mro()
254 for k, v in c.__dict__.items()
255 if isinstance(v, DynamicClassAttribute)
256 }
Ethan Furman354ecf12015-03-11 08:43:12 -0700257
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700258 # Reverse value->name map for hashable values.
Ethan Furman520ad572013-07-19 19:47:21 -0700259 enum_class._value2member_map_ = {}
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700260
Ethan Furman2da95042014-03-03 12:42:52 -0800261 # If a custom type is mixed into the Enum, and it does not know how
262 # to pickle itself, pickle.dumps will succeed but pickle.loads will
263 # fail. Rather than have the error show up later and possibly far
264 # from the source, sabotage the pickle protocol for this class so
265 # that pickle.dumps also fails.
266 #
267 # However, if the new class implements its own __reduce_ex__, do not
268 # sabotage -- it's on them to make sure it works correctly. We use
269 # __reduce_ex__ instead of any of the others as it is preferred by
270 # pickle over __reduce__, and it handles all pickle protocols.
271 if '__reduce_ex__' not in classdict:
Ethan Furmandc870522014-02-18 12:37:12 -0800272 if member_type is not object:
273 methods = ('__getnewargs_ex__', '__getnewargs__',
274 '__reduce_ex__', '__reduce__')
Ethan Furman2da95042014-03-03 12:42:52 -0800275 if not any(m in member_type.__dict__ for m in methods):
Ethan Furmandc870522014-02-18 12:37:12 -0800276 _make_class_unpicklable(enum_class)
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700277
278 # instantiate them, checking for duplicates as we go
279 # we instantiate first instead of checking for duplicates first in case
280 # a custom __new__ is doing something funky with the values -- such as
281 # auto-numbering ;)
282 for member_name in classdict._member_names:
Ethan Furmanee47e5c2016-08-31 00:12:15 -0700283 value = enum_members[member_name]
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700284 if not isinstance(value, tuple):
285 args = (value, )
286 else:
287 args = value
288 if member_type is tuple: # special case for tuple enums
289 args = (args, ) # wrap it one more time
290 if not use_args:
291 enum_member = __new__(enum_class)
Ethan Furmanb41803e2013-07-25 13:50:45 -0700292 if not hasattr(enum_member, '_value_'):
293 enum_member._value_ = value
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700294 else:
295 enum_member = __new__(enum_class, *args)
Ethan Furmanb41803e2013-07-25 13:50:45 -0700296 if not hasattr(enum_member, '_value_'):
Ethan Furmanee47e5c2016-08-31 00:12:15 -0700297 if member_type is object:
298 enum_member._value_ = value
299 else:
300 enum_member._value_ = member_type(*args)
Ethan Furman520ad572013-07-19 19:47:21 -0700301 value = enum_member._value_
Ethan Furman520ad572013-07-19 19:47:21 -0700302 enum_member._name_ = member_name
Ethan Furmanc850f342013-09-15 16:59:35 -0700303 enum_member.__objclass__ = enum_class
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700304 enum_member.__init__(*args)
305 # If another member with the same value was already defined, the
306 # new member becomes an alias to the existing one.
Ethan Furman520ad572013-07-19 19:47:21 -0700307 for name, canonical_member in enum_class._member_map_.items():
Ethan Furman0081f232014-09-16 17:31:23 -0700308 if canonical_member._value_ == enum_member._value_:
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700309 enum_member = canonical_member
310 break
311 else:
312 # Aliases don't appear in member names (only in __members__).
Ethan Furman520ad572013-07-19 19:47:21 -0700313 enum_class._member_names_.append(member_name)
Ethan Furman354ecf12015-03-11 08:43:12 -0700314 # performance boost for any member that would not shadow
315 # a DynamicClassAttribute
orlnub1230fb9fad2018-09-12 20:28:53 +0300316 if member_name not in dynamic_attributes:
Ethan Furman354ecf12015-03-11 08:43:12 -0700317 setattr(enum_class, member_name, enum_member)
318 # now add to _member_map_
Ethan Furman520ad572013-07-19 19:47:21 -0700319 enum_class._member_map_[member_name] = enum_member
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700320 try:
321 # This may fail if value is not hashable. We can't add the value
322 # to the map, and by-value lookups for this value will be
323 # linear.
Ethan Furman520ad572013-07-19 19:47:21 -0700324 enum_class._value2member_map_[value] = enum_member
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700325 except TypeError:
326 pass
327
328 # double check that repr and friends are not the mixin's or various
329 # things break (such as pickle)
Ethan Furman22415ad2020-09-15 16:28:25 -0700330 # however, if the method is defined in the Enum itself, don't replace
331 # it
Ethan Furmandc870522014-02-18 12:37:12 -0800332 for name in ('__repr__', '__str__', '__format__', '__reduce_ex__'):
Ethan Furman22415ad2020-09-15 16:28:25 -0700333 if name in classdict:
334 continue
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700335 class_method = getattr(enum_class, name)
336 obj_method = getattr(member_type, name, None)
337 enum_method = getattr(first_enum, name, None)
338 if obj_method is not None and obj_method is class_method:
339 setattr(enum_class, name, enum_method)
340
341 # replace any other __new__ with our own (as long as Enum is not None,
342 # anyway) -- again, this is to support pickle
343 if Enum is not None:
344 # if the user defined their own __new__, save it before it gets
345 # clobbered in case they subclass later
346 if save_new:
347 enum_class.__new_member__ = __new__
348 enum_class.__new__ = Enum.__new__
Ethan Furmane8e61272016-08-20 07:19:31 -0700349
350 # py3 support for definition order (helps keep py2/py3 code in sync)
351 if _order_ is not None:
352 if isinstance(_order_, str):
353 _order_ = _order_.replace(',', ' ').split()
354 if _order_ != enum_class._member_names_:
355 raise TypeError('member order does not match _order_')
356
Ethan Furman6bd94de2020-12-09 16:41:22 -0800357 # finally, call parents' __init_subclass__
358 if Enum is not None and old_init_subclass is not None:
359 old_init_subclass(**kwds)
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700360 return enum_class
361
Ethan Furman5de67b12016-04-13 23:52:09 -0700362 def __bool__(self):
363 """
364 classes/types should always be True.
365 """
366 return True
367
Ethan Furmand9925a12014-09-16 20:35:55 -0700368 def __call__(cls, value, names=None, *, module=None, qualname=None, type=None, start=1):
Ethan Furman6d3dfee2020-12-08 12:26:56 -0800369 """
370 Either returns an existing member, or creates a new enum class.
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700371
372 This method is used both when an enum class is given a value to match
373 to an enumeration member (i.e. Color(3)) and for the functional API
Ethan Furman23bb6f42016-11-21 09:22:05 -0800374 (i.e. Color = Enum('Color', names='RED GREEN BLUE')).
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700375
Ethan Furman2da95042014-03-03 12:42:52 -0800376 When used for the functional API:
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700377
Ethan Furman2da95042014-03-03 12:42:52 -0800378 `value` will be the name of the new class.
379
380 `names` should be either a string of white-space/comma delimited names
Ethan Furmand9925a12014-09-16 20:35:55 -0700381 (values will start at `start`), or an iterator/mapping of name, value pairs.
Ethan Furman2da95042014-03-03 12:42:52 -0800382
383 `module` should be set to the module this class is being created in;
384 if it is not set, an attempt to find that module will be made, but if
385 it fails the class will not be picklable.
386
387 `qualname` should be set to the actual location this class can be found
388 at in its module; by default it is set to the global scope. If this is
389 not correct, unpickling will fail in some circumstances.
390
391 `type`, if set, will be mixed in as the first base class.
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700392 """
393 if names is None: # simple value lookup
394 return cls.__new__(cls, value)
395 # otherwise, functional API: we're creating a new Enum type
Ethan Furman6d3dfee2020-12-08 12:26:56 -0800396 return cls._create_(
397 value,
398 names,
399 module=module,
400 qualname=qualname,
401 type=type,
402 start=start,
403 )
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700404
405 def __contains__(cls, member):
Rahul Jha94306522018-09-10 23:51:04 +0530406 if not isinstance(member, Enum):
407 raise TypeError(
408 "unsupported operand type(s) for 'in': '%s' and '%s'" % (
409 type(member).__qualname__, cls.__class__.__qualname__))
Ethan Furman0081f232014-09-16 17:31:23 -0700410 return isinstance(member, cls) and member._name_ in cls._member_map_
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700411
Ethan Furman64a99722013-09-22 16:18:19 -0700412 def __delattr__(cls, attr):
413 # nicer error message when someone tries to delete an attribute
414 # (see issue19025).
415 if attr in cls._member_map_:
Ethan Furman6d3dfee2020-12-08 12:26:56 -0800416 raise AttributeError("%s: cannot delete Enum member %r." % (cls.__name__, attr))
Ethan Furman64a99722013-09-22 16:18:19 -0700417 super().__delattr__(attr)
418
Ethan Furman388a3922013-08-12 06:51:41 -0700419 def __dir__(self):
Ethan Furman6d3dfee2020-12-08 12:26:56 -0800420 return (
421 ['__class__', '__doc__', '__members__', '__module__']
422 + self._member_names_
423 )
Ethan Furman388a3922013-08-12 06:51:41 -0700424
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700425 def __getattr__(cls, name):
Ethan Furman6d3dfee2020-12-08 12:26:56 -0800426 """
427 Return the enum member matching `name`
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700428
429 We use __getattr__ instead of descriptors or inserting into the enum
430 class' __dict__ in order to support `name` and `value` being both
431 properties for enum members (which live in the class' __dict__) and
432 enum members themselves.
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700433 """
434 if _is_dunder(name):
435 raise AttributeError(name)
436 try:
Ethan Furman520ad572013-07-19 19:47:21 -0700437 return cls._member_map_[name]
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700438 except KeyError:
439 raise AttributeError(name) from None
440
441 def __getitem__(cls, name):
Ethan Furman520ad572013-07-19 19:47:21 -0700442 return cls._member_map_[name]
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700443
444 def __iter__(cls):
Ethan Furman6d3dfee2020-12-08 12:26:56 -0800445 """
446 Returns members in definition order.
447 """
Ethan Furman520ad572013-07-19 19:47:21 -0700448 return (cls._member_map_[name] for name in cls._member_names_)
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700449
450 def __len__(cls):
Ethan Furman520ad572013-07-19 19:47:21 -0700451 return len(cls._member_names_)
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700452
Ethan Furman2131a4a2013-09-14 18:11:24 -0700453 @property
454 def __members__(cls):
Ethan Furman6d3dfee2020-12-08 12:26:56 -0800455 """
456 Returns a mapping of member name->value.
Ethan Furman2131a4a2013-09-14 18:11:24 -0700457
458 This mapping lists all enum members, including aliases. Note that this
459 is a read-only view of the internal mapping.
Ethan Furman2131a4a2013-09-14 18:11:24 -0700460 """
461 return MappingProxyType(cls._member_map_)
462
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700463 def __repr__(cls):
464 return "<enum %r>" % cls.__name__
465
Ethan Furman2131a4a2013-09-14 18:11:24 -0700466 def __reversed__(cls):
Ethan Furman6d3dfee2020-12-08 12:26:56 -0800467 """
468 Returns members in reverse definition order.
469 """
Ethan Furman2131a4a2013-09-14 18:11:24 -0700470 return (cls._member_map_[name] for name in reversed(cls._member_names_))
471
Ethan Furmanf203f2d2013-09-06 07:16:48 -0700472 def __setattr__(cls, name, value):
Ethan Furman6d3dfee2020-12-08 12:26:56 -0800473 """
474 Block attempts to reassign Enum members.
Ethan Furmanf203f2d2013-09-06 07:16:48 -0700475
476 A simple assignment to the class namespace only changes one of the
477 several possible ways to get an Enum member from the Enum class,
478 resulting in an inconsistent Enumeration.
Ethan Furmanf203f2d2013-09-06 07:16:48 -0700479 """
480 member_map = cls.__dict__.get('_member_map_', {})
481 if name in member_map:
482 raise AttributeError('Cannot reassign members.')
483 super().__setattr__(name, value)
484
anentropicb8e21f12018-04-16 04:40:35 +0100485 def _create_(cls, class_name, names, *, module=None, qualname=None, type=None, start=1):
Ethan Furman6d3dfee2020-12-08 12:26:56 -0800486 """
487 Convenience method to create a new Enum class.
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700488
489 `names` can be:
490
491 * A string containing member names, separated either with spaces or
Ethan Furmand9925a12014-09-16 20:35:55 -0700492 commas. Values are incremented by 1 from `start`.
493 * An iterable of member names. Values are incremented by 1 from `start`.
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700494 * An iterable of (member name, value) pairs.
Ethan Furmand9925a12014-09-16 20:35:55 -0700495 * A mapping of member name -> value pairs.
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700496 """
497 metacls = cls.__class__
498 bases = (cls, ) if type is None else (type, cls)
Ethan Furman3064dbf2020-09-16 07:11:57 -0700499 _, first_enum = cls._get_mixins_(cls, bases)
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700500 classdict = metacls.__prepare__(class_name, bases)
501
502 # special processing needed for names?
503 if isinstance(names, str):
504 names = names.replace(',', ' ').split()
Dong-hee Nadcc8ce42017-06-22 01:52:32 +0900505 if isinstance(names, (tuple, list)) and names and isinstance(names[0], str):
Ethan Furmanee47e5c2016-08-31 00:12:15 -0700506 original_names, names = names, []
Ethan Furmanc16595e2016-09-10 23:36:59 -0700507 last_values = []
Ethan Furmanee47e5c2016-08-31 00:12:15 -0700508 for count, name in enumerate(original_names):
Ethan Furmanc16595e2016-09-10 23:36:59 -0700509 value = first_enum._generate_next_value_(name, start, count, last_values[:])
510 last_values.append(value)
511 names.append((name, value))
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700512
513 # Here, names is either an iterable of (name, value) or a mapping.
514 for item in names:
515 if isinstance(item, str):
516 member_name, member_value = item, names[item]
517 else:
518 member_name, member_value = item
519 classdict[member_name] = member_value
520 enum_class = metacls.__new__(metacls, class_name, bases, classdict)
521
522 # TODO: replace the frame hack if a blessed way to know the calling
523 # module is ever developed
524 if module is None:
525 try:
526 module = sys._getframe(2).f_globals['__name__']
Pablo Galindo293dd232019-11-19 21:34:03 +0000527 except (AttributeError, ValueError, KeyError):
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700528 pass
529 if module is None:
530 _make_class_unpicklable(enum_class)
531 else:
532 enum_class.__module__ = module
Ethan Furmanca1b7942014-02-08 11:36:27 -0800533 if qualname is not None:
534 enum_class.__qualname__ = qualname
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700535
536 return enum_class
537
orlnub1230fb9fad2018-09-12 20:28:53 +0300538 def _convert_(cls, name, module, filter, source=None):
539 """
540 Create a new Enum subclass that replaces a collection of global constants
541 """
542 # convert all constants from source (or module) that pass filter() to
543 # a new Enum called name, and export the enum and its members back to
544 # module;
545 # also, replace the __reduce_ex__ method so unpickling works in
546 # previous Python versions
547 module_globals = vars(sys.modules[module])
548 if source:
549 source = vars(source)
550 else:
551 source = module_globals
552 # _value2member_map_ is populated in the same order every time
553 # for a consistent reverse mapping of number to name when there
554 # are multiple names for the same number.
555 members = [
556 (name, value)
557 for name, value in source.items()
558 if filter(name)]
559 try:
560 # sort by value
561 members.sort(key=lambda t: (t[1], t[0]))
562 except TypeError:
563 # unless some values aren't comparable, in which case sort by name
564 members.sort(key=lambda t: t[0])
565 cls = cls(name, members, module=module)
566 cls.__reduce_ex__ = _reduce_ex_by_name
567 module_globals.update(cls.__members__)
568 module_globals[name] = cls
569 return cls
570
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700571 @staticmethod
Ethan Furman3064dbf2020-09-16 07:11:57 -0700572 def _check_for_existing_members(class_name, bases):
573 for chain in bases:
574 for base in chain.__mro__:
575 if issubclass(base, Enum) and base._member_names_:
Ethan Furman6d3dfee2020-12-08 12:26:56 -0800576 raise TypeError(
577 "%s: cannot extend enumeration %r"
578 % (class_name, base.__name__)
579 )
Ethan Furman3064dbf2020-09-16 07:11:57 -0700580
581 @staticmethod
582 def _get_mixins_(class_name, bases):
Ethan Furman6d3dfee2020-12-08 12:26:56 -0800583 """
584 Returns the type for creating enum members, and the first inherited
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700585 enum class.
586
587 bases: the tuple of bases that was given to __new__
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700588 """
589 if not bases:
590 return object, Enum
591
Ethan Furman5bdab642018-09-21 19:03:09 -0700592 def _find_data_type(bases):
Ethan Furmanbff01f32020-09-15 15:56:26 -0700593 data_types = []
Ethan Furman5bdab642018-09-21 19:03:09 -0700594 for chain in bases:
Ethan Furmanbff01f32020-09-15 15:56:26 -0700595 candidate = None
Ethan Furman5bdab642018-09-21 19:03:09 -0700596 for base in chain.__mro__:
597 if base is object:
598 continue
Ethan Furmanc2667362020-12-07 00:17:31 -0800599 elif issubclass(base, Enum):
600 if base._member_type_ is not object:
601 data_types.append(base._member_type_)
602 break
Ethan Furman5bdab642018-09-21 19:03:09 -0700603 elif '__new__' in base.__dict__:
Ethan Furmancd453852018-10-05 23:29:36 -0700604 if issubclass(base, Enum):
Ethan Furman5bdab642018-09-21 19:03:09 -0700605 continue
Ethan Furmanbff01f32020-09-15 15:56:26 -0700606 data_types.append(candidate or base)
607 break
Ethan Furmanc2667362020-12-07 00:17:31 -0800608 else:
Ethan Furmanbff01f32020-09-15 15:56:26 -0700609 candidate = base
610 if len(data_types) > 1:
Ethan Furman3064dbf2020-09-16 07:11:57 -0700611 raise TypeError('%r: too many data types: %r' % (class_name, data_types))
Ethan Furmanbff01f32020-09-15 15:56:26 -0700612 elif data_types:
613 return data_types[0]
614 else:
615 return None
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700616
Ethan Furman5bdab642018-09-21 19:03:09 -0700617 # ensure final parent class is an Enum derivative, find any concrete
618 # data type, and check that Enum has no members
619 first_enum = bases[-1]
620 if not issubclass(first_enum, Enum):
621 raise TypeError("new enumerations should be created as "
622 "`EnumName([mixin_type, ...] [data_type,] enum_type)`")
623 member_type = _find_data_type(bases) or object
624 if first_enum._member_names_:
625 raise TypeError("Cannot extend enumerations")
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700626 return member_type, first_enum
627
628 @staticmethod
629 def _find_new_(classdict, member_type, first_enum):
Ethan Furman6d3dfee2020-12-08 12:26:56 -0800630 """
631 Returns the __new__ to be used for creating the enum members.
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700632
633 classdict: the class dictionary given to __new__
634 member_type: the data type whose __new__ will be used by default
635 first_enum: enumeration to check for an overriding __new__
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700636 """
637 # now find the correct __new__, checking to see of one was defined
638 # by the user; also check earlier enum classes in case a __new__ was
639 # saved as __new_member__
640 __new__ = classdict.get('__new__', None)
641
642 # should __new__ be saved as __new_member__ later?
643 save_new = __new__ is not None
644
645 if __new__ is None:
646 # check all possibles for __new_member__ before falling back to
647 # __new__
648 for method in ('__new_member__', '__new__'):
649 for possible in (member_type, first_enum):
650 target = getattr(possible, method, None)
651 if target not in {
652 None,
653 None.__new__,
654 object.__new__,
655 Enum.__new__,
656 }:
657 __new__ = target
658 break
659 if __new__ is not None:
660 break
661 else:
662 __new__ = object.__new__
663
664 # if a non-object.__new__ is used then whatever value/tuple was
665 # assigned to the enum member name will be passed to __new__ and to the
666 # new enum member's __init__
667 if __new__ is object.__new__:
668 use_args = False
669 else:
670 use_args = True
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700671 return __new__, save_new, use_args
672
673
674class Enum(metaclass=EnumMeta):
Ethan Furman6d3dfee2020-12-08 12:26:56 -0800675 """
676 Generic enumeration.
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700677
678 Derive from this class to define new enumerations.
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700679 """
680 def __new__(cls, value):
681 # all enum instances are actually created during class construction
682 # without calling this method; this method is called by the metaclass'
683 # __call__ (i.e. Color(3) ), and by pickle
684 if type(value) is cls:
Ethan Furman23bb6f42016-11-21 09:22:05 -0800685 # For lookups like Color(Color.RED)
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700686 return value
687 # by-value search for a matching enum member
688 # see if it's in the reverse mapping (for hashable values)
Ethan Furman2aa27322013-07-19 19:35:56 -0700689 try:
Andrew Svetlov34ae04f2018-12-26 20:45:33 +0200690 return cls._value2member_map_[value]
691 except KeyError:
692 # Not found, no need to do long O(n) search
693 pass
Ethan Furman2aa27322013-07-19 19:35:56 -0700694 except TypeError:
695 # not there, now do long search -- O(n) behavior
Ethan Furman520ad572013-07-19 19:47:21 -0700696 for member in cls._member_map_.values():
Ethan Furman0081f232014-09-16 17:31:23 -0700697 if member._value_ == value:
Ethan Furman2aa27322013-07-19 19:35:56 -0700698 return member
Ethan Furmanee47e5c2016-08-31 00:12:15 -0700699 # still not found -- try _missing_ hook
Ethan Furman019f0a02018-09-12 11:43:34 -0700700 try:
701 exc = None
702 result = cls._missing_(value)
703 except Exception as e:
704 exc = e
705 result = None
706 if isinstance(result, cls):
707 return result
708 else:
Walter Dörwald323842c2019-07-18 20:37:13 +0200709 ve_exc = ValueError("%r is not a valid %s" % (value, cls.__qualname__))
Ethan Furman019f0a02018-09-12 11:43:34 -0700710 if result is None and exc is None:
711 raise ve_exc
712 elif exc is None:
713 exc = TypeError(
714 'error in %s._missing_: returned %r instead of None or a valid member'
715 % (cls.__name__, result)
716 )
717 exc.__context__ = ve_exc
718 raise exc
Ethan Furmanee47e5c2016-08-31 00:12:15 -0700719
Ethan Furmanc16595e2016-09-10 23:36:59 -0700720 def _generate_next_value_(name, start, count, last_values):
Ethan Furman6d3dfee2020-12-08 12:26:56 -0800721 """
722 Generate the next value when not given.
723
724 name: the name of the member
725 start: the initial start value or None
726 count: the number of existing members
727 last_value: the last value assigned or None
728 """
Ethan Furmanc16595e2016-09-10 23:36:59 -0700729 for last_value in reversed(last_values):
730 try:
731 return last_value + 1
732 except TypeError:
733 pass
734 else:
Ethan Furmanee47e5c2016-08-31 00:12:15 -0700735 return start
Ethan Furmanc16595e2016-09-10 23:36:59 -0700736
Ethan Furman6bd94de2020-12-09 16:41:22 -0800737 def __init_subclass__(cls, **kwds):
738 super().__init_subclass__(**kwds)
739
Ethan Furmanee47e5c2016-08-31 00:12:15 -0700740 @classmethod
741 def _missing_(cls, value):
Ethan Furmanc95ad7a2020-09-16 10:26:50 -0700742 return None
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700743
744 def __repr__(self):
745 return "<%s.%s: %r>" % (
Ethan Furman520ad572013-07-19 19:47:21 -0700746 self.__class__.__name__, self._name_, self._value_)
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700747
748 def __str__(self):
Ethan Furman520ad572013-07-19 19:47:21 -0700749 return "%s.%s" % (self.__class__.__name__, self._name_)
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700750
Ethan Furman388a3922013-08-12 06:51:41 -0700751 def __dir__(self):
Ethan Furman6d3dfee2020-12-08 12:26:56 -0800752 """
753 Returns all members and all public methods
754 """
Ethan Furman0ae550b2014-10-14 08:58:32 -0700755 added_behavior = [
756 m
757 for cls in self.__class__.mro()
758 for m in cls.__dict__
Ethan Furman354ecf12015-03-11 08:43:12 -0700759 if m[0] != '_' and m not in self._member_map_
Angelin BOOZ68526fe2020-09-21 15:11:06 +0200760 ] + [m for m in self.__dict__ if m[0] != '_']
Ethan Furmanec5f8eb2014-10-21 13:40:35 -0700761 return (['__class__', '__doc__', '__module__'] + added_behavior)
Ethan Furman388a3922013-08-12 06:51:41 -0700762
Ethan Furmanec15a822013-08-31 19:17:41 -0700763 def __format__(self, format_spec):
Ethan Furman6d3dfee2020-12-08 12:26:56 -0800764 """
765 Returns format using actual value type unless __str__ has been overridden.
766 """
Ethan Furmanec15a822013-08-31 19:17:41 -0700767 # mixed-in Enums should use the mixed-in type's __format__, otherwise
768 # we can get strange results with the Enum name showing up instead of
769 # the value
770
thatneat2f19e822019-07-04 11:28:37 -0700771 # pure Enum branch, or branch with __str__ explicitly overridden
Ethan Furman37440ee2020-12-08 11:14:10 -0800772 str_overridden = type(self).__str__ not in (Enum.__str__, Flag.__str__)
thatneat2f19e822019-07-04 11:28:37 -0700773 if self._member_type_ is object or str_overridden:
Ethan Furmanec15a822013-08-31 19:17:41 -0700774 cls = str
775 val = str(self)
776 # mix-in branch
777 else:
778 cls = self._member_type_
Ethan Furman0081f232014-09-16 17:31:23 -0700779 val = self._value_
Ethan Furmanec15a822013-08-31 19:17:41 -0700780 return cls.__format__(val, format_spec)
781
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700782 def __hash__(self):
Ethan Furman520ad572013-07-19 19:47:21 -0700783 return hash(self._name_)
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700784
Ethan Furmanca1b7942014-02-08 11:36:27 -0800785 def __reduce_ex__(self, proto):
Ethan Furmandc870522014-02-18 12:37:12 -0800786 return self.__class__, (self._value_, )
Ethan Furmanca1b7942014-02-08 11:36:27 -0800787
Ethan Furman33918c12013-09-27 23:02:02 -0700788 # DynamicClassAttribute is used to provide access to the `name` and
789 # `value` properties of enum members while keeping some measure of
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700790 # protection from modification, while still allowing for an enumeration
791 # to have members named `name` and `value`. This works because enumeration
792 # members are not set directly on the enum class -- __getattr__ is
793 # used to look them up.
794
Ethan Furmane03ea372013-09-25 07:14:41 -0700795 @DynamicClassAttribute
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700796 def name(self):
Ethan Furmanc850f342013-09-15 16:59:35 -0700797 """The name of the Enum member."""
Ethan Furman520ad572013-07-19 19:47:21 -0700798 return self._name_
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700799
Ethan Furmane03ea372013-09-25 07:14:41 -0700800 @DynamicClassAttribute
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700801 def value(self):
Ethan Furmanc850f342013-09-15 16:59:35 -0700802 """The value of the Enum member."""
Ethan Furman520ad572013-07-19 19:47:21 -0700803 return self._value_
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700804
805
806class IntEnum(int, Enum):
Ethan Furman0063ff42020-09-21 17:23:13 -0700807 """
808 Enum where members are also (and must be) ints
809 """
810
811
812class StrEnum(str, Enum):
813 """
814 Enum where members are also (and must be) strings
815 """
816
817 def __new__(cls, *values):
818 if len(values) > 3:
819 raise TypeError('too many arguments for str(): %r' % (values, ))
820 if len(values) == 1:
821 # it must be a string
822 if not isinstance(values[0], str):
823 raise TypeError('%r is not a string' % (values[0], ))
824 if len(values) > 1:
825 # check that encoding argument is a string
826 if not isinstance(values[1], str):
827 raise TypeError('encoding must be a string, not %r' % (values[1], ))
828 if len(values) > 2:
829 # check that errors argument is a string
830 if not isinstance(values[2], str):
831 raise TypeError('errors must be a string, not %r' % (values[2], ))
832 value = str(*values)
833 member = str.__new__(cls, value)
834 member._value_ = value
835 return member
Ethan Furmanf24bb352013-07-18 17:05:39 -0700836
Ethan Furmand986d162020-09-22 13:00:07 -0700837 __str__ = str.__str__
838
Ethan Furmanefb13be2020-12-10 12:20:06 -0800839 def _generate_next_value_(name, start, count, last_values):
840 """
841 Return the lower-cased version of the member name.
842 """
843 return name.lower()
844
Ethan Furmanf24bb352013-07-18 17:05:39 -0700845
Ethan Furman24e837f2015-03-18 17:27:57 -0700846def _reduce_ex_by_name(self, proto):
847 return self.name
848
Ethan Furman65a5a472016-09-01 23:55:19 -0700849class Flag(Enum):
Ethan Furman6d3dfee2020-12-08 12:26:56 -0800850 """
851 Support for flags
852 """
Ethan Furmanc16595e2016-09-10 23:36:59 -0700853
854 def _generate_next_value_(name, start, count, last_values):
Ethan Furmanee47e5c2016-08-31 00:12:15 -0700855 """
856 Generate the next value when not given.
857
858 name: the name of the member
HongWeipengbb16fb22019-09-21 13:22:54 +0800859 start: the initial start value or None
Ethan Furmanee47e5c2016-08-31 00:12:15 -0700860 count: the number of existing members
861 last_value: the last value assigned or None
862 """
863 if not count:
864 return start if start is not None else 1
Ethan Furmanc16595e2016-09-10 23:36:59 -0700865 for last_value in reversed(last_values):
866 try:
867 high_bit = _high_bit(last_value)
868 break
Ethan Furman3515dcc2016-09-18 13:15:41 -0700869 except Exception:
Ethan Furmanc16595e2016-09-10 23:36:59 -0700870 raise TypeError('Invalid Flag value: %r' % last_value) from None
Ethan Furmanee47e5c2016-08-31 00:12:15 -0700871 return 2 ** (high_bit+1)
872
873 @classmethod
874 def _missing_(cls, value):
Ethan Furman6d3dfee2020-12-08 12:26:56 -0800875 """
876 Returns member (possibly creating it) if one can be found for value.
877 """
Ethan Furmanee47e5c2016-08-31 00:12:15 -0700878 original_value = value
879 if value < 0:
880 value = ~value
881 possible_member = cls._create_pseudo_member_(value)
Ethan Furmanee47e5c2016-08-31 00:12:15 -0700882 if original_value < 0:
883 possible_member = ~possible_member
884 return possible_member
885
886 @classmethod
887 def _create_pseudo_member_(cls, value):
Ethan Furman3515dcc2016-09-18 13:15:41 -0700888 """
889 Create a composite member iff value contains only members.
890 """
Ethan Furmanee47e5c2016-08-31 00:12:15 -0700891 pseudo_member = cls._value2member_map_.get(value, None)
892 if pseudo_member is None:
Ethan Furman3515dcc2016-09-18 13:15:41 -0700893 # verify all bits are accounted for
894 _, extra_flags = _decompose(cls, value)
895 if extra_flags:
Walter Dörwald323842c2019-07-18 20:37:13 +0200896 raise ValueError("%r is not a valid %s" % (value, cls.__qualname__))
Ethan Furman3515dcc2016-09-18 13:15:41 -0700897 # construct a singleton enum pseudo-member
Ethan Furmanee47e5c2016-08-31 00:12:15 -0700898 pseudo_member = object.__new__(cls)
899 pseudo_member._name_ = None
900 pseudo_member._value_ = value
Ethan Furman28cf6632017-01-24 12:12:06 -0800901 # use setdefault in case another thread already created a composite
902 # with this value
903 pseudo_member = cls._value2member_map_.setdefault(value, pseudo_member)
Ethan Furmanee47e5c2016-08-31 00:12:15 -0700904 return pseudo_member
905
Ethan Furmanee47e5c2016-08-31 00:12:15 -0700906 def __contains__(self, other):
Ethan Furman6d3dfee2020-12-08 12:26:56 -0800907 """
908 Returns True if self has at least the same flags set as other.
909 """
Ethan Furmanee47e5c2016-08-31 00:12:15 -0700910 if not isinstance(other, self.__class__):
Rahul Jha94306522018-09-10 23:51:04 +0530911 raise TypeError(
912 "unsupported operand type(s) for 'in': '%s' and '%s'" % (
913 type(other).__qualname__, self.__class__.__qualname__))
Ethan Furmanee47e5c2016-08-31 00:12:15 -0700914 return other._value_ & self._value_ == other._value_
915
Ethan Furman7219e272020-09-16 13:01:00 -0700916 def __iter__(self):
Ethan Furman6d3dfee2020-12-08 12:26:56 -0800917 """
918 Returns flags in decreasing value order.
919 """
Ethan Furman7219e272020-09-16 13:01:00 -0700920 members, extra_flags = _decompose(self.__class__, self.value)
921 return (m for m in members if m._value_ != 0)
922
Ethan Furmanee47e5c2016-08-31 00:12:15 -0700923 def __repr__(self):
924 cls = self.__class__
925 if self._name_ is not None:
926 return '<%s.%s: %r>' % (cls.__name__, self._name_, self._value_)
Ethan Furman3515dcc2016-09-18 13:15:41 -0700927 members, uncovered = _decompose(cls, self._value_)
Ethan Furman27682d22016-09-04 11:39:01 -0700928 return '<%s.%s: %r>' % (
929 cls.__name__,
930 '|'.join([str(m._name_ or m._value_) for m in members]),
931 self._value_,
932 )
Ethan Furmanee47e5c2016-08-31 00:12:15 -0700933
934 def __str__(self):
935 cls = self.__class__
936 if self._name_ is not None:
937 return '%s.%s' % (cls.__name__, self._name_)
Ethan Furman3515dcc2016-09-18 13:15:41 -0700938 members, uncovered = _decompose(cls, self._value_)
Ethan Furmanee47e5c2016-08-31 00:12:15 -0700939 if len(members) == 1 and members[0]._name_ is None:
940 return '%s.%r' % (cls.__name__, members[0]._value_)
941 else:
942 return '%s.%s' % (
943 cls.__name__,
944 '|'.join([str(m._name_ or m._value_) for m in members]),
945 )
946
Ethan Furman25d94bb2016-09-02 16:32:32 -0700947 def __bool__(self):
948 return bool(self._value_)
949
Ethan Furmanee47e5c2016-08-31 00:12:15 -0700950 def __or__(self, other):
951 if not isinstance(other, self.__class__):
952 return NotImplemented
953 return self.__class__(self._value_ | other._value_)
954
955 def __and__(self, other):
956 if not isinstance(other, self.__class__):
957 return NotImplemented
958 return self.__class__(self._value_ & other._value_)
959
960 def __xor__(self, other):
961 if not isinstance(other, self.__class__):
962 return NotImplemented
963 return self.__class__(self._value_ ^ other._value_)
964
965 def __invert__(self):
Ethan Furman3515dcc2016-09-18 13:15:41 -0700966 members, uncovered = _decompose(self.__class__, self._value_)
Serhiy Storchaka81108372017-09-26 00:55:55 +0300967 inverted = self.__class__(0)
968 for m in self.__class__:
969 if m not in members and not (m._value_ & self._value_):
970 inverted = inverted | m
Ethan Furmanee47e5c2016-08-31 00:12:15 -0700971 return self.__class__(inverted)
972
973
Ethan Furman65a5a472016-09-01 23:55:19 -0700974class IntFlag(int, Flag):
Ethan Furman6d3dfee2020-12-08 12:26:56 -0800975 """
976 Support for integer-based Flags
977 """
Ethan Furmanee47e5c2016-08-31 00:12:15 -0700978
979 @classmethod
Ethan Furman3515dcc2016-09-18 13:15:41 -0700980 def _missing_(cls, value):
Ethan Furman6d3dfee2020-12-08 12:26:56 -0800981 """
982 Returns member (possibly creating it) if one can be found for value.
983 """
Ethan Furman3515dcc2016-09-18 13:15:41 -0700984 if not isinstance(value, int):
Walter Dörwald323842c2019-07-18 20:37:13 +0200985 raise ValueError("%r is not a valid %s" % (value, cls.__qualname__))
Ethan Furman3515dcc2016-09-18 13:15:41 -0700986 new_member = cls._create_pseudo_member_(value)
987 return new_member
988
989 @classmethod
Ethan Furmanee47e5c2016-08-31 00:12:15 -0700990 def _create_pseudo_member_(cls, value):
Ethan Furman6d3dfee2020-12-08 12:26:56 -0800991 """
992 Create a composite member iff value contains only members.
993 """
Ethan Furmanee47e5c2016-08-31 00:12:15 -0700994 pseudo_member = cls._value2member_map_.get(value, None)
995 if pseudo_member is None:
Ethan Furman3515dcc2016-09-18 13:15:41 -0700996 need_to_create = [value]
997 # get unaccounted for bits
998 _, extra_flags = _decompose(cls, value)
999 # timer = 10
1000 while extra_flags:
1001 # timer -= 1
1002 bit = _high_bit(extra_flags)
1003 flag_value = 2 ** bit
1004 if (flag_value not in cls._value2member_map_ and
1005 flag_value not in need_to_create
1006 ):
1007 need_to_create.append(flag_value)
1008 if extra_flags == -flag_value:
1009 extra_flags = 0
1010 else:
1011 extra_flags ^= flag_value
1012 for value in reversed(need_to_create):
1013 # construct singleton pseudo-members
1014 pseudo_member = int.__new__(cls, value)
1015 pseudo_member._name_ = None
1016 pseudo_member._value_ = value
Ethan Furman28cf6632017-01-24 12:12:06 -08001017 # use setdefault in case another thread already created a composite
1018 # with this value
1019 pseudo_member = cls._value2member_map_.setdefault(value, pseudo_member)
Ethan Furmanee47e5c2016-08-31 00:12:15 -07001020 return pseudo_member
1021
Ethan Furmanee47e5c2016-08-31 00:12:15 -07001022 def __or__(self, other):
1023 if not isinstance(other, (self.__class__, int)):
1024 return NotImplemented
Ethan Furman3515dcc2016-09-18 13:15:41 -07001025 result = self.__class__(self._value_ | self.__class__(other)._value_)
1026 return result
Ethan Furmanee47e5c2016-08-31 00:12:15 -07001027
1028 def __and__(self, other):
1029 if not isinstance(other, (self.__class__, int)):
1030 return NotImplemented
1031 return self.__class__(self._value_ & self.__class__(other)._value_)
1032
1033 def __xor__(self, other):
1034 if not isinstance(other, (self.__class__, int)):
1035 return NotImplemented
1036 return self.__class__(self._value_ ^ self.__class__(other)._value_)
1037
1038 __ror__ = __or__
1039 __rand__ = __and__
1040 __rxor__ = __xor__
1041
1042 def __invert__(self):
Ethan Furman3515dcc2016-09-18 13:15:41 -07001043 result = self.__class__(~self._value_)
1044 return result
Ethan Furmanee47e5c2016-08-31 00:12:15 -07001045
1046
1047def _high_bit(value):
Ethan Furman6d3dfee2020-12-08 12:26:56 -08001048 """
1049 returns index of highest bit, or -1 if value is zero or negative
1050 """
Ethan Furman3515dcc2016-09-18 13:15:41 -07001051 return value.bit_length() - 1
Ethan Furmanee47e5c2016-08-31 00:12:15 -07001052
Ethan Furmanf24bb352013-07-18 17:05:39 -07001053def unique(enumeration):
Ethan Furman6d3dfee2020-12-08 12:26:56 -08001054 """
1055 Class decorator for enumerations ensuring unique member values.
1056 """
Ethan Furmanf24bb352013-07-18 17:05:39 -07001057 duplicates = []
1058 for name, member in enumeration.__members__.items():
1059 if name != member.name:
1060 duplicates.append((name, member.name))
1061 if duplicates:
1062 alias_details = ', '.join(
1063 ["%s -> %s" % (alias, name) for (alias, name) in duplicates])
1064 raise ValueError('duplicate values found in %r: %s' %
1065 (enumeration, alias_details))
1066 return enumeration
Ethan Furman3515dcc2016-09-18 13:15:41 -07001067
1068def _decompose(flag, value):
Ethan Furman6d3dfee2020-12-08 12:26:56 -08001069 """
1070 Extract all members from the value.
1071 """
Ethan Furman3515dcc2016-09-18 13:15:41 -07001072 # _decompose is only called if the value is not named
1073 not_covered = value
1074 negative = value < 0
Ethan Furman3515dcc2016-09-18 13:15:41 -07001075 members = []
HongWeipeng0b41a922019-11-27 06:36:02 +08001076 for member in flag:
1077 member_value = member.value
Ethan Furman3515dcc2016-09-18 13:15:41 -07001078 if member_value and member_value & value == member_value:
1079 members.append(member)
1080 not_covered &= ~member_value
HongWeipeng0b41a922019-11-27 06:36:02 +08001081 if not negative:
1082 tmp = not_covered
1083 while tmp:
1084 flag_value = 2 ** _high_bit(tmp)
1085 if flag_value in flag._value2member_map_:
1086 members.append(flag._value2member_map_[flag_value])
1087 not_covered &= ~flag_value
1088 tmp &= ~flag_value
Ethan Furman3515dcc2016-09-18 13:15:41 -07001089 if not members and value in flag._value2member_map_:
1090 members.append(flag._value2member_map_[value])
1091 members.sort(key=lambda m: m._value_, reverse=True)
1092 if len(members) > 1 and members[0].value == value:
1093 # we have the breakdown, don't need the value member itself
1094 members.pop(0)
1095 return members, not_covered