blob: db79e66903ff4bfaf7e56ca2acda3749753d788c [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__ = [
Ethan Furman9bf7c2d2021-07-03 21:08:42 -07006 'EnumMeta',
7 'Enum', 'IntEnum', 'Flag', 'IntFlag',
8 'auto', 'unique',
Ethan Furmanc16595e2016-09-10 23:36:59 -07009 ]
Ethan Furman6b3d64a2013-06-14 16:55:46 -070010
11
Ethan Furman101e0742013-09-15 12:34:36 -070012def _is_descriptor(obj):
Ethan Furman6d3dfee2020-12-08 12:26:56 -080013 """
14 Returns True if obj is a descriptor, False otherwise.
15 """
Ethan Furman101e0742013-09-15 12:34:36 -070016 return (
17 hasattr(obj, '__get__') or
18 hasattr(obj, '__set__') or
Ethan Furman6d3dfee2020-12-08 12:26:56 -080019 hasattr(obj, '__delete__')
20 )
Ethan Furman101e0742013-09-15 12:34:36 -070021
Ethan Furman6b3d64a2013-06-14 16:55:46 -070022def _is_dunder(name):
Ethan Furman6d3dfee2020-12-08 12:26:56 -080023 """
24 Returns True if a __dunder__ name, False otherwise.
25 """
26 return (
27 len(name) > 4 and
Brennan D Baraban8b914d22019-03-03 14:09:11 -080028 name[:2] == name[-2:] == '__' and
29 name[2] != '_' and
Ethan Furman6d3dfee2020-12-08 12:26:56 -080030 name[-3] != '_'
31 )
Ethan Furman6b3d64a2013-06-14 16:55:46 -070032
33def _is_sunder(name):
Ethan Furman6d3dfee2020-12-08 12:26:56 -080034 """
35 Returns True if a _sunder_ name, False otherwise.
36 """
37 return (
38 len(name) > 2 and
Brennan D Baraban8b914d22019-03-03 14:09:11 -080039 name[0] == name[-1] == '_' and
Ethan Furman6b3d64a2013-06-14 16:55:46 -070040 name[1:2] != '_' and
Ethan Furman6d3dfee2020-12-08 12:26:56 -080041 name[-2:-1] != '_'
42 )
Ethan Furman6b3d64a2013-06-14 16:55:46 -070043
Ethan Furman7cf0aad2020-12-09 17:12:11 -080044def _is_private(cls_name, name):
45 # do not use `re` as `re` imports `enum`
46 pattern = '_%s__' % (cls_name, )
47 if (
Ethan Furman9bf7c2d2021-07-03 21:08:42 -070048 len(name) >= 5
Ethan Furman7cf0aad2020-12-09 17:12:11 -080049 and name.startswith(pattern)
Ethan Furman9bf7c2d2021-07-03 21:08:42 -070050 and name[len(pattern)] != '_'
Ethan Furman7cf0aad2020-12-09 17:12:11 -080051 and (name[-1] != '_' or name[-2] != '_')
52 ):
53 return True
54 else:
55 return False
56
Ethan Furman9bf7c2d2021-07-03 21:08:42 -070057def _make_class_unpicklable(cls):
Ethan Furman7aaeb2a2021-01-25 14:26:19 -080058 """
Ethan Furman9bf7c2d2021-07-03 21:08:42 -070059 Make the given class un-picklable.
Ethan Furman6d3dfee2020-12-08 12:26:56 -080060 """
Ethan Furmanca1b7942014-02-08 11:36:27 -080061 def _break_on_call_reduce(self, proto):
Ethan Furman6b3d64a2013-06-14 16:55:46 -070062 raise TypeError('%r cannot be pickled' % self)
Ethan Furman9bf7c2d2021-07-03 21:08:42 -070063 cls.__reduce_ex__ = _break_on_call_reduce
64 cls.__module__ = '<unknown>'
Ethan Furman7aaeb2a2021-01-25 14:26:19 -080065
Ethan Furman3515dcc2016-09-18 13:15:41 -070066_auto_null = object()
Ethan Furmanc16595e2016-09-10 23:36:59 -070067class auto:
68 """
69 Instances are replaced with an appropriate value in Enum class suites.
70 """
Ethan Furman3515dcc2016-09-18 13:15:41 -070071 value = _auto_null
Ethan Furmanc16595e2016-09-10 23:36:59 -070072
Ethan Furman101e0742013-09-15 12:34:36 -070073
Ethan Furman6b3d64a2013-06-14 16:55:46 -070074class _EnumDict(dict):
Ethan Furman6d3dfee2020-12-08 12:26:56 -080075 """
76 Track enum member order and ensure member names are not reused.
Ethan Furman6b3d64a2013-06-14 16:55:46 -070077
Ethan Furman9bf7c2d2021-07-03 21:08:42 -070078 EnumMeta will use the names found in self._member_names as the
Ethan Furman6b3d64a2013-06-14 16:55:46 -070079 enumeration member names.
Ethan Furman6b3d64a2013-06-14 16:55:46 -070080 """
81 def __init__(self):
82 super().__init__()
83 self._member_names = []
Ethan Furmanc16595e2016-09-10 23:36:59 -070084 self._last_values = []
Ethan Furmana4b1bb42018-01-22 07:56:37 -080085 self._ignore = []
Ethan Onstottd9a43e22020-04-28 13:20:55 -040086 self._auto_called = False
Ethan Furman6b3d64a2013-06-14 16:55:46 -070087
88 def __setitem__(self, key, value):
Ethan Furman6d3dfee2020-12-08 12:26:56 -080089 """
90 Changes anything not dundered or not a descriptor.
Ethan Furman6b3d64a2013-06-14 16:55:46 -070091
92 If an enum member name is used twice, an error is raised; duplicate
93 values are not checked for.
94
95 Single underscore (sunder) names are reserved.
Ethan Furman6b3d64a2013-06-14 16:55:46 -070096 """
Ethan Furman7cf0aad2020-12-09 17:12:11 -080097 if _is_private(self._cls_name, key):
Ethan Furman9bf7c2d2021-07-03 21:08:42 -070098 import warnings
99 warnings.warn(
100 "private variables, such as %r, will be normal attributes in 3.11"
101 % (key, ),
102 DeprecationWarning,
103 stacklevel=2,
104 )
105 if _is_sunder(key):
Ethan Furmanee47e5c2016-08-31 00:12:15 -0700106 if key not in (
Ethan Furman9bf7c2d2021-07-03 21:08:42 -0700107 '_order_', '_create_pseudo_member_',
Ethan Furmana4b1bb42018-01-22 07:56:37 -0800108 '_generate_next_value_', '_missing_', '_ignore_',
Ethan Furmanee47e5c2016-08-31 00:12:15 -0700109 ):
Ethan Furman9bf7c2d2021-07-03 21:08:42 -0700110 raise ValueError('_names_ are reserved for future Enum use')
Ethan Furmanc16595e2016-09-10 23:36:59 -0700111 if key == '_generate_next_value_':
Ethan Onstottd9a43e22020-04-28 13:20:55 -0400112 # check if members already defined as auto()
113 if self._auto_called:
114 raise TypeError("_generate_next_value_ must be defined before members")
Ethan Furman9bf7c2d2021-07-03 21:08:42 -0700115 setattr(self, '_generate_next_value', value)
Ethan Furmana4b1bb42018-01-22 07:56:37 -0800116 elif key == '_ignore_':
117 if isinstance(value, str):
118 value = value.replace(',',' ').split()
119 else:
120 value = list(value)
121 self._ignore = value
122 already = set(value) & set(self._member_names)
123 if already:
Ethan Furman6d3dfee2020-12-08 12:26:56 -0800124 raise ValueError(
125 '_ignore_ cannot specify already set names: %r'
126 % (already, )
127 )
Ethan Furman101e0742013-09-15 12:34:36 -0700128 elif _is_dunder(key):
Ethan Furmane8e61272016-08-20 07:19:31 -0700129 if key == '__order__':
130 key = '_order_'
Ethan Furman101e0742013-09-15 12:34:36 -0700131 elif key in self._member_names:
132 # descriptor overwriting an enum?
Ethan Furman9bf7c2d2021-07-03 21:08:42 -0700133 raise TypeError('Attempted to reuse key: %r' % key)
Ethan Furmana4b1bb42018-01-22 07:56:37 -0800134 elif key in self._ignore:
135 pass
Ethan Furman101e0742013-09-15 12:34:36 -0700136 elif not _is_descriptor(value):
137 if key in self:
138 # enum overwriting a descriptor?
Ethan Furmanee47e5c2016-08-31 00:12:15 -0700139 raise TypeError('%r already defined as: %r' % (key, self[key]))
Ethan Furmanc16595e2016-09-10 23:36:59 -0700140 if isinstance(value, auto):
Ethan Furman3515dcc2016-09-18 13:15:41 -0700141 if value.value == _auto_null:
Ethan Furman6d3dfee2020-12-08 12:26:56 -0800142 value.value = self._generate_next_value(
Ethan Furman9bf7c2d2021-07-03 21:08:42 -0700143 key,
144 1,
145 len(self._member_names),
146 self._last_values[:],
Ethan Furman6d3dfee2020-12-08 12:26:56 -0800147 )
Ethan Furmanfc23a942020-09-16 12:37:54 -0700148 self._auto_called = True
Ethan Furman3515dcc2016-09-18 13:15:41 -0700149 value = value.value
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700150 self._member_names.append(key)
Ethan Furmanc16595e2016-09-10 23:36:59 -0700151 self._last_values.append(value)
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700152 super().__setitem__(key, value)
153
Ethan Furmana6582872020-12-10 13:07:00 -0800154
Ethan Furman9bf7c2d2021-07-03 21:08:42 -0700155# Dummy value for Enum as EnumMeta explicitly checks for it, but of course
156# until EnumMeta finishes running the first time the Enum class doesn't exist.
157# This is also why there are checks in EnumMeta like `if Enum is not None`
158Enum = None
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700159
Ethan Furman9bf7c2d2021-07-03 21:08:42 -0700160class EnumMeta(type):
Ethan Furman6d3dfee2020-12-08 12:26:56 -0800161 """
162 Metaclass for Enum
163 """
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700164 @classmethod
Ethan Furman6ec0ade2020-12-24 10:05:02 -0800165 def __prepare__(metacls, cls, bases, **kwds):
Ethan Furman3064dbf2020-09-16 07:11:57 -0700166 # check that previous enum members do not exist
167 metacls._check_for_existing_members(cls, bases)
Ethan Furmanee47e5c2016-08-31 00:12:15 -0700168 # create the namespace dict
169 enum_dict = _EnumDict()
Ethan Furman7cf0aad2020-12-09 17:12:11 -0800170 enum_dict._cls_name = cls
Ethan Furmanee47e5c2016-08-31 00:12:15 -0700171 # inherit previous flags and _generate_next_value_ function
Ethan Furman3064dbf2020-09-16 07:11:57 -0700172 member_type, first_enum = metacls._get_mixins_(cls, bases)
Ethan Furmanee47e5c2016-08-31 00:12:15 -0700173 if first_enum is not None:
Ethan Furman6d3dfee2020-12-08 12:26:56 -0800174 enum_dict['_generate_next_value_'] = getattr(
175 first_enum, '_generate_next_value_', None,
176 )
Ethan Furmanee47e5c2016-08-31 00:12:15 -0700177 return enum_dict
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700178
Ethan Furman9bf7c2d2021-07-03 21:08:42 -0700179 def __new__(metacls, cls, bases, classdict, **kwds):
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700180 # an Enum class is final once enumeration items have been defined; it
181 # cannot be mixed with other types (int, float, etc.) if it has an
182 # inherited __new__ unless a new __new__ is defined (or the resulting
183 # class will fail).
Ethan Furmana4b1bb42018-01-22 07:56:37 -0800184 #
185 # remove any keys listed in _ignore_
186 classdict.setdefault('_ignore_', []).append('_ignore_')
187 ignore = classdict['_ignore_']
188 for key in ignore:
189 classdict.pop(key, None)
Ethan Furman3064dbf2020-09-16 07:11:57 -0700190 member_type, first_enum = metacls._get_mixins_(cls, bases)
Ethan Furmanc2667362020-12-07 00:17:31 -0800191 __new__, save_new, use_args = metacls._find_new_(
192 classdict, member_type, first_enum,
193 )
Ethan Furman9bf7c2d2021-07-03 21:08:42 -0700194
195 # save enum items into separate mapping so they don't get baked into
196 # the new class
197 enum_members = {k: classdict[k] for k in classdict._member_names}
198 for name in classdict._member_names:
199 del classdict[name]
200
201 # adjust the sunders
202 _order_ = classdict.pop('_order_', None)
203
204 # check for illegal enum names (any others?)
205 invalid_names = set(enum_members) & {'mro', ''}
206 if invalid_names:
207 raise ValueError('Invalid enum member name: {0}'.format(
208 ','.join(invalid_names)))
209
Ethan Furmanc314e602021-01-12 23:47:57 -0800210 # create a default docstring if one has not been provided
211 if '__doc__' not in classdict:
212 classdict['__doc__'] = 'An enumeration.'
Ethan Furman9bf7c2d2021-07-03 21:08:42 -0700213
214 enum_class = super().__new__(metacls, cls, bases, classdict, **kwds)
215 enum_class._member_names_ = [] # names in definition order
216 enum_class._member_map_ = {} # name->value map
217 enum_class._member_type_ = member_type
218
219 # save DynamicClassAttribute attributes from super classes so we know
220 # if we can take the shortcut of storing members in the class dict
221 dynamic_attributes = {
222 k for c in enum_class.mro()
223 for k, v in c.__dict__.items()
224 if isinstance(v, DynamicClassAttribute)
225 }
226
227 # Reverse value->name map for hashable values.
228 enum_class._value2member_map_ = {}
229
230 # If a custom type is mixed into the Enum, and it does not know how
231 # to pickle itself, pickle.dumps will succeed but pickle.loads will
232 # fail. Rather than have the error show up later and possibly far
233 # from the source, sabotage the pickle protocol for this class so
234 # that pickle.dumps also fails.
Ethan Furmanc314e602021-01-12 23:47:57 -0800235 #
Ethan Furman9bf7c2d2021-07-03 21:08:42 -0700236 # However, if the new class implements its own __reduce_ex__, do not
237 # sabotage -- it's on them to make sure it works correctly. We use
238 # __reduce_ex__ instead of any of the others as it is preferred by
239 # pickle over __reduce__, and it handles all pickle protocols.
240 if '__reduce_ex__' not in classdict:
241 if member_type is not object:
242 methods = ('__getnewargs_ex__', '__getnewargs__',
243 '__reduce_ex__', '__reduce__')
244 if not any(m in member_type.__dict__ for m in methods):
245 if '__new__' in classdict:
246 # too late, sabotage
247 _make_class_unpicklable(enum_class)
248 else:
249 # final attempt to verify that pickling would work:
250 # travel mro until __new__ is found, checking for
251 # __reduce__ and friends along the way -- if any of them
252 # are found before/when __new__ is found, pickling should
253 # work
254 sabotage = None
255 for chain in bases:
256 for base in chain.__mro__:
257 if base is object:
258 continue
259 elif any(m in base.__dict__ for m in methods):
260 # found one, we're good
261 sabotage = False
262 break
263 elif '__new__' in base.__dict__:
264 # not good
265 sabotage = True
266 break
267 if sabotage is not None:
268 break
269 if sabotage:
270 _make_class_unpicklable(enum_class)
271 # instantiate them, checking for duplicates as we go
272 # we instantiate first instead of checking for duplicates first in case
273 # a custom __new__ is doing something funky with the values -- such as
274 # auto-numbering ;)
275 for member_name in classdict._member_names:
276 value = enum_members[member_name]
277 if not isinstance(value, tuple):
278 args = (value, )
279 else:
280 args = value
281 if member_type is tuple: # special case for tuple enums
282 args = (args, ) # wrap it one more time
283 if not use_args:
284 enum_member = __new__(enum_class)
285 if not hasattr(enum_member, '_value_'):
286 enum_member._value_ = value
287 else:
288 enum_member = __new__(enum_class, *args)
289 if not hasattr(enum_member, '_value_'):
290 if member_type is object:
291 enum_member._value_ = value
292 else:
293 enum_member._value_ = member_type(*args)
294 value = enum_member._value_
295 enum_member._name_ = member_name
296 enum_member.__objclass__ = enum_class
297 enum_member.__init__(*args)
298 # If another member with the same value was already defined, the
299 # new member becomes an alias to the existing one.
300 for name, canonical_member in enum_class._member_map_.items():
301 if canonical_member._value_ == enum_member._value_:
302 enum_member = canonical_member
303 break
304 else:
305 # Aliases don't appear in member names (only in __members__).
306 enum_class._member_names_.append(member_name)
307 # performance boost for any member that would not shadow
308 # a DynamicClassAttribute
309 if member_name not in dynamic_attributes:
310 setattr(enum_class, member_name, enum_member)
311 # now add to _member_map_
312 enum_class._member_map_[member_name] = enum_member
313 try:
314 # This may fail if value is not hashable. We can't add the value
315 # to the map, and by-value lookups for this value will be
316 # linear.
317 enum_class._value2member_map_[value] = enum_member
318 except TypeError:
319 pass
320
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700321 # double check that repr and friends are not the mixin's or various
322 # things break (such as pickle)
Ethan Furman22415ad2020-09-15 16:28:25 -0700323 # however, if the method is defined in the Enum itself, don't replace
324 # it
Ethan Furmandc870522014-02-18 12:37:12 -0800325 for name in ('__repr__', '__str__', '__format__', '__reduce_ex__'):
Ethan Furman22415ad2020-09-15 16:28:25 -0700326 if name in classdict:
327 continue
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700328 class_method = getattr(enum_class, name)
329 obj_method = getattr(member_type, name, None)
330 enum_method = getattr(first_enum, name, None)
331 if obj_method is not None and obj_method is class_method:
332 setattr(enum_class, name, enum_method)
Ethan Furman9bf7c2d2021-07-03 21:08:42 -0700333
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700334 # replace any other __new__ with our own (as long as Enum is not None,
335 # anyway) -- again, this is to support pickle
336 if Enum is not None:
337 # if the user defined their own __new__, save it before it gets
338 # clobbered in case they subclass later
339 if save_new:
340 enum_class.__new_member__ = __new__
341 enum_class.__new__ = Enum.__new__
Ethan Furman9bf7c2d2021-07-03 21:08:42 -0700342
Ethan Furmane8e61272016-08-20 07:19:31 -0700343 # py3 support for definition order (helps keep py2/py3 code in sync)
344 if _order_ is not None:
345 if isinstance(_order_, str):
346 _order_ = _order_.replace(',', ' ').split()
347 if _order_ != enum_class._member_names_:
Ethan Furman9bf7c2d2021-07-03 21:08:42 -0700348 raise TypeError('member order does not match _order_')
349
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700350 return enum_class
351
Ethan Furman5de67b12016-04-13 23:52:09 -0700352 def __bool__(self):
353 """
354 classes/types should always be True.
355 """
356 return True
357
Ethan Furman9bf7c2d2021-07-03 21:08:42 -0700358 def __call__(cls, value, names=None, *, module=None, qualname=None, type=None, start=1):
Ethan Furman6d3dfee2020-12-08 12:26:56 -0800359 """
360 Either returns an existing member, or creates a new enum class.
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700361
362 This method is used both when an enum class is given a value to match
363 to an enumeration member (i.e. Color(3)) and for the functional API
Ethan Furman23bb6f42016-11-21 09:22:05 -0800364 (i.e. Color = Enum('Color', names='RED GREEN BLUE')).
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700365
Ethan Furman2da95042014-03-03 12:42:52 -0800366 When used for the functional API:
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700367
Ethan Furman2da95042014-03-03 12:42:52 -0800368 `value` will be the name of the new class.
369
370 `names` should be either a string of white-space/comma delimited names
Ethan Furmand9925a12014-09-16 20:35:55 -0700371 (values will start at `start`), or an iterator/mapping of name, value pairs.
Ethan Furman2da95042014-03-03 12:42:52 -0800372
373 `module` should be set to the module this class is being created in;
374 if it is not set, an attempt to find that module will be made, but if
375 it fails the class will not be picklable.
376
377 `qualname` should be set to the actual location this class can be found
378 at in its module; by default it is set to the global scope. If this is
379 not correct, unpickling will fail in some circumstances.
380
381 `type`, if set, will be mixed in as the first base class.
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700382 """
383 if names is None: # simple value lookup
384 return cls.__new__(cls, value)
385 # otherwise, functional API: we're creating a new Enum type
Ethan Furman6d3dfee2020-12-08 12:26:56 -0800386 return cls._create_(
387 value,
388 names,
389 module=module,
390 qualname=qualname,
391 type=type,
392 start=start,
393 )
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700394
395 def __contains__(cls, member):
Rahul Jha94306522018-09-10 23:51:04 +0530396 if not isinstance(member, Enum):
397 raise TypeError(
398 "unsupported operand type(s) for 'in': '%s' and '%s'" % (
399 type(member).__qualname__, cls.__class__.__qualname__))
Ethan Furman0081f232014-09-16 17:31:23 -0700400 return isinstance(member, cls) and member._name_ in cls._member_map_
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700401
Ethan Furman64a99722013-09-22 16:18:19 -0700402 def __delattr__(cls, attr):
403 # nicer error message when someone tries to delete an attribute
404 # (see issue19025).
405 if attr in cls._member_map_:
Ethan Furman9bf7c2d2021-07-03 21:08:42 -0700406 raise AttributeError("%s: cannot delete Enum member." % cls.__name__)
Ethan Furman64a99722013-09-22 16:18:19 -0700407 super().__delattr__(attr)
408
Ethan Furman388a3922013-08-12 06:51:41 -0700409 def __dir__(self):
Ethan Furman6d3dfee2020-12-08 12:26:56 -0800410 return (
411 ['__class__', '__doc__', '__members__', '__module__']
412 + self._member_names_
413 )
Ethan Furman388a3922013-08-12 06:51:41 -0700414
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700415 def __getattr__(cls, name):
Ethan Furman6d3dfee2020-12-08 12:26:56 -0800416 """
417 Return the enum member matching `name`
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700418
419 We use __getattr__ instead of descriptors or inserting into the enum
420 class' __dict__ in order to support `name` and `value` being both
421 properties for enum members (which live in the class' __dict__) and
422 enum members themselves.
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700423 """
424 if _is_dunder(name):
425 raise AttributeError(name)
426 try:
Ethan Furman520ad572013-07-19 19:47:21 -0700427 return cls._member_map_[name]
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700428 except KeyError:
429 raise AttributeError(name) from None
430
431 def __getitem__(cls, name):
Ethan Furman520ad572013-07-19 19:47:21 -0700432 return cls._member_map_[name]
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700433
434 def __iter__(cls):
Ethan Furman6d3dfee2020-12-08 12:26:56 -0800435 """
436 Returns members in definition order.
437 """
Ethan Furman520ad572013-07-19 19:47:21 -0700438 return (cls._member_map_[name] for name in cls._member_names_)
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700439
440 def __len__(cls):
Ethan Furman520ad572013-07-19 19:47:21 -0700441 return len(cls._member_names_)
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700442
Ethan Furman9bf7c2d2021-07-03 21:08:42 -0700443 @property
Ethan Furman2131a4a2013-09-14 18:11:24 -0700444 def __members__(cls):
Ethan Furman6d3dfee2020-12-08 12:26:56 -0800445 """
446 Returns a mapping of member name->value.
Ethan Furman2131a4a2013-09-14 18:11:24 -0700447
448 This mapping lists all enum members, including aliases. Note that this
449 is a read-only view of the internal mapping.
Ethan Furman2131a4a2013-09-14 18:11:24 -0700450 """
451 return MappingProxyType(cls._member_map_)
452
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700453 def __repr__(cls):
Ethan Furman9bf7c2d2021-07-03 21:08:42 -0700454 return "<enum %r>" % cls.__name__
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700455
Ethan Furman2131a4a2013-09-14 18:11:24 -0700456 def __reversed__(cls):
Ethan Furman6d3dfee2020-12-08 12:26:56 -0800457 """
458 Returns members in reverse definition order.
459 """
Ethan Furman2131a4a2013-09-14 18:11:24 -0700460 return (cls._member_map_[name] for name in reversed(cls._member_names_))
461
Ethan Furmanf203f2d2013-09-06 07:16:48 -0700462 def __setattr__(cls, name, value):
Ethan Furman6d3dfee2020-12-08 12:26:56 -0800463 """
464 Block attempts to reassign Enum members.
Ethan Furmanf203f2d2013-09-06 07:16:48 -0700465
466 A simple assignment to the class namespace only changes one of the
467 several possible ways to get an Enum member from the Enum class,
468 resulting in an inconsistent Enumeration.
Ethan Furmanf203f2d2013-09-06 07:16:48 -0700469 """
470 member_map = cls.__dict__.get('_member_map_', {})
471 if name in member_map:
Ethan Furman9bf7c2d2021-07-03 21:08:42 -0700472 raise AttributeError('Cannot reassign members.')
Ethan Furmanf203f2d2013-09-06 07:16:48 -0700473 super().__setattr__(name, value)
474
Ethan Furman9bf7c2d2021-07-03 21:08:42 -0700475 def _create_(cls, class_name, names, *, module=None, qualname=None, type=None, start=1):
Ethan Furman6d3dfee2020-12-08 12:26:56 -0800476 """
477 Convenience method to create a new Enum class.
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700478
479 `names` can be:
480
481 * A string containing member names, separated either with spaces or
Ethan Furmand9925a12014-09-16 20:35:55 -0700482 commas. Values are incremented by 1 from `start`.
483 * An iterable of member names. Values are incremented by 1 from `start`.
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700484 * An iterable of (member name, value) pairs.
Ethan Furmand9925a12014-09-16 20:35:55 -0700485 * A mapping of member name -> value pairs.
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700486 """
487 metacls = cls.__class__
488 bases = (cls, ) if type is None else (type, cls)
Ethan Furman3064dbf2020-09-16 07:11:57 -0700489 _, first_enum = cls._get_mixins_(cls, bases)
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700490 classdict = metacls.__prepare__(class_name, bases)
491
492 # special processing needed for names?
493 if isinstance(names, str):
494 names = names.replace(',', ' ').split()
Dong-hee Nadcc8ce42017-06-22 01:52:32 +0900495 if isinstance(names, (tuple, list)) and names and isinstance(names[0], str):
Ethan Furmanee47e5c2016-08-31 00:12:15 -0700496 original_names, names = names, []
Ethan Furmanc16595e2016-09-10 23:36:59 -0700497 last_values = []
Ethan Furmanee47e5c2016-08-31 00:12:15 -0700498 for count, name in enumerate(original_names):
Ethan Furmanc16595e2016-09-10 23:36:59 -0700499 value = first_enum._generate_next_value_(name, start, count, last_values[:])
500 last_values.append(value)
501 names.append((name, value))
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700502
503 # Here, names is either an iterable of (name, value) or a mapping.
504 for item in names:
505 if isinstance(item, str):
506 member_name, member_value = item, names[item]
507 else:
508 member_name, member_value = item
509 classdict[member_name] = member_value
Ethan Furman9bf7c2d2021-07-03 21:08:42 -0700510 enum_class = metacls.__new__(metacls, class_name, bases, classdict)
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700511
512 # TODO: replace the frame hack if a blessed way to know the calling
513 # module is ever developed
514 if module is None:
515 try:
516 module = sys._getframe(2).f_globals['__name__']
Pablo Galindo293dd232019-11-19 21:34:03 +0000517 except (AttributeError, ValueError, KeyError):
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700518 pass
519 if module is None:
Ethan Furman9bf7c2d2021-07-03 21:08:42 -0700520 _make_class_unpicklable(enum_class)
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700521 else:
Ethan Furman9bf7c2d2021-07-03 21:08:42 -0700522 enum_class.__module__ = module
Ethan Furmanca1b7942014-02-08 11:36:27 -0800523 if qualname is not None:
Ethan Furman9bf7c2d2021-07-03 21:08:42 -0700524 enum_class.__qualname__ = qualname
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700525
Ethan Furman9bf7c2d2021-07-03 21:08:42 -0700526 return enum_class
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700527
Ethan Furman9bf7c2d2021-07-03 21:08:42 -0700528 def _convert_(cls, name, module, filter, source=None):
orlnub1230fb9fad2018-09-12 20:28:53 +0300529 """
530 Create a new Enum subclass that replaces a collection of global constants
531 """
532 # convert all constants from source (or module) that pass filter() to
533 # a new Enum called name, and export the enum and its members back to
534 # module;
535 # also, replace the __reduce_ex__ method so unpickling works in
536 # previous Python versions
Ethan Furman9bf7c2d2021-07-03 21:08:42 -0700537 module_globals = vars(sys.modules[module])
orlnub1230fb9fad2018-09-12 20:28:53 +0300538 if source:
Ethan Furman9bf7c2d2021-07-03 21:08:42 -0700539 source = vars(source)
orlnub1230fb9fad2018-09-12 20:28:53 +0300540 else:
541 source = module_globals
542 # _value2member_map_ is populated in the same order every time
543 # for a consistent reverse mapping of number to name when there
544 # are multiple names for the same number.
545 members = [
546 (name, value)
547 for name, value in source.items()
548 if filter(name)]
549 try:
550 # sort by value
551 members.sort(key=lambda t: (t[1], t[0]))
552 except TypeError:
553 # unless some values aren't comparable, in which case sort by name
554 members.sort(key=lambda t: t[0])
Ethan Furman9bf7c2d2021-07-03 21:08:42 -0700555 cls = cls(name, members, module=module)
556 cls.__reduce_ex__ = _reduce_ex_by_name
557 module_globals.update(cls.__members__)
orlnub1230fb9fad2018-09-12 20:28:53 +0300558 module_globals[name] = cls
559 return cls
560
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700561 @staticmethod
Ethan Furman3064dbf2020-09-16 07:11:57 -0700562 def _check_for_existing_members(class_name, bases):
563 for chain in bases:
564 for base in chain.__mro__:
565 if issubclass(base, Enum) and base._member_names_:
Ethan Furman6d3dfee2020-12-08 12:26:56 -0800566 raise TypeError(
567 "%s: cannot extend enumeration %r"
568 % (class_name, base.__name__)
569 )
Ethan Furman3064dbf2020-09-16 07:11:57 -0700570
571 @staticmethod
572 def _get_mixins_(class_name, bases):
Ethan Furman6d3dfee2020-12-08 12:26:56 -0800573 """
574 Returns the type for creating enum members, and the first inherited
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700575 enum class.
576
577 bases: the tuple of bases that was given to __new__
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700578 """
579 if not bases:
580 return object, Enum
581
Ethan Furman5bdab642018-09-21 19:03:09 -0700582 def _find_data_type(bases):
Ethan Furman9bf7c2d2021-07-03 21:08:42 -0700583 data_types = []
Ethan Furman5bdab642018-09-21 19:03:09 -0700584 for chain in bases:
Ethan Furmanbff01f32020-09-15 15:56:26 -0700585 candidate = None
Ethan Furman5bdab642018-09-21 19:03:09 -0700586 for base in chain.__mro__:
587 if base is object:
588 continue
Ethan Furmanc2667362020-12-07 00:17:31 -0800589 elif issubclass(base, Enum):
590 if base._member_type_ is not object:
Ethan Furman9bf7c2d2021-07-03 21:08:42 -0700591 data_types.append(base._member_type_)
Ethan Furmanc2667362020-12-07 00:17:31 -0800592 break
Ethan Furman5bdab642018-09-21 19:03:09 -0700593 elif '__new__' in base.__dict__:
Ethan Furmancd453852018-10-05 23:29:36 -0700594 if issubclass(base, Enum):
Ethan Furman5bdab642018-09-21 19:03:09 -0700595 continue
Ethan Furman9bf7c2d2021-07-03 21:08:42 -0700596 data_types.append(candidate or base)
Ethan Furmanbff01f32020-09-15 15:56:26 -0700597 break
Ethan Furmanc2667362020-12-07 00:17:31 -0800598 else:
Miss Islington (bot)0f993242021-06-15 14:07:37 -0700599 candidate = candidate or base
Ethan Furmanbff01f32020-09-15 15:56:26 -0700600 if len(data_types) > 1:
Ethan Furman3064dbf2020-09-16 07:11:57 -0700601 raise TypeError('%r: too many data types: %r' % (class_name, data_types))
Ethan Furmanbff01f32020-09-15 15:56:26 -0700602 elif data_types:
Ethan Furman9bf7c2d2021-07-03 21:08:42 -0700603 return data_types[0]
Ethan Furmanbff01f32020-09-15 15:56:26 -0700604 else:
605 return None
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700606
Ethan Furman5bdab642018-09-21 19:03:09 -0700607 # ensure final parent class is an Enum derivative, find any concrete
608 # data type, and check that Enum has no members
609 first_enum = bases[-1]
610 if not issubclass(first_enum, Enum):
611 raise TypeError("new enumerations should be created as "
612 "`EnumName([mixin_type, ...] [data_type,] enum_type)`")
613 member_type = _find_data_type(bases) or object
614 if first_enum._member_names_:
615 raise TypeError("Cannot extend enumerations")
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700616 return member_type, first_enum
617
618 @staticmethod
619 def _find_new_(classdict, member_type, first_enum):
Ethan Furman6d3dfee2020-12-08 12:26:56 -0800620 """
621 Returns the __new__ to be used for creating the enum members.
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700622
623 classdict: the class dictionary given to __new__
624 member_type: the data type whose __new__ will be used by default
625 first_enum: enumeration to check for an overriding __new__
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700626 """
627 # now find the correct __new__, checking to see of one was defined
628 # by the user; also check earlier enum classes in case a __new__ was
629 # saved as __new_member__
630 __new__ = classdict.get('__new__', None)
631
632 # should __new__ be saved as __new_member__ later?
Ethan Furman9bf7c2d2021-07-03 21:08:42 -0700633 save_new = __new__ is not None
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700634
635 if __new__ is None:
636 # check all possibles for __new_member__ before falling back to
637 # __new__
638 for method in ('__new_member__', '__new__'):
639 for possible in (member_type, first_enum):
640 target = getattr(possible, method, None)
641 if target not in {
642 None,
643 None.__new__,
644 object.__new__,
645 Enum.__new__,
646 }:
647 __new__ = target
648 break
649 if __new__ is not None:
650 break
651 else:
652 __new__ = object.__new__
653
654 # if a non-object.__new__ is used then whatever value/tuple was
655 # assigned to the enum member name will be passed to __new__ and to the
656 # new enum member's __init__
Ethan Furman9bf7c2d2021-07-03 21:08:42 -0700657 if __new__ is object.__new__:
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700658 use_args = False
659 else:
660 use_args = True
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700661 return __new__, save_new, use_args
662
663
Ethan Furman9bf7c2d2021-07-03 21:08:42 -0700664class Enum(metaclass=EnumMeta):
Ethan Furman6d3dfee2020-12-08 12:26:56 -0800665 """
666 Generic enumeration.
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700667
668 Derive from this class to define new enumerations.
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700669 """
670 def __new__(cls, value):
671 # all enum instances are actually created during class construction
672 # without calling this method; this method is called by the metaclass'
673 # __call__ (i.e. Color(3) ), and by pickle
674 if type(value) is cls:
Ethan Furman23bb6f42016-11-21 09:22:05 -0800675 # For lookups like Color(Color.RED)
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700676 return value
677 # by-value search for a matching enum member
678 # see if it's in the reverse mapping (for hashable values)
Ethan Furman2aa27322013-07-19 19:35:56 -0700679 try:
Andrew Svetlov34ae04f2018-12-26 20:45:33 +0200680 return cls._value2member_map_[value]
681 except KeyError:
682 # Not found, no need to do long O(n) search
683 pass
Ethan Furman2aa27322013-07-19 19:35:56 -0700684 except TypeError:
685 # not there, now do long search -- O(n) behavior
Ethan Furman520ad572013-07-19 19:47:21 -0700686 for member in cls._member_map_.values():
Ethan Furman0081f232014-09-16 17:31:23 -0700687 if member._value_ == value:
Ethan Furman2aa27322013-07-19 19:35:56 -0700688 return member
Ethan Furmanee47e5c2016-08-31 00:12:15 -0700689 # still not found -- try _missing_ hook
Ethan Furman019f0a02018-09-12 11:43:34 -0700690 try:
691 exc = None
692 result = cls._missing_(value)
693 except Exception as e:
694 exc = e
695 result = None
Ethan Furman9bf7c2d2021-07-03 21:08:42 -0700696 if isinstance(result, cls):
697 return result
698 else:
699 ve_exc = ValueError("%r is not a valid %s" % (value, cls.__qualname__))
700 if result is None and exc is None:
701 raise ve_exc
702 elif exc is None:
703 exc = TypeError(
704 'error in %s._missing_: returned %r instead of None or a valid member'
705 % (cls.__name__, result)
706 )
707 exc.__context__ = ve_exc
708 raise exc
Ethan Furmanee47e5c2016-08-31 00:12:15 -0700709
Ethan Furmanc16595e2016-09-10 23:36:59 -0700710 def _generate_next_value_(name, start, count, last_values):
Ethan Furman6d3dfee2020-12-08 12:26:56 -0800711 """
712 Generate the next value when not given.
713
714 name: the name of the member
715 start: the initial start value or None
716 count: the number of existing members
717 last_value: the last value assigned or None
718 """
Ethan Furmanc16595e2016-09-10 23:36:59 -0700719 for last_value in reversed(last_values):
720 try:
721 return last_value + 1
722 except TypeError:
723 pass
724 else:
Ethan Furmanee47e5c2016-08-31 00:12:15 -0700725 return start
Ethan Furmanc16595e2016-09-10 23:36:59 -0700726
Ethan Furmanee47e5c2016-08-31 00:12:15 -0700727 @classmethod
728 def _missing_(cls, value):
Ethan Furmanc95ad7a2020-09-16 10:26:50 -0700729 return None
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700730
731 def __repr__(self):
Ethan Furman9bf7c2d2021-07-03 21:08:42 -0700732 return "<%s.%s: %r>" % (
733 self.__class__.__name__, self._name_, self._value_)
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700734
735 def __str__(self):
Ethan Furman9bf7c2d2021-07-03 21:08:42 -0700736 return "%s.%s" % (self.__class__.__name__, self._name_)
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700737
Ethan Furman388a3922013-08-12 06:51:41 -0700738 def __dir__(self):
Ethan Furman6d3dfee2020-12-08 12:26:56 -0800739 """
740 Returns all members and all public methods
741 """
Ethan Furman0ae550b2014-10-14 08:58:32 -0700742 added_behavior = [
743 m
744 for cls in self.__class__.mro()
745 for m in cls.__dict__
Ethan Furman354ecf12015-03-11 08:43:12 -0700746 if m[0] != '_' and m not in self._member_map_
Angelin BOOZ68526fe2020-09-21 15:11:06 +0200747 ] + [m for m in self.__dict__ if m[0] != '_']
Ethan Furmanec5f8eb2014-10-21 13:40:35 -0700748 return (['__class__', '__doc__', '__module__'] + added_behavior)
Ethan Furman388a3922013-08-12 06:51:41 -0700749
Ethan Furmanec15a822013-08-31 19:17:41 -0700750 def __format__(self, format_spec):
Ethan Furman6d3dfee2020-12-08 12:26:56 -0800751 """
752 Returns format using actual value type unless __str__ has been overridden.
753 """
Ethan Furmanec15a822013-08-31 19:17:41 -0700754 # mixed-in Enums should use the mixed-in type's __format__, otherwise
755 # we can get strange results with the Enum name showing up instead of
756 # the value
Ethan Furman9bf7c2d2021-07-03 21:08:42 -0700757
thatneat2f19e822019-07-04 11:28:37 -0700758 # pure Enum branch, or branch with __str__ explicitly overridden
Ethan Furman9bf7c2d2021-07-03 21:08:42 -0700759 str_overridden = type(self).__str__ not in (Enum.__str__, Flag.__str__)
thatneat2f19e822019-07-04 11:28:37 -0700760 if self._member_type_ is object or str_overridden:
Ethan Furmanec15a822013-08-31 19:17:41 -0700761 cls = str
762 val = str(self)
763 # mix-in branch
764 else:
765 cls = self._member_type_
Ethan Furman0081f232014-09-16 17:31:23 -0700766 val = self._value_
Ethan Furmanec15a822013-08-31 19:17:41 -0700767 return cls.__format__(val, format_spec)
768
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700769 def __hash__(self):
Ethan Furman520ad572013-07-19 19:47:21 -0700770 return hash(self._name_)
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700771
Ethan Furmanca1b7942014-02-08 11:36:27 -0800772 def __reduce_ex__(self, proto):
Ethan Furman9bf7c2d2021-07-03 21:08:42 -0700773 return self.__class__, (self._value_, )
Ethan Furmanca1b7942014-02-08 11:36:27 -0800774
Ethan Furman9bf7c2d2021-07-03 21:08:42 -0700775 # DynamicClassAttribute is used to provide access to the `name` and
776 # `value` properties of enum members while keeping some measure of
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700777 # protection from modification, while still allowing for an enumeration
778 # to have members named `name` and `value`. This works because enumeration
Ethan Furman9bf7c2d2021-07-03 21:08:42 -0700779 # members are not set directly on the enum class -- __getattr__ is
780 # used to look them up.
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700781
Ethan Furman9bf7c2d2021-07-03 21:08:42 -0700782 @DynamicClassAttribute
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700783 def name(self):
Ethan Furmanc850f342013-09-15 16:59:35 -0700784 """The name of the Enum member."""
Ethan Furman520ad572013-07-19 19:47:21 -0700785 return self._name_
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700786
Ethan Furman9bf7c2d2021-07-03 21:08:42 -0700787 @DynamicClassAttribute
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700788 def value(self):
Ethan Furmanc850f342013-09-15 16:59:35 -0700789 """The value of the Enum member."""
Ethan Furman520ad572013-07-19 19:47:21 -0700790 return self._value_
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700791
792
793class IntEnum(int, Enum):
Ethan Furman9bf7c2d2021-07-03 21:08:42 -0700794 """Enum where members are also (and must be) ints"""
Ethan Furman1b4addf2021-06-18 14:25:42 -0700795
Ethan Furman0063ff42020-09-21 17:23:13 -0700796
Ethan Furman9bf7c2d2021-07-03 21:08:42 -0700797def _reduce_ex_by_name(self, proto):
Ethan Furman24e837f2015-03-18 17:27:57 -0700798 return self.name
799
Ethan Furman9bf7c2d2021-07-03 21:08:42 -0700800class Flag(Enum):
Ethan Furman6d3dfee2020-12-08 12:26:56 -0800801 """
802 Support for flags
803 """
Ethan Furmanc16595e2016-09-10 23:36:59 -0700804
805 def _generate_next_value_(name, start, count, last_values):
Ethan Furmanee47e5c2016-08-31 00:12:15 -0700806 """
807 Generate the next value when not given.
808
809 name: the name of the member
HongWeipengbb16fb22019-09-21 13:22:54 +0800810 start: the initial start value or None
Ethan Furmanee47e5c2016-08-31 00:12:15 -0700811 count: the number of existing members
812 last_value: the last value assigned or None
813 """
814 if not count:
815 return start if start is not None else 1
Ethan Furman9bf7c2d2021-07-03 21:08:42 -0700816 for last_value in reversed(last_values):
817 try:
818 high_bit = _high_bit(last_value)
819 break
820 except Exception:
821 raise TypeError('Invalid Flag value: %r' % last_value) from None
Ethan Furmanee47e5c2016-08-31 00:12:15 -0700822 return 2 ** (high_bit+1)
823
824 @classmethod
825 def _missing_(cls, value):
Ethan Furman6d3dfee2020-12-08 12:26:56 -0800826 """
Ethan Furman9bf7c2d2021-07-03 21:08:42 -0700827 Returns member (possibly creating it) if one can be found for value.
Ethan Furman3515dcc2016-09-18 13:15:41 -0700828 """
Ethan Furman9bf7c2d2021-07-03 21:08:42 -0700829 original_value = value
Ethan Furman7aaeb2a2021-01-25 14:26:19 -0800830 if value < 0:
Ethan Furman9bf7c2d2021-07-03 21:08:42 -0700831 value = ~value
832 possible_member = cls._create_pseudo_member_(value)
833 if original_value < 0:
834 possible_member = ~possible_member
835 return possible_member
836
837 @classmethod
838 def _create_pseudo_member_(cls, value):
839 """
840 Create a composite member iff value contains only members.
841 """
842 pseudo_member = cls._value2member_map_.get(value, None)
843 if pseudo_member is None:
844 # verify all bits are accounted for
845 _, extra_flags = _decompose(cls, value)
846 if extra_flags:
847 raise ValueError("%r is not a valid %s" % (value, cls.__qualname__))
Ethan Furman3515dcc2016-09-18 13:15:41 -0700848 # construct a singleton enum pseudo-member
Ethan Furmanee47e5c2016-08-31 00:12:15 -0700849 pseudo_member = object.__new__(cls)
Ethan Furman7aaeb2a2021-01-25 14:26:19 -0800850 pseudo_member._name_ = None
Ethan Furman9bf7c2d2021-07-03 21:08:42 -0700851 pseudo_member._value_ = value
852 # use setdefault in case another thread already created a composite
853 # with this value
Ethan Furman28cf6632017-01-24 12:12:06 -0800854 pseudo_member = cls._value2member_map_.setdefault(value, pseudo_member)
Ethan Furmanee47e5c2016-08-31 00:12:15 -0700855 return pseudo_member
856
Ethan Furmanee47e5c2016-08-31 00:12:15 -0700857 def __contains__(self, other):
Ethan Furman6d3dfee2020-12-08 12:26:56 -0800858 """
859 Returns True if self has at least the same flags set as other.
860 """
Ethan Furmanee47e5c2016-08-31 00:12:15 -0700861 if not isinstance(other, self.__class__):
Rahul Jha94306522018-09-10 23:51:04 +0530862 raise TypeError(
863 "unsupported operand type(s) for 'in': '%s' and '%s'" % (
864 type(other).__qualname__, self.__class__.__qualname__))
Ethan Furmanee47e5c2016-08-31 00:12:15 -0700865 return other._value_ & self._value_ == other._value_
866
Ethan Furmanee47e5c2016-08-31 00:12:15 -0700867 def __repr__(self):
Ethan Furman9bf7c2d2021-07-03 21:08:42 -0700868 cls = self.__class__
869 if self._name_ is not None:
870 return '<%s.%s: %r>' % (cls.__name__, self._name_, self._value_)
871 members, uncovered = _decompose(cls, self._value_)
872 return '<%s.%s: %r>' % (
873 cls.__name__,
874 '|'.join([str(m._name_ or m._value_) for m in members]),
875 self._value_,
876 )
Ethan Furmanee47e5c2016-08-31 00:12:15 -0700877
878 def __str__(self):
879 cls = self.__class__
Ethan Furman9bf7c2d2021-07-03 21:08:42 -0700880 if self._name_ is not None:
881 return '%s.%s' % (cls.__name__, self._name_)
882 members, uncovered = _decompose(cls, self._value_)
883 if len(members) == 1 and members[0]._name_ is None:
884 return '%s.%r' % (cls.__name__, members[0]._value_)
Ethan Furmanee47e5c2016-08-31 00:12:15 -0700885 else:
Ethan Furman9bf7c2d2021-07-03 21:08:42 -0700886 return '%s.%s' % (
887 cls.__name__,
888 '|'.join([str(m._name_ or m._value_) for m in members]),
889 )
Ethan Furmanee47e5c2016-08-31 00:12:15 -0700890
Ethan Furman25d94bb2016-09-02 16:32:32 -0700891 def __bool__(self):
892 return bool(self._value_)
893
Ethan Furmanee47e5c2016-08-31 00:12:15 -0700894 def __or__(self, other):
895 if not isinstance(other, self.__class__):
896 return NotImplemented
897 return self.__class__(self._value_ | other._value_)
898
899 def __and__(self, other):
900 if not isinstance(other, self.__class__):
901 return NotImplemented
902 return self.__class__(self._value_ & other._value_)
903
904 def __xor__(self, other):
905 if not isinstance(other, self.__class__):
906 return NotImplemented
907 return self.__class__(self._value_ ^ other._value_)
908
909 def __invert__(self):
Ethan Furman9bf7c2d2021-07-03 21:08:42 -0700910 members, uncovered = _decompose(self.__class__, self._value_)
911 inverted = self.__class__(0)
912 for m in self.__class__:
913 if m not in members and not (m._value_ & self._value_):
914 inverted = inverted | m
915 return self.__class__(inverted)
Ethan Furmanee47e5c2016-08-31 00:12:15 -0700916
917
Ethan Furman9bf7c2d2021-07-03 21:08:42 -0700918class IntFlag(int, Flag):
Ethan Furman6d3dfee2020-12-08 12:26:56 -0800919 """
920 Support for integer-based Flags
921 """
Ethan Furmanee47e5c2016-08-31 00:12:15 -0700922
Ethan Furman9bf7c2d2021-07-03 21:08:42 -0700923 @classmethod
924 def _missing_(cls, value):
Ethan Furman1b4addf2021-06-18 14:25:42 -0700925 """
Ethan Furman9bf7c2d2021-07-03 21:08:42 -0700926 Returns member (possibly creating it) if one can be found for value.
Ethan Furman1b4addf2021-06-18 14:25:42 -0700927 """
Ethan Furman9bf7c2d2021-07-03 21:08:42 -0700928 if not isinstance(value, int):
929 raise ValueError("%r is not a valid %s" % (value, cls.__qualname__))
930 new_member = cls._create_pseudo_member_(value)
931 return new_member
932
933 @classmethod
934 def _create_pseudo_member_(cls, value):
935 """
936 Create a composite member iff value contains only members.
937 """
938 pseudo_member = cls._value2member_map_.get(value, None)
939 if pseudo_member is None:
940 need_to_create = [value]
941 # get unaccounted for bits
942 _, extra_flags = _decompose(cls, value)
943 # timer = 10
944 while extra_flags:
945 # timer -= 1
946 bit = _high_bit(extra_flags)
947 flag_value = 2 ** bit
948 if (flag_value not in cls._value2member_map_ and
949 flag_value not in need_to_create
950 ):
951 need_to_create.append(flag_value)
952 if extra_flags == -flag_value:
953 extra_flags = 0
954 else:
955 extra_flags ^= flag_value
956 for value in reversed(need_to_create):
957 # construct singleton pseudo-members
958 pseudo_member = int.__new__(cls, value)
959 pseudo_member._name_ = None
960 pseudo_member._value_ = value
961 # use setdefault in case another thread already created a composite
962 # with this value
963 pseudo_member = cls._value2member_map_.setdefault(value, pseudo_member)
964 return pseudo_member
Ethan Furman1b4addf2021-06-18 14:25:42 -0700965
Ethan Furmanee47e5c2016-08-31 00:12:15 -0700966 def __or__(self, other):
Ethan Furman9bf7c2d2021-07-03 21:08:42 -0700967 if not isinstance(other, (self.__class__, int)):
Ethan Furmanee47e5c2016-08-31 00:12:15 -0700968 return NotImplemented
Ethan Furman9bf7c2d2021-07-03 21:08:42 -0700969 result = self.__class__(self._value_ | self.__class__(other)._value_)
970 return result
Ethan Furmanee47e5c2016-08-31 00:12:15 -0700971
972 def __and__(self, other):
Ethan Furman9bf7c2d2021-07-03 21:08:42 -0700973 if not isinstance(other, (self.__class__, int)):
Ethan Furmanee47e5c2016-08-31 00:12:15 -0700974 return NotImplemented
Ethan Furman9bf7c2d2021-07-03 21:08:42 -0700975 return self.__class__(self._value_ & self.__class__(other)._value_)
Ethan Furmanee47e5c2016-08-31 00:12:15 -0700976
977 def __xor__(self, other):
Ethan Furman9bf7c2d2021-07-03 21:08:42 -0700978 if not isinstance(other, (self.__class__, int)):
Ethan Furmanee47e5c2016-08-31 00:12:15 -0700979 return NotImplemented
Ethan Furman9bf7c2d2021-07-03 21:08:42 -0700980 return self.__class__(self._value_ ^ self.__class__(other)._value_)
Ethan Furmanee47e5c2016-08-31 00:12:15 -0700981
982 __ror__ = __or__
983 __rand__ = __and__
984 __rxor__ = __xor__
Ethan Furman9bf7c2d2021-07-03 21:08:42 -0700985
986 def __invert__(self):
987 result = self.__class__(~self._value_)
988 return result
989
Ethan Furmanee47e5c2016-08-31 00:12:15 -0700990
991def _high_bit(value):
Ethan Furman6d3dfee2020-12-08 12:26:56 -0800992 """
993 returns index of highest bit, or -1 if value is zero or negative
994 """
Ethan Furman3515dcc2016-09-18 13:15:41 -0700995 return value.bit_length() - 1
Ethan Furmanee47e5c2016-08-31 00:12:15 -0700996
Ethan Furmanf24bb352013-07-18 17:05:39 -0700997def unique(enumeration):
Ethan Furman6d3dfee2020-12-08 12:26:56 -0800998 """
999 Class decorator for enumerations ensuring unique member values.
1000 """
Ethan Furmanf24bb352013-07-18 17:05:39 -07001001 duplicates = []
1002 for name, member in enumeration.__members__.items():
1003 if name != member.name:
1004 duplicates.append((name, member.name))
1005 if duplicates:
1006 alias_details = ', '.join(
1007 ["%s -> %s" % (alias, name) for (alias, name) in duplicates])
1008 raise ValueError('duplicate values found in %r: %s' %
1009 (enumeration, alias_details))
1010 return enumeration
Ethan Furman3515dcc2016-09-18 13:15:41 -07001011
Ethan Furman9bf7c2d2021-07-03 21:08:42 -07001012def _decompose(flag, value):
Ethan Furmanb7751062021-03-30 21:17:26 -07001013 """
Ethan Furman9bf7c2d2021-07-03 21:08:42 -07001014 Extract all members from the value.
Ethan Furmanb7751062021-03-30 21:17:26 -07001015 """
Ethan Furman9bf7c2d2021-07-03 21:08:42 -07001016 # _decompose is only called if the value is not named
1017 not_covered = value
1018 negative = value < 0
1019 members = []
1020 for member in flag:
1021 member_value = member.value
1022 if member_value and member_value & value == member_value:
1023 members.append(member)
1024 not_covered &= ~member_value
1025 if not negative:
1026 tmp = not_covered
1027 while tmp:
1028 flag_value = 2 ** _high_bit(tmp)
1029 if flag_value in flag._value2member_map_:
1030 members.append(flag._value2member_map_[flag_value])
1031 not_covered &= ~flag_value
1032 tmp &= ~flag_value
1033 if not members and value in flag._value2member_map_:
1034 members.append(flag._value2member_map_[value])
1035 members.sort(key=lambda m: m._value_, reverse=True)
1036 if len(members) > 1 and members[0].value == value:
1037 # we have the breakdown, don't need the value member itself
1038 members.pop(0)
1039 return members, not_covered