blob: 83e050e1e41cde32213516db03a5b61ad2bc6a34 [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 Furman6b3d64a2013-06-14 16:55:46 -070052def _make_class_unpicklable(cls):
Ethan Furman6d3dfee2020-12-08 12:26:56 -080053 """
54 Make the given class un-picklable.
55 """
Ethan Furmanca1b7942014-02-08 11:36:27 -080056 def _break_on_call_reduce(self, proto):
Ethan Furman6b3d64a2013-06-14 16:55:46 -070057 raise TypeError('%r cannot be pickled' % self)
Ethan Furmanca1b7942014-02-08 11:36:27 -080058 cls.__reduce_ex__ = _break_on_call_reduce
Ethan Furman6b3d64a2013-06-14 16:55:46 -070059 cls.__module__ = '<unknown>'
60
Ethan Furman3515dcc2016-09-18 13:15:41 -070061_auto_null = object()
Ethan Furmanc16595e2016-09-10 23:36:59 -070062class auto:
63 """
64 Instances are replaced with an appropriate value in Enum class suites.
65 """
Ethan Furman3515dcc2016-09-18 13:15:41 -070066 value = _auto_null
Ethan Furmanc16595e2016-09-10 23:36:59 -070067
Ethan Furman101e0742013-09-15 12:34:36 -070068
Ethan Furman6b3d64a2013-06-14 16:55:46 -070069class _EnumDict(dict):
Ethan Furman6d3dfee2020-12-08 12:26:56 -080070 """
71 Track enum member order and ensure member names are not reused.
Ethan Furman6b3d64a2013-06-14 16:55:46 -070072
73 EnumMeta will use the names found in self._member_names as the
74 enumeration member names.
Ethan Furman6b3d64a2013-06-14 16:55:46 -070075 """
76 def __init__(self):
77 super().__init__()
78 self._member_names = []
Ethan Furmanc16595e2016-09-10 23:36:59 -070079 self._last_values = []
Ethan Furmana4b1bb42018-01-22 07:56:37 -080080 self._ignore = []
Ethan Onstottd9a43e22020-04-28 13:20:55 -040081 self._auto_called = False
Ethan Furman6b3d64a2013-06-14 16:55:46 -070082
83 def __setitem__(self, key, value):
Ethan Furman6d3dfee2020-12-08 12:26:56 -080084 """
85 Changes anything not dundered or not a descriptor.
Ethan Furman6b3d64a2013-06-14 16:55:46 -070086
87 If an enum member name is used twice, an error is raised; duplicate
88 values are not checked for.
89
90 Single underscore (sunder) names are reserved.
Ethan Furman6b3d64a2013-06-14 16:55:46 -070091 """
92 if _is_sunder(key):
Ethan Furmanee47e5c2016-08-31 00:12:15 -070093 if key not in (
Ethan Furman3515dcc2016-09-18 13:15:41 -070094 '_order_', '_create_pseudo_member_',
Ethan Furmana4b1bb42018-01-22 07:56:37 -080095 '_generate_next_value_', '_missing_', '_ignore_',
Ethan Furmanee47e5c2016-08-31 00:12:15 -070096 ):
Ethan Furman6d3dfee2020-12-08 12:26:56 -080097 raise ValueError(
98 '_sunder_ names, such as %r, are reserved for future Enum use'
99 % (key, )
100 )
Ethan Furmanc16595e2016-09-10 23:36:59 -0700101 if key == '_generate_next_value_':
Ethan Onstottd9a43e22020-04-28 13:20:55 -0400102 # check if members already defined as auto()
103 if self._auto_called:
104 raise TypeError("_generate_next_value_ must be defined before members")
Ethan Furmanc16595e2016-09-10 23:36:59 -0700105 setattr(self, '_generate_next_value', value)
Ethan Furmana4b1bb42018-01-22 07:56:37 -0800106 elif key == '_ignore_':
107 if isinstance(value, str):
108 value = value.replace(',',' ').split()
109 else:
110 value = list(value)
111 self._ignore = value
112 already = set(value) & set(self._member_names)
113 if already:
Ethan Furman6d3dfee2020-12-08 12:26:56 -0800114 raise ValueError(
115 '_ignore_ cannot specify already set names: %r'
116 % (already, )
117 )
Ethan Furman101e0742013-09-15 12:34:36 -0700118 elif _is_dunder(key):
Ethan Furmane8e61272016-08-20 07:19:31 -0700119 if key == '__order__':
120 key = '_order_'
Ethan Furman101e0742013-09-15 12:34:36 -0700121 elif key in self._member_names:
122 # descriptor overwriting an enum?
123 raise TypeError('Attempted to reuse key: %r' % key)
Ethan Furmana4b1bb42018-01-22 07:56:37 -0800124 elif key in self._ignore:
125 pass
Ethan Furman101e0742013-09-15 12:34:36 -0700126 elif not _is_descriptor(value):
127 if key in self:
128 # enum overwriting a descriptor?
Ethan Furmanee47e5c2016-08-31 00:12:15 -0700129 raise TypeError('%r already defined as: %r' % (key, self[key]))
Ethan Furmanc16595e2016-09-10 23:36:59 -0700130 if isinstance(value, auto):
Ethan Furman3515dcc2016-09-18 13:15:41 -0700131 if value.value == _auto_null:
Ethan Furman6d3dfee2020-12-08 12:26:56 -0800132 value.value = self._generate_next_value(
133 key,
134 1,
135 len(self._member_names),
136 self._last_values[:],
137 )
Ethan Furmanfc23a942020-09-16 12:37:54 -0700138 self._auto_called = True
Ethan Furman3515dcc2016-09-18 13:15:41 -0700139 value = value.value
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700140 self._member_names.append(key)
Ethan Furmanc16595e2016-09-10 23:36:59 -0700141 self._last_values.append(value)
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700142 super().__setitem__(key, value)
143
144
Ezio Melotti9a3777e2013-08-17 15:53:55 +0300145# Dummy value for Enum as EnumMeta explicitly checks for it, but of course
146# until EnumMeta finishes running the first time the Enum class doesn't exist.
147# This is also why there are checks in EnumMeta like `if Enum is not None`
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700148Enum = None
149
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700150class EnumMeta(type):
Ethan Furman6d3dfee2020-12-08 12:26:56 -0800151 """
152 Metaclass for Enum
153 """
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700154 @classmethod
Ethan Furman332dbc72016-08-20 00:00:52 -0700155 def __prepare__(metacls, cls, bases):
Ethan Furman3064dbf2020-09-16 07:11:57 -0700156 # check that previous enum members do not exist
157 metacls._check_for_existing_members(cls, bases)
Ethan Furmanee47e5c2016-08-31 00:12:15 -0700158 # create the namespace dict
159 enum_dict = _EnumDict()
160 # inherit previous flags and _generate_next_value_ function
Ethan Furman3064dbf2020-09-16 07:11:57 -0700161 member_type, first_enum = metacls._get_mixins_(cls, bases)
Ethan Furmanee47e5c2016-08-31 00:12:15 -0700162 if first_enum is not None:
Ethan Furman6d3dfee2020-12-08 12:26:56 -0800163 enum_dict['_generate_next_value_'] = getattr(
164 first_enum, '_generate_next_value_', None,
165 )
Ethan Furmanee47e5c2016-08-31 00:12:15 -0700166 return enum_dict
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700167
Ethan Furman6bd94de2020-12-09 16:41:22 -0800168 def __new__(metacls, cls, bases, classdict, **kwds):
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700169 # an Enum class is final once enumeration items have been defined; it
170 # cannot be mixed with other types (int, float, etc.) if it has an
171 # inherited __new__ unless a new __new__ is defined (or the resulting
172 # class will fail).
Ethan Furmana4b1bb42018-01-22 07:56:37 -0800173 #
174 # remove any keys listed in _ignore_
175 classdict.setdefault('_ignore_', []).append('_ignore_')
176 ignore = classdict['_ignore_']
177 for key in ignore:
178 classdict.pop(key, None)
Ethan Furman3064dbf2020-09-16 07:11:57 -0700179 member_type, first_enum = metacls._get_mixins_(cls, bases)
Ethan Furmanc2667362020-12-07 00:17:31 -0800180 __new__, save_new, use_args = metacls._find_new_(
181 classdict, member_type, first_enum,
182 )
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700183
184 # save enum items into separate mapping so they don't get baked into
185 # the new class
Ethan Furmanee47e5c2016-08-31 00:12:15 -0700186 enum_members = {k: classdict[k] for k in classdict._member_names}
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700187 for name in classdict._member_names:
188 del classdict[name]
189
Ethan Furmane8e61272016-08-20 07:19:31 -0700190 # adjust the sunders
191 _order_ = classdict.pop('_order_', None)
192
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700193 # check for illegal enum names (any others?)
Brennan D Baraban8b914d22019-03-03 14:09:11 -0800194 invalid_names = set(enum_members) & {'mro', ''}
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700195 if invalid_names:
196 raise ValueError('Invalid enum member name: {0}'.format(
197 ','.join(invalid_names)))
198
Ethan Furman48a724f2015-04-11 23:23:06 -0700199 # create a default docstring if one has not been provided
200 if '__doc__' not in classdict:
201 classdict['__doc__'] = 'An enumeration.'
202
Ethan Furman6bd94de2020-12-09 16:41:22 -0800203 # postpone calling __init_subclass__
204 if '__init_subclass__' in classdict and classdict['__init_subclass__'] is None:
205 raise TypeError('%s.__init_subclass__ cannot be None')
206 # remove current __init_subclass__ so previous one can be found with getattr
207 new_init_subclass = classdict.pop('__init_subclass__', None)
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700208 # create our new Enum type
Ethan Furman6bd94de2020-12-09 16:41:22 -0800209 if bases:
210 bases = (_NoInitSubclass, ) + bases
211 enum_class = type.__new__(metacls, cls, bases, classdict)
212 enum_class.__bases__ = enum_class.__bases__[1:] #or (object, )
213 else:
214 enum_class = type.__new__(metacls, cls, bases, classdict)
215 old_init_subclass = getattr(enum_class, '__init_subclass__', None)
216 # and restore the new one (if there was one)
217 if new_init_subclass is not None:
218 enum_class.__init_subclass__ = classmethod(new_init_subclass)
Ethan Furman520ad572013-07-19 19:47:21 -0700219 enum_class._member_names_ = [] # names in definition order
INADA Naokie57f91a2018-06-19 01:14:26 +0900220 enum_class._member_map_ = {} # name->value map
Ethan Furman5e5a8232013-08-04 08:42:23 -0700221 enum_class._member_type_ = member_type
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700222
orlnub1230fb9fad2018-09-12 20:28:53 +0300223 # save DynamicClassAttribute attributes from super classes so we know
224 # if we can take the shortcut of storing members in the class dict
Ethan Furman6d3dfee2020-12-08 12:26:56 -0800225 dynamic_attributes = {
226 k for c in enum_class.mro()
227 for k, v in c.__dict__.items()
228 if isinstance(v, DynamicClassAttribute)
229 }
Ethan Furman354ecf12015-03-11 08:43:12 -0700230
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700231 # Reverse value->name map for hashable values.
Ethan Furman520ad572013-07-19 19:47:21 -0700232 enum_class._value2member_map_ = {}
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700233
Ethan Furman2da95042014-03-03 12:42:52 -0800234 # If a custom type is mixed into the Enum, and it does not know how
235 # to pickle itself, pickle.dumps will succeed but pickle.loads will
236 # fail. Rather than have the error show up later and possibly far
237 # from the source, sabotage the pickle protocol for this class so
238 # that pickle.dumps also fails.
239 #
240 # However, if the new class implements its own __reduce_ex__, do not
241 # sabotage -- it's on them to make sure it works correctly. We use
242 # __reduce_ex__ instead of any of the others as it is preferred by
243 # pickle over __reduce__, and it handles all pickle protocols.
244 if '__reduce_ex__' not in classdict:
Ethan Furmandc870522014-02-18 12:37:12 -0800245 if member_type is not object:
246 methods = ('__getnewargs_ex__', '__getnewargs__',
247 '__reduce_ex__', '__reduce__')
Ethan Furman2da95042014-03-03 12:42:52 -0800248 if not any(m in member_type.__dict__ for m in methods):
Ethan Furmandc870522014-02-18 12:37:12 -0800249 _make_class_unpicklable(enum_class)
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700250
251 # instantiate them, checking for duplicates as we go
252 # we instantiate first instead of checking for duplicates first in case
253 # a custom __new__ is doing something funky with the values -- such as
254 # auto-numbering ;)
255 for member_name in classdict._member_names:
Ethan Furmanee47e5c2016-08-31 00:12:15 -0700256 value = enum_members[member_name]
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700257 if not isinstance(value, tuple):
258 args = (value, )
259 else:
260 args = value
261 if member_type is tuple: # special case for tuple enums
262 args = (args, ) # wrap it one more time
263 if not use_args:
264 enum_member = __new__(enum_class)
Ethan Furmanb41803e2013-07-25 13:50:45 -0700265 if not hasattr(enum_member, '_value_'):
266 enum_member._value_ = value
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700267 else:
268 enum_member = __new__(enum_class, *args)
Ethan Furmanb41803e2013-07-25 13:50:45 -0700269 if not hasattr(enum_member, '_value_'):
Ethan Furmanee47e5c2016-08-31 00:12:15 -0700270 if member_type is object:
271 enum_member._value_ = value
272 else:
273 enum_member._value_ = member_type(*args)
Ethan Furman520ad572013-07-19 19:47:21 -0700274 value = enum_member._value_
Ethan Furman520ad572013-07-19 19:47:21 -0700275 enum_member._name_ = member_name
Ethan Furmanc850f342013-09-15 16:59:35 -0700276 enum_member.__objclass__ = enum_class
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700277 enum_member.__init__(*args)
278 # If another member with the same value was already defined, the
279 # new member becomes an alias to the existing one.
Ethan Furman520ad572013-07-19 19:47:21 -0700280 for name, canonical_member in enum_class._member_map_.items():
Ethan Furman0081f232014-09-16 17:31:23 -0700281 if canonical_member._value_ == enum_member._value_:
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700282 enum_member = canonical_member
283 break
284 else:
285 # Aliases don't appear in member names (only in __members__).
Ethan Furman520ad572013-07-19 19:47:21 -0700286 enum_class._member_names_.append(member_name)
Ethan Furman354ecf12015-03-11 08:43:12 -0700287 # performance boost for any member that would not shadow
288 # a DynamicClassAttribute
orlnub1230fb9fad2018-09-12 20:28:53 +0300289 if member_name not in dynamic_attributes:
Ethan Furman354ecf12015-03-11 08:43:12 -0700290 setattr(enum_class, member_name, enum_member)
291 # now add to _member_map_
Ethan Furman520ad572013-07-19 19:47:21 -0700292 enum_class._member_map_[member_name] = enum_member
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700293 try:
294 # This may fail if value is not hashable. We can't add the value
295 # to the map, and by-value lookups for this value will be
296 # linear.
Ethan Furman520ad572013-07-19 19:47:21 -0700297 enum_class._value2member_map_[value] = enum_member
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700298 except TypeError:
299 pass
300
301 # double check that repr and friends are not the mixin's or various
302 # things break (such as pickle)
Ethan Furman22415ad2020-09-15 16:28:25 -0700303 # however, if the method is defined in the Enum itself, don't replace
304 # it
Ethan Furmandc870522014-02-18 12:37:12 -0800305 for name in ('__repr__', '__str__', '__format__', '__reduce_ex__'):
Ethan Furman22415ad2020-09-15 16:28:25 -0700306 if name in classdict:
307 continue
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700308 class_method = getattr(enum_class, name)
309 obj_method = getattr(member_type, name, None)
310 enum_method = getattr(first_enum, name, None)
311 if obj_method is not None and obj_method is class_method:
312 setattr(enum_class, name, enum_method)
313
314 # replace any other __new__ with our own (as long as Enum is not None,
315 # anyway) -- again, this is to support pickle
316 if Enum is not None:
317 # if the user defined their own __new__, save it before it gets
318 # clobbered in case they subclass later
319 if save_new:
320 enum_class.__new_member__ = __new__
321 enum_class.__new__ = Enum.__new__
Ethan Furmane8e61272016-08-20 07:19:31 -0700322
323 # py3 support for definition order (helps keep py2/py3 code in sync)
324 if _order_ is not None:
325 if isinstance(_order_, str):
326 _order_ = _order_.replace(',', ' ').split()
327 if _order_ != enum_class._member_names_:
328 raise TypeError('member order does not match _order_')
329
Ethan Furman6bd94de2020-12-09 16:41:22 -0800330 # finally, call parents' __init_subclass__
331 if Enum is not None and old_init_subclass is not None:
332 old_init_subclass(**kwds)
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700333 return enum_class
334
Ethan Furman5de67b12016-04-13 23:52:09 -0700335 def __bool__(self):
336 """
337 classes/types should always be True.
338 """
339 return True
340
Ethan Furmand9925a12014-09-16 20:35:55 -0700341 def __call__(cls, value, names=None, *, module=None, qualname=None, type=None, start=1):
Ethan Furman6d3dfee2020-12-08 12:26:56 -0800342 """
343 Either returns an existing member, or creates a new enum class.
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700344
345 This method is used both when an enum class is given a value to match
346 to an enumeration member (i.e. Color(3)) and for the functional API
Ethan Furman23bb6f42016-11-21 09:22:05 -0800347 (i.e. Color = Enum('Color', names='RED GREEN BLUE')).
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700348
Ethan Furman2da95042014-03-03 12:42:52 -0800349 When used for the functional API:
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700350
Ethan Furman2da95042014-03-03 12:42:52 -0800351 `value` will be the name of the new class.
352
353 `names` should be either a string of white-space/comma delimited names
Ethan Furmand9925a12014-09-16 20:35:55 -0700354 (values will start at `start`), or an iterator/mapping of name, value pairs.
Ethan Furman2da95042014-03-03 12:42:52 -0800355
356 `module` should be set to the module this class is being created in;
357 if it is not set, an attempt to find that module will be made, but if
358 it fails the class will not be picklable.
359
360 `qualname` should be set to the actual location this class can be found
361 at in its module; by default it is set to the global scope. If this is
362 not correct, unpickling will fail in some circumstances.
363
364 `type`, if set, will be mixed in as the first base class.
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700365 """
366 if names is None: # simple value lookup
367 return cls.__new__(cls, value)
368 # otherwise, functional API: we're creating a new Enum type
Ethan Furman6d3dfee2020-12-08 12:26:56 -0800369 return cls._create_(
370 value,
371 names,
372 module=module,
373 qualname=qualname,
374 type=type,
375 start=start,
376 )
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700377
378 def __contains__(cls, member):
Rahul Jha94306522018-09-10 23:51:04 +0530379 if not isinstance(member, Enum):
380 raise TypeError(
381 "unsupported operand type(s) for 'in': '%s' and '%s'" % (
382 type(member).__qualname__, cls.__class__.__qualname__))
Ethan Furman0081f232014-09-16 17:31:23 -0700383 return isinstance(member, cls) and member._name_ in cls._member_map_
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700384
Ethan Furman64a99722013-09-22 16:18:19 -0700385 def __delattr__(cls, attr):
386 # nicer error message when someone tries to delete an attribute
387 # (see issue19025).
388 if attr in cls._member_map_:
Ethan Furman6d3dfee2020-12-08 12:26:56 -0800389 raise AttributeError("%s: cannot delete Enum member %r." % (cls.__name__, attr))
Ethan Furman64a99722013-09-22 16:18:19 -0700390 super().__delattr__(attr)
391
Ethan Furman388a3922013-08-12 06:51:41 -0700392 def __dir__(self):
Ethan Furman6d3dfee2020-12-08 12:26:56 -0800393 return (
394 ['__class__', '__doc__', '__members__', '__module__']
395 + self._member_names_
396 )
Ethan Furman388a3922013-08-12 06:51:41 -0700397
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700398 def __getattr__(cls, name):
Ethan Furman6d3dfee2020-12-08 12:26:56 -0800399 """
400 Return the enum member matching `name`
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700401
402 We use __getattr__ instead of descriptors or inserting into the enum
403 class' __dict__ in order to support `name` and `value` being both
404 properties for enum members (which live in the class' __dict__) and
405 enum members themselves.
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700406 """
407 if _is_dunder(name):
408 raise AttributeError(name)
409 try:
Ethan Furman520ad572013-07-19 19:47:21 -0700410 return cls._member_map_[name]
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700411 except KeyError:
412 raise AttributeError(name) from None
413
414 def __getitem__(cls, name):
Ethan Furman520ad572013-07-19 19:47:21 -0700415 return cls._member_map_[name]
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700416
417 def __iter__(cls):
Ethan Furman6d3dfee2020-12-08 12:26:56 -0800418 """
419 Returns members in definition order.
420 """
Ethan Furman520ad572013-07-19 19:47:21 -0700421 return (cls._member_map_[name] for name in cls._member_names_)
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700422
423 def __len__(cls):
Ethan Furman520ad572013-07-19 19:47:21 -0700424 return len(cls._member_names_)
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700425
Ethan Furman2131a4a2013-09-14 18:11:24 -0700426 @property
427 def __members__(cls):
Ethan Furman6d3dfee2020-12-08 12:26:56 -0800428 """
429 Returns a mapping of member name->value.
Ethan Furman2131a4a2013-09-14 18:11:24 -0700430
431 This mapping lists all enum members, including aliases. Note that this
432 is a read-only view of the internal mapping.
Ethan Furman2131a4a2013-09-14 18:11:24 -0700433 """
434 return MappingProxyType(cls._member_map_)
435
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700436 def __repr__(cls):
437 return "<enum %r>" % cls.__name__
438
Ethan Furman2131a4a2013-09-14 18:11:24 -0700439 def __reversed__(cls):
Ethan Furman6d3dfee2020-12-08 12:26:56 -0800440 """
441 Returns members in reverse definition order.
442 """
Ethan Furman2131a4a2013-09-14 18:11:24 -0700443 return (cls._member_map_[name] for name in reversed(cls._member_names_))
444
Ethan Furmanf203f2d2013-09-06 07:16:48 -0700445 def __setattr__(cls, name, value):
Ethan Furman6d3dfee2020-12-08 12:26:56 -0800446 """
447 Block attempts to reassign Enum members.
Ethan Furmanf203f2d2013-09-06 07:16:48 -0700448
449 A simple assignment to the class namespace only changes one of the
450 several possible ways to get an Enum member from the Enum class,
451 resulting in an inconsistent Enumeration.
Ethan Furmanf203f2d2013-09-06 07:16:48 -0700452 """
453 member_map = cls.__dict__.get('_member_map_', {})
454 if name in member_map:
455 raise AttributeError('Cannot reassign members.')
456 super().__setattr__(name, value)
457
anentropicb8e21f12018-04-16 04:40:35 +0100458 def _create_(cls, class_name, names, *, module=None, qualname=None, type=None, start=1):
Ethan Furman6d3dfee2020-12-08 12:26:56 -0800459 """
460 Convenience method to create a new Enum class.
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700461
462 `names` can be:
463
464 * A string containing member names, separated either with spaces or
Ethan Furmand9925a12014-09-16 20:35:55 -0700465 commas. Values are incremented by 1 from `start`.
466 * An iterable of member names. Values are incremented by 1 from `start`.
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700467 * An iterable of (member name, value) pairs.
Ethan Furmand9925a12014-09-16 20:35:55 -0700468 * A mapping of member name -> value pairs.
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700469 """
470 metacls = cls.__class__
471 bases = (cls, ) if type is None else (type, cls)
Ethan Furman3064dbf2020-09-16 07:11:57 -0700472 _, first_enum = cls._get_mixins_(cls, bases)
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700473 classdict = metacls.__prepare__(class_name, bases)
474
475 # special processing needed for names?
476 if isinstance(names, str):
477 names = names.replace(',', ' ').split()
Dong-hee Nadcc8ce42017-06-22 01:52:32 +0900478 if isinstance(names, (tuple, list)) and names and isinstance(names[0], str):
Ethan Furmanee47e5c2016-08-31 00:12:15 -0700479 original_names, names = names, []
Ethan Furmanc16595e2016-09-10 23:36:59 -0700480 last_values = []
Ethan Furmanee47e5c2016-08-31 00:12:15 -0700481 for count, name in enumerate(original_names):
Ethan Furmanc16595e2016-09-10 23:36:59 -0700482 value = first_enum._generate_next_value_(name, start, count, last_values[:])
483 last_values.append(value)
484 names.append((name, value))
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700485
486 # Here, names is either an iterable of (name, value) or a mapping.
487 for item in names:
488 if isinstance(item, str):
489 member_name, member_value = item, names[item]
490 else:
491 member_name, member_value = item
492 classdict[member_name] = member_value
493 enum_class = metacls.__new__(metacls, class_name, bases, classdict)
494
495 # TODO: replace the frame hack if a blessed way to know the calling
496 # module is ever developed
497 if module is None:
498 try:
499 module = sys._getframe(2).f_globals['__name__']
Pablo Galindo293dd232019-11-19 21:34:03 +0000500 except (AttributeError, ValueError, KeyError):
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700501 pass
502 if module is None:
503 _make_class_unpicklable(enum_class)
504 else:
505 enum_class.__module__ = module
Ethan Furmanca1b7942014-02-08 11:36:27 -0800506 if qualname is not None:
507 enum_class.__qualname__ = qualname
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700508
509 return enum_class
510
orlnub1230fb9fad2018-09-12 20:28:53 +0300511 def _convert_(cls, name, module, filter, source=None):
512 """
513 Create a new Enum subclass that replaces a collection of global constants
514 """
515 # convert all constants from source (or module) that pass filter() to
516 # a new Enum called name, and export the enum and its members back to
517 # module;
518 # also, replace the __reduce_ex__ method so unpickling works in
519 # previous Python versions
520 module_globals = vars(sys.modules[module])
521 if source:
522 source = vars(source)
523 else:
524 source = module_globals
525 # _value2member_map_ is populated in the same order every time
526 # for a consistent reverse mapping of number to name when there
527 # are multiple names for the same number.
528 members = [
529 (name, value)
530 for name, value in source.items()
531 if filter(name)]
532 try:
533 # sort by value
534 members.sort(key=lambda t: (t[1], t[0]))
535 except TypeError:
536 # unless some values aren't comparable, in which case sort by name
537 members.sort(key=lambda t: t[0])
538 cls = cls(name, members, module=module)
539 cls.__reduce_ex__ = _reduce_ex_by_name
540 module_globals.update(cls.__members__)
541 module_globals[name] = cls
542 return cls
543
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700544 @staticmethod
Ethan Furman3064dbf2020-09-16 07:11:57 -0700545 def _check_for_existing_members(class_name, bases):
546 for chain in bases:
547 for base in chain.__mro__:
548 if issubclass(base, Enum) and base._member_names_:
Ethan Furman6d3dfee2020-12-08 12:26:56 -0800549 raise TypeError(
550 "%s: cannot extend enumeration %r"
551 % (class_name, base.__name__)
552 )
Ethan Furman3064dbf2020-09-16 07:11:57 -0700553
554 @staticmethod
555 def _get_mixins_(class_name, bases):
Ethan Furman6d3dfee2020-12-08 12:26:56 -0800556 """
557 Returns the type for creating enum members, and the first inherited
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700558 enum class.
559
560 bases: the tuple of bases that was given to __new__
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700561 """
562 if not bases:
563 return object, Enum
564
Ethan Furman5bdab642018-09-21 19:03:09 -0700565 def _find_data_type(bases):
Ethan Furmanbff01f32020-09-15 15:56:26 -0700566 data_types = []
Ethan Furman5bdab642018-09-21 19:03:09 -0700567 for chain in bases:
Ethan Furmanbff01f32020-09-15 15:56:26 -0700568 candidate = None
Ethan Furman5bdab642018-09-21 19:03:09 -0700569 for base in chain.__mro__:
570 if base is object:
571 continue
Ethan Furmanc2667362020-12-07 00:17:31 -0800572 elif issubclass(base, Enum):
573 if base._member_type_ is not object:
574 data_types.append(base._member_type_)
575 break
Ethan Furman5bdab642018-09-21 19:03:09 -0700576 elif '__new__' in base.__dict__:
Ethan Furmancd453852018-10-05 23:29:36 -0700577 if issubclass(base, Enum):
Ethan Furman5bdab642018-09-21 19:03:09 -0700578 continue
Ethan Furmanbff01f32020-09-15 15:56:26 -0700579 data_types.append(candidate or base)
580 break
Ethan Furmanc2667362020-12-07 00:17:31 -0800581 else:
Ethan Furmanbff01f32020-09-15 15:56:26 -0700582 candidate = base
583 if len(data_types) > 1:
Ethan Furman3064dbf2020-09-16 07:11:57 -0700584 raise TypeError('%r: too many data types: %r' % (class_name, data_types))
Ethan Furmanbff01f32020-09-15 15:56:26 -0700585 elif data_types:
586 return data_types[0]
587 else:
588 return None
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700589
Ethan Furman5bdab642018-09-21 19:03:09 -0700590 # ensure final parent class is an Enum derivative, find any concrete
591 # data type, and check that Enum has no members
592 first_enum = bases[-1]
593 if not issubclass(first_enum, Enum):
594 raise TypeError("new enumerations should be created as "
595 "`EnumName([mixin_type, ...] [data_type,] enum_type)`")
596 member_type = _find_data_type(bases) or object
597 if first_enum._member_names_:
598 raise TypeError("Cannot extend enumerations")
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700599 return member_type, first_enum
600
601 @staticmethod
602 def _find_new_(classdict, member_type, first_enum):
Ethan Furman6d3dfee2020-12-08 12:26:56 -0800603 """
604 Returns the __new__ to be used for creating the enum members.
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700605
606 classdict: the class dictionary given to __new__
607 member_type: the data type whose __new__ will be used by default
608 first_enum: enumeration to check for an overriding __new__
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700609 """
610 # now find the correct __new__, checking to see of one was defined
611 # by the user; also check earlier enum classes in case a __new__ was
612 # saved as __new_member__
613 __new__ = classdict.get('__new__', None)
614
615 # should __new__ be saved as __new_member__ later?
616 save_new = __new__ is not None
617
618 if __new__ is None:
619 # check all possibles for __new_member__ before falling back to
620 # __new__
621 for method in ('__new_member__', '__new__'):
622 for possible in (member_type, first_enum):
623 target = getattr(possible, method, None)
624 if target not in {
625 None,
626 None.__new__,
627 object.__new__,
628 Enum.__new__,
629 }:
630 __new__ = target
631 break
632 if __new__ is not None:
633 break
634 else:
635 __new__ = object.__new__
636
637 # if a non-object.__new__ is used then whatever value/tuple was
638 # assigned to the enum member name will be passed to __new__ and to the
639 # new enum member's __init__
640 if __new__ is object.__new__:
641 use_args = False
642 else:
643 use_args = True
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700644 return __new__, save_new, use_args
645
646
647class Enum(metaclass=EnumMeta):
Ethan Furman6d3dfee2020-12-08 12:26:56 -0800648 """
649 Generic enumeration.
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700650
651 Derive from this class to define new enumerations.
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700652 """
653 def __new__(cls, value):
654 # all enum instances are actually created during class construction
655 # without calling this method; this method is called by the metaclass'
656 # __call__ (i.e. Color(3) ), and by pickle
657 if type(value) is cls:
Ethan Furman23bb6f42016-11-21 09:22:05 -0800658 # For lookups like Color(Color.RED)
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700659 return value
660 # by-value search for a matching enum member
661 # see if it's in the reverse mapping (for hashable values)
Ethan Furman2aa27322013-07-19 19:35:56 -0700662 try:
Andrew Svetlov34ae04f2018-12-26 20:45:33 +0200663 return cls._value2member_map_[value]
664 except KeyError:
665 # Not found, no need to do long O(n) search
666 pass
Ethan Furman2aa27322013-07-19 19:35:56 -0700667 except TypeError:
668 # not there, now do long search -- O(n) behavior
Ethan Furman520ad572013-07-19 19:47:21 -0700669 for member in cls._member_map_.values():
Ethan Furman0081f232014-09-16 17:31:23 -0700670 if member._value_ == value:
Ethan Furman2aa27322013-07-19 19:35:56 -0700671 return member
Ethan Furmanee47e5c2016-08-31 00:12:15 -0700672 # still not found -- try _missing_ hook
Ethan Furman019f0a02018-09-12 11:43:34 -0700673 try:
674 exc = None
675 result = cls._missing_(value)
676 except Exception as e:
677 exc = e
678 result = None
679 if isinstance(result, cls):
680 return result
681 else:
Walter Dörwald323842c2019-07-18 20:37:13 +0200682 ve_exc = ValueError("%r is not a valid %s" % (value, cls.__qualname__))
Ethan Furman019f0a02018-09-12 11:43:34 -0700683 if result is None and exc is None:
684 raise ve_exc
685 elif exc is None:
686 exc = TypeError(
687 'error in %s._missing_: returned %r instead of None or a valid member'
688 % (cls.__name__, result)
689 )
690 exc.__context__ = ve_exc
691 raise exc
Ethan Furmanee47e5c2016-08-31 00:12:15 -0700692
Ethan Furmanc16595e2016-09-10 23:36:59 -0700693 def _generate_next_value_(name, start, count, last_values):
Ethan Furman6d3dfee2020-12-08 12:26:56 -0800694 """
695 Generate the next value when not given.
696
697 name: the name of the member
698 start: the initial start value or None
699 count: the number of existing members
700 last_value: the last value assigned or None
701 """
Ethan Furmanc16595e2016-09-10 23:36:59 -0700702 for last_value in reversed(last_values):
703 try:
704 return last_value + 1
705 except TypeError:
706 pass
707 else:
Ethan Furmanee47e5c2016-08-31 00:12:15 -0700708 return start
Ethan Furmanc16595e2016-09-10 23:36:59 -0700709
Ethan Furman6bd94de2020-12-09 16:41:22 -0800710 def __init_subclass__(cls, **kwds):
711 super().__init_subclass__(**kwds)
712
Ethan Furmanee47e5c2016-08-31 00:12:15 -0700713 @classmethod
714 def _missing_(cls, value):
Ethan Furmanc95ad7a2020-09-16 10:26:50 -0700715 return None
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700716
717 def __repr__(self):
718 return "<%s.%s: %r>" % (
Ethan Furman520ad572013-07-19 19:47:21 -0700719 self.__class__.__name__, self._name_, self._value_)
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700720
721 def __str__(self):
Ethan Furman520ad572013-07-19 19:47:21 -0700722 return "%s.%s" % (self.__class__.__name__, self._name_)
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700723
Ethan Furman388a3922013-08-12 06:51:41 -0700724 def __dir__(self):
Ethan Furman6d3dfee2020-12-08 12:26:56 -0800725 """
726 Returns all members and all public methods
727 """
Ethan Furman0ae550b2014-10-14 08:58:32 -0700728 added_behavior = [
729 m
730 for cls in self.__class__.mro()
731 for m in cls.__dict__
Ethan Furman354ecf12015-03-11 08:43:12 -0700732 if m[0] != '_' and m not in self._member_map_
Angelin BOOZ68526fe2020-09-21 15:11:06 +0200733 ] + [m for m in self.__dict__ if m[0] != '_']
Ethan Furmanec5f8eb2014-10-21 13:40:35 -0700734 return (['__class__', '__doc__', '__module__'] + added_behavior)
Ethan Furman388a3922013-08-12 06:51:41 -0700735
Ethan Furmanec15a822013-08-31 19:17:41 -0700736 def __format__(self, format_spec):
Ethan Furman6d3dfee2020-12-08 12:26:56 -0800737 """
738 Returns format using actual value type unless __str__ has been overridden.
739 """
Ethan Furmanec15a822013-08-31 19:17:41 -0700740 # mixed-in Enums should use the mixed-in type's __format__, otherwise
741 # we can get strange results with the Enum name showing up instead of
742 # the value
743
thatneat2f19e822019-07-04 11:28:37 -0700744 # pure Enum branch, or branch with __str__ explicitly overridden
Ethan Furman37440ee2020-12-08 11:14:10 -0800745 str_overridden = type(self).__str__ not in (Enum.__str__, Flag.__str__)
thatneat2f19e822019-07-04 11:28:37 -0700746 if self._member_type_ is object or str_overridden:
Ethan Furmanec15a822013-08-31 19:17:41 -0700747 cls = str
748 val = str(self)
749 # mix-in branch
750 else:
751 cls = self._member_type_
Ethan Furman0081f232014-09-16 17:31:23 -0700752 val = self._value_
Ethan Furmanec15a822013-08-31 19:17:41 -0700753 return cls.__format__(val, format_spec)
754
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700755 def __hash__(self):
Ethan Furman520ad572013-07-19 19:47:21 -0700756 return hash(self._name_)
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700757
Ethan Furmanca1b7942014-02-08 11:36:27 -0800758 def __reduce_ex__(self, proto):
Ethan Furmandc870522014-02-18 12:37:12 -0800759 return self.__class__, (self._value_, )
Ethan Furmanca1b7942014-02-08 11:36:27 -0800760
Ethan Furman33918c12013-09-27 23:02:02 -0700761 # DynamicClassAttribute is used to provide access to the `name` and
762 # `value` properties of enum members while keeping some measure of
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700763 # protection from modification, while still allowing for an enumeration
764 # to have members named `name` and `value`. This works because enumeration
765 # members are not set directly on the enum class -- __getattr__ is
766 # used to look them up.
767
Ethan Furmane03ea372013-09-25 07:14:41 -0700768 @DynamicClassAttribute
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700769 def name(self):
Ethan Furmanc850f342013-09-15 16:59:35 -0700770 """The name of the Enum member."""
Ethan Furman520ad572013-07-19 19:47:21 -0700771 return self._name_
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700772
Ethan Furmane03ea372013-09-25 07:14:41 -0700773 @DynamicClassAttribute
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700774 def value(self):
Ethan Furmanc850f342013-09-15 16:59:35 -0700775 """The value of the Enum member."""
Ethan Furman520ad572013-07-19 19:47:21 -0700776 return self._value_
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700777
778
779class IntEnum(int, Enum):
Ethan Furman0063ff42020-09-21 17:23:13 -0700780 """
781 Enum where members are also (and must be) ints
782 """
783
784
785class StrEnum(str, Enum):
786 """
787 Enum where members are also (and must be) strings
788 """
789
790 def __new__(cls, *values):
791 if len(values) > 3:
792 raise TypeError('too many arguments for str(): %r' % (values, ))
793 if len(values) == 1:
794 # it must be a string
795 if not isinstance(values[0], str):
796 raise TypeError('%r is not a string' % (values[0], ))
797 if len(values) > 1:
798 # check that encoding argument is a string
799 if not isinstance(values[1], str):
800 raise TypeError('encoding must be a string, not %r' % (values[1], ))
801 if len(values) > 2:
802 # check that errors argument is a string
803 if not isinstance(values[2], str):
804 raise TypeError('errors must be a string, not %r' % (values[2], ))
805 value = str(*values)
806 member = str.__new__(cls, value)
807 member._value_ = value
808 return member
Ethan Furmanf24bb352013-07-18 17:05:39 -0700809
Ethan Furmand986d162020-09-22 13:00:07 -0700810 __str__ = str.__str__
811
Ethan Furmanf24bb352013-07-18 17:05:39 -0700812
Ethan Furman24e837f2015-03-18 17:27:57 -0700813def _reduce_ex_by_name(self, proto):
814 return self.name
815
Ethan Furman65a5a472016-09-01 23:55:19 -0700816class Flag(Enum):
Ethan Furman6d3dfee2020-12-08 12:26:56 -0800817 """
818 Support for flags
819 """
Ethan Furmanc16595e2016-09-10 23:36:59 -0700820
821 def _generate_next_value_(name, start, count, last_values):
Ethan Furmanee47e5c2016-08-31 00:12:15 -0700822 """
823 Generate the next value when not given.
824
825 name: the name of the member
HongWeipengbb16fb22019-09-21 13:22:54 +0800826 start: the initial start value or None
Ethan Furmanee47e5c2016-08-31 00:12:15 -0700827 count: the number of existing members
828 last_value: the last value assigned or None
829 """
830 if not count:
831 return start if start is not None else 1
Ethan Furmanc16595e2016-09-10 23:36:59 -0700832 for last_value in reversed(last_values):
833 try:
834 high_bit = _high_bit(last_value)
835 break
Ethan Furman3515dcc2016-09-18 13:15:41 -0700836 except Exception:
Ethan Furmanc16595e2016-09-10 23:36:59 -0700837 raise TypeError('Invalid Flag value: %r' % last_value) from None
Ethan Furmanee47e5c2016-08-31 00:12:15 -0700838 return 2 ** (high_bit+1)
839
840 @classmethod
841 def _missing_(cls, value):
Ethan Furman6d3dfee2020-12-08 12:26:56 -0800842 """
843 Returns member (possibly creating it) if one can be found for value.
844 """
Ethan Furmanee47e5c2016-08-31 00:12:15 -0700845 original_value = value
846 if value < 0:
847 value = ~value
848 possible_member = cls._create_pseudo_member_(value)
Ethan Furmanee47e5c2016-08-31 00:12:15 -0700849 if original_value < 0:
850 possible_member = ~possible_member
851 return possible_member
852
853 @classmethod
854 def _create_pseudo_member_(cls, value):
Ethan Furman3515dcc2016-09-18 13:15:41 -0700855 """
856 Create a composite member iff value contains only members.
857 """
Ethan Furmanee47e5c2016-08-31 00:12:15 -0700858 pseudo_member = cls._value2member_map_.get(value, None)
859 if pseudo_member is None:
Ethan Furman3515dcc2016-09-18 13:15:41 -0700860 # verify all bits are accounted for
861 _, extra_flags = _decompose(cls, value)
862 if extra_flags:
Walter Dörwald323842c2019-07-18 20:37:13 +0200863 raise ValueError("%r is not a valid %s" % (value, cls.__qualname__))
Ethan Furman3515dcc2016-09-18 13:15:41 -0700864 # construct a singleton enum pseudo-member
Ethan Furmanee47e5c2016-08-31 00:12:15 -0700865 pseudo_member = object.__new__(cls)
866 pseudo_member._name_ = None
867 pseudo_member._value_ = value
Ethan Furman28cf6632017-01-24 12:12:06 -0800868 # use setdefault in case another thread already created a composite
869 # with this value
870 pseudo_member = cls._value2member_map_.setdefault(value, pseudo_member)
Ethan Furmanee47e5c2016-08-31 00:12:15 -0700871 return pseudo_member
872
Ethan Furmanee47e5c2016-08-31 00:12:15 -0700873 def __contains__(self, other):
Ethan Furman6d3dfee2020-12-08 12:26:56 -0800874 """
875 Returns True if self has at least the same flags set as other.
876 """
Ethan Furmanee47e5c2016-08-31 00:12:15 -0700877 if not isinstance(other, self.__class__):
Rahul Jha94306522018-09-10 23:51:04 +0530878 raise TypeError(
879 "unsupported operand type(s) for 'in': '%s' and '%s'" % (
880 type(other).__qualname__, self.__class__.__qualname__))
Ethan Furmanee47e5c2016-08-31 00:12:15 -0700881 return other._value_ & self._value_ == other._value_
882
Ethan Furman7219e272020-09-16 13:01:00 -0700883 def __iter__(self):
Ethan Furman6d3dfee2020-12-08 12:26:56 -0800884 """
885 Returns flags in decreasing value order.
886 """
Ethan Furman7219e272020-09-16 13:01:00 -0700887 members, extra_flags = _decompose(self.__class__, self.value)
888 return (m for m in members if m._value_ != 0)
889
Ethan Furmanee47e5c2016-08-31 00:12:15 -0700890 def __repr__(self):
891 cls = self.__class__
892 if self._name_ is not None:
893 return '<%s.%s: %r>' % (cls.__name__, self._name_, self._value_)
Ethan Furman3515dcc2016-09-18 13:15:41 -0700894 members, uncovered = _decompose(cls, self._value_)
Ethan Furman27682d22016-09-04 11:39:01 -0700895 return '<%s.%s: %r>' % (
896 cls.__name__,
897 '|'.join([str(m._name_ or m._value_) for m in members]),
898 self._value_,
899 )
Ethan Furmanee47e5c2016-08-31 00:12:15 -0700900
901 def __str__(self):
902 cls = self.__class__
903 if self._name_ is not None:
904 return '%s.%s' % (cls.__name__, self._name_)
Ethan Furman3515dcc2016-09-18 13:15:41 -0700905 members, uncovered = _decompose(cls, self._value_)
Ethan Furmanee47e5c2016-08-31 00:12:15 -0700906 if len(members) == 1 and members[0]._name_ is None:
907 return '%s.%r' % (cls.__name__, members[0]._value_)
908 else:
909 return '%s.%s' % (
910 cls.__name__,
911 '|'.join([str(m._name_ or m._value_) for m in members]),
912 )
913
Ethan Furman25d94bb2016-09-02 16:32:32 -0700914 def __bool__(self):
915 return bool(self._value_)
916
Ethan Furmanee47e5c2016-08-31 00:12:15 -0700917 def __or__(self, other):
918 if not isinstance(other, self.__class__):
919 return NotImplemented
920 return self.__class__(self._value_ | other._value_)
921
922 def __and__(self, other):
923 if not isinstance(other, self.__class__):
924 return NotImplemented
925 return self.__class__(self._value_ & other._value_)
926
927 def __xor__(self, other):
928 if not isinstance(other, self.__class__):
929 return NotImplemented
930 return self.__class__(self._value_ ^ other._value_)
931
932 def __invert__(self):
Ethan Furman3515dcc2016-09-18 13:15:41 -0700933 members, uncovered = _decompose(self.__class__, self._value_)
Serhiy Storchaka81108372017-09-26 00:55:55 +0300934 inverted = self.__class__(0)
935 for m in self.__class__:
936 if m not in members and not (m._value_ & self._value_):
937 inverted = inverted | m
Ethan Furmanee47e5c2016-08-31 00:12:15 -0700938 return self.__class__(inverted)
939
940
Ethan Furman65a5a472016-09-01 23:55:19 -0700941class IntFlag(int, Flag):
Ethan Furman6d3dfee2020-12-08 12:26:56 -0800942 """
943 Support for integer-based Flags
944 """
Ethan Furmanee47e5c2016-08-31 00:12:15 -0700945
946 @classmethod
Ethan Furman3515dcc2016-09-18 13:15:41 -0700947 def _missing_(cls, value):
Ethan Furman6d3dfee2020-12-08 12:26:56 -0800948 """
949 Returns member (possibly creating it) if one can be found for value.
950 """
Ethan Furman3515dcc2016-09-18 13:15:41 -0700951 if not isinstance(value, int):
Walter Dörwald323842c2019-07-18 20:37:13 +0200952 raise ValueError("%r is not a valid %s" % (value, cls.__qualname__))
Ethan Furman3515dcc2016-09-18 13:15:41 -0700953 new_member = cls._create_pseudo_member_(value)
954 return new_member
955
956 @classmethod
Ethan Furmanee47e5c2016-08-31 00:12:15 -0700957 def _create_pseudo_member_(cls, value):
Ethan Furman6d3dfee2020-12-08 12:26:56 -0800958 """
959 Create a composite member iff value contains only members.
960 """
Ethan Furmanee47e5c2016-08-31 00:12:15 -0700961 pseudo_member = cls._value2member_map_.get(value, None)
962 if pseudo_member is None:
Ethan Furman3515dcc2016-09-18 13:15:41 -0700963 need_to_create = [value]
964 # get unaccounted for bits
965 _, extra_flags = _decompose(cls, value)
966 # timer = 10
967 while extra_flags:
968 # timer -= 1
969 bit = _high_bit(extra_flags)
970 flag_value = 2 ** bit
971 if (flag_value not in cls._value2member_map_ and
972 flag_value not in need_to_create
973 ):
974 need_to_create.append(flag_value)
975 if extra_flags == -flag_value:
976 extra_flags = 0
977 else:
978 extra_flags ^= flag_value
979 for value in reversed(need_to_create):
980 # construct singleton pseudo-members
981 pseudo_member = int.__new__(cls, value)
982 pseudo_member._name_ = None
983 pseudo_member._value_ = value
Ethan Furman28cf6632017-01-24 12:12:06 -0800984 # use setdefault in case another thread already created a composite
985 # with this value
986 pseudo_member = cls._value2member_map_.setdefault(value, pseudo_member)
Ethan Furmanee47e5c2016-08-31 00:12:15 -0700987 return pseudo_member
988
Ethan Furmanee47e5c2016-08-31 00:12:15 -0700989 def __or__(self, other):
990 if not isinstance(other, (self.__class__, int)):
991 return NotImplemented
Ethan Furman3515dcc2016-09-18 13:15:41 -0700992 result = self.__class__(self._value_ | self.__class__(other)._value_)
993 return result
Ethan Furmanee47e5c2016-08-31 00:12:15 -0700994
995 def __and__(self, other):
996 if not isinstance(other, (self.__class__, int)):
997 return NotImplemented
998 return self.__class__(self._value_ & self.__class__(other)._value_)
999
1000 def __xor__(self, other):
1001 if not isinstance(other, (self.__class__, int)):
1002 return NotImplemented
1003 return self.__class__(self._value_ ^ self.__class__(other)._value_)
1004
1005 __ror__ = __or__
1006 __rand__ = __and__
1007 __rxor__ = __xor__
1008
1009 def __invert__(self):
Ethan Furman3515dcc2016-09-18 13:15:41 -07001010 result = self.__class__(~self._value_)
1011 return result
Ethan Furmanee47e5c2016-08-31 00:12:15 -07001012
1013
1014def _high_bit(value):
Ethan Furman6d3dfee2020-12-08 12:26:56 -08001015 """
1016 returns index of highest bit, or -1 if value is zero or negative
1017 """
Ethan Furman3515dcc2016-09-18 13:15:41 -07001018 return value.bit_length() - 1
Ethan Furmanee47e5c2016-08-31 00:12:15 -07001019
Ethan Furmanf24bb352013-07-18 17:05:39 -07001020def unique(enumeration):
Ethan Furman6d3dfee2020-12-08 12:26:56 -08001021 """
1022 Class decorator for enumerations ensuring unique member values.
1023 """
Ethan Furmanf24bb352013-07-18 17:05:39 -07001024 duplicates = []
1025 for name, member in enumeration.__members__.items():
1026 if name != member.name:
1027 duplicates.append((name, member.name))
1028 if duplicates:
1029 alias_details = ', '.join(
1030 ["%s -> %s" % (alias, name) for (alias, name) in duplicates])
1031 raise ValueError('duplicate values found in %r: %s' %
1032 (enumeration, alias_details))
1033 return enumeration
Ethan Furman3515dcc2016-09-18 13:15:41 -07001034
1035def _decompose(flag, value):
Ethan Furman6d3dfee2020-12-08 12:26:56 -08001036 """
1037 Extract all members from the value.
1038 """
Ethan Furman3515dcc2016-09-18 13:15:41 -07001039 # _decompose is only called if the value is not named
1040 not_covered = value
1041 negative = value < 0
Ethan Furman3515dcc2016-09-18 13:15:41 -07001042 members = []
HongWeipeng0b41a922019-11-27 06:36:02 +08001043 for member in flag:
1044 member_value = member.value
Ethan Furman3515dcc2016-09-18 13:15:41 -07001045 if member_value and member_value & value == member_value:
1046 members.append(member)
1047 not_covered &= ~member_value
HongWeipeng0b41a922019-11-27 06:36:02 +08001048 if not negative:
1049 tmp = not_covered
1050 while tmp:
1051 flag_value = 2 ** _high_bit(tmp)
1052 if flag_value in flag._value2member_map_:
1053 members.append(flag._value2member_map_[flag_value])
1054 not_covered &= ~flag_value
1055 tmp &= ~flag_value
Ethan Furman3515dcc2016-09-18 13:15:41 -07001056 if not members and value in flag._value2member_map_:
1057 members.append(flag._value2member_map_[value])
1058 members.sort(key=lambda m: m._value_, reverse=True)
1059 if len(members) > 1 and members[0].value == value:
1060 # we have the breakdown, don't need the value member itself
1061 members.pop(0)
1062 return members, not_covered