blob: a93642068f7437ff91200c7ecd988e255e46b384 [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 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 (
48 len(name) >= 5
49 and name.startswith(pattern)
50 and name[len(pattern)] != '_'
51 and (name[-1] != '_' or name[-2] != '_')
52 ):
53 return True
54 else:
55 return False
56
Ethan Furman6b3d64a2013-06-14 16:55:46 -070057def _make_class_unpicklable(cls):
Ethan Furman6d3dfee2020-12-08 12:26:56 -080058 """
59 Make the given class un-picklable.
60 """
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 Furmanca1b7942014-02-08 11:36:27 -080063 cls.__reduce_ex__ = _break_on_call_reduce
Ethan Furman6b3d64a2013-06-14 16:55:46 -070064 cls.__module__ = '<unknown>'
65
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
78 EnumMeta will use the names found in self._member_names as the
79 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):
98 # do nothing, name will be a normal attribute
99 pass
100 elif _is_sunder(key):
Ethan Furmanee47e5c2016-08-31 00:12:15 -0700101 if key not in (
Ethan Furman3515dcc2016-09-18 13:15:41 -0700102 '_order_', '_create_pseudo_member_',
Ethan Furmana4b1bb42018-01-22 07:56:37 -0800103 '_generate_next_value_', '_missing_', '_ignore_',
Ethan Furmanee47e5c2016-08-31 00:12:15 -0700104 ):
Ethan Furman6d3dfee2020-12-08 12:26:56 -0800105 raise ValueError(
106 '_sunder_ names, such as %r, are reserved for future Enum use'
107 % (key, )
108 )
Ethan Furmanc16595e2016-09-10 23:36:59 -0700109 if key == '_generate_next_value_':
Ethan Onstottd9a43e22020-04-28 13:20:55 -0400110 # check if members already defined as auto()
111 if self._auto_called:
112 raise TypeError("_generate_next_value_ must be defined before members")
Ethan Furmanc16595e2016-09-10 23:36:59 -0700113 setattr(self, '_generate_next_value', value)
Ethan Furmana4b1bb42018-01-22 07:56:37 -0800114 elif key == '_ignore_':
115 if isinstance(value, str):
116 value = value.replace(',',' ').split()
117 else:
118 value = list(value)
119 self._ignore = value
120 already = set(value) & set(self._member_names)
121 if already:
Ethan Furman6d3dfee2020-12-08 12:26:56 -0800122 raise ValueError(
123 '_ignore_ cannot specify already set names: %r'
124 % (already, )
125 )
Ethan Furman101e0742013-09-15 12:34:36 -0700126 elif _is_dunder(key):
Ethan Furmane8e61272016-08-20 07:19:31 -0700127 if key == '__order__':
128 key = '_order_'
Ethan Furman101e0742013-09-15 12:34:36 -0700129 elif key in self._member_names:
130 # descriptor overwriting an enum?
Ethan Furmana6582872020-12-10 13:07:00 -0800131 raise TypeError('%r already defined as: %r' % (key, self[key]))
Ethan Furmana4b1bb42018-01-22 07:56:37 -0800132 elif key in self._ignore:
133 pass
Ethan Furman101e0742013-09-15 12:34:36 -0700134 elif not _is_descriptor(value):
135 if key in self:
136 # enum overwriting a descriptor?
Ethan Furmanee47e5c2016-08-31 00:12:15 -0700137 raise TypeError('%r already defined as: %r' % (key, self[key]))
Ethan Furmanc16595e2016-09-10 23:36:59 -0700138 if isinstance(value, auto):
Ethan Furman3515dcc2016-09-18 13:15:41 -0700139 if value.value == _auto_null:
Ethan Furman6d3dfee2020-12-08 12:26:56 -0800140 value.value = self._generate_next_value(
141 key,
142 1,
143 len(self._member_names),
144 self._last_values[:],
145 )
Ethan Furmanfc23a942020-09-16 12:37:54 -0700146 self._auto_called = True
Ethan Furman3515dcc2016-09-18 13:15:41 -0700147 value = value.value
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700148 self._member_names.append(key)
Ethan Furmanc16595e2016-09-10 23:36:59 -0700149 self._last_values.append(value)
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700150 super().__setitem__(key, value)
151
Ethan Furmana6582872020-12-10 13:07:00 -0800152 def update(self, members, **more_members):
153 try:
154 for name in members.keys():
155 self[name] = members[name]
156 except AttributeError:
157 for name, value in members:
158 self[name] = value
159 for name, value in more_members.items():
160 self[name] = value
161
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700162
Ezio Melotti9a3777e2013-08-17 15:53:55 +0300163# Dummy value for Enum as EnumMeta explicitly checks for it, but of course
164# until EnumMeta finishes running the first time the Enum class doesn't exist.
165# This is also why there are checks in EnumMeta like `if Enum is not None`
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700166Enum = None
167
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700168class EnumMeta(type):
Ethan Furman6d3dfee2020-12-08 12:26:56 -0800169 """
170 Metaclass for Enum
171 """
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700172 @classmethod
Ethan Furman6ec0ade2020-12-24 10:05:02 -0800173 def __prepare__(metacls, cls, bases, **kwds):
Ethan Furman3064dbf2020-09-16 07:11:57 -0700174 # check that previous enum members do not exist
175 metacls._check_for_existing_members(cls, bases)
Ethan Furmanee47e5c2016-08-31 00:12:15 -0700176 # create the namespace dict
177 enum_dict = _EnumDict()
Ethan Furman7cf0aad2020-12-09 17:12:11 -0800178 enum_dict._cls_name = cls
Ethan Furmanee47e5c2016-08-31 00:12:15 -0700179 # inherit previous flags and _generate_next_value_ function
Ethan Furman3064dbf2020-09-16 07:11:57 -0700180 member_type, first_enum = metacls._get_mixins_(cls, bases)
Ethan Furmanee47e5c2016-08-31 00:12:15 -0700181 if first_enum is not None:
Ethan Furman6d3dfee2020-12-08 12:26:56 -0800182 enum_dict['_generate_next_value_'] = getattr(
183 first_enum, '_generate_next_value_', None,
184 )
Ethan Furmanee47e5c2016-08-31 00:12:15 -0700185 return enum_dict
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700186
Ethan Furman6bd94de2020-12-09 16:41:22 -0800187 def __new__(metacls, cls, bases, classdict, **kwds):
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700188 # an Enum class is final once enumeration items have been defined; it
189 # cannot be mixed with other types (int, float, etc.) if it has an
190 # inherited __new__ unless a new __new__ is defined (or the resulting
191 # class will fail).
Ethan Furmana4b1bb42018-01-22 07:56:37 -0800192 #
193 # remove any keys listed in _ignore_
194 classdict.setdefault('_ignore_', []).append('_ignore_')
195 ignore = classdict['_ignore_']
196 for key in ignore:
197 classdict.pop(key, None)
Ethan Furman3064dbf2020-09-16 07:11:57 -0700198 member_type, first_enum = metacls._get_mixins_(cls, bases)
Ethan Furmanc2667362020-12-07 00:17:31 -0800199 __new__, save_new, use_args = metacls._find_new_(
200 classdict, member_type, first_enum,
201 )
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700202
203 # save enum items into separate mapping so they don't get baked into
204 # the new class
Ethan Furmanee47e5c2016-08-31 00:12:15 -0700205 enum_members = {k: classdict[k] for k in classdict._member_names}
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700206 for name in classdict._member_names:
207 del classdict[name]
208
Ethan Furmane8e61272016-08-20 07:19:31 -0700209 # adjust the sunders
210 _order_ = classdict.pop('_order_', None)
211
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700212 # check for illegal enum names (any others?)
Brennan D Baraban8b914d22019-03-03 14:09:11 -0800213 invalid_names = set(enum_members) & {'mro', ''}
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700214 if invalid_names:
215 raise ValueError('Invalid enum member name: {0}'.format(
216 ','.join(invalid_names)))
217
Ethan Furman48a724f2015-04-11 23:23:06 -0700218 # create a default docstring if one has not been provided
219 if '__doc__' not in classdict:
220 classdict['__doc__'] = 'An enumeration.'
221
Ethan Furmana581a862021-01-07 13:17:55 -0800222 enum_class = super().__new__(metacls, cls, bases, classdict, **kwds)
Ethan Furman520ad572013-07-19 19:47:21 -0700223 enum_class._member_names_ = [] # names in definition order
INADA Naokie57f91a2018-06-19 01:14:26 +0900224 enum_class._member_map_ = {} # name->value map
Ethan Furman5e5a8232013-08-04 08:42:23 -0700225 enum_class._member_type_ = member_type
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700226
orlnub1230fb9fad2018-09-12 20:28:53 +0300227 # save DynamicClassAttribute attributes from super classes so we know
228 # if we can take the shortcut of storing members in the class dict
Ethan Furman6d3dfee2020-12-08 12:26:56 -0800229 dynamic_attributes = {
230 k for c in enum_class.mro()
231 for k, v in c.__dict__.items()
232 if isinstance(v, DynamicClassAttribute)
233 }
Ethan Furman354ecf12015-03-11 08:43:12 -0700234
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700235 # Reverse value->name map for hashable values.
Ethan Furman520ad572013-07-19 19:47:21 -0700236 enum_class._value2member_map_ = {}
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700237
Ethan Furman2da95042014-03-03 12:42:52 -0800238 # If a custom type is mixed into the Enum, and it does not know how
239 # to pickle itself, pickle.dumps will succeed but pickle.loads will
240 # fail. Rather than have the error show up later and possibly far
241 # from the source, sabotage the pickle protocol for this class so
242 # that pickle.dumps also fails.
243 #
244 # However, if the new class implements its own __reduce_ex__, do not
245 # sabotage -- it's on them to make sure it works correctly. We use
246 # __reduce_ex__ instead of any of the others as it is preferred by
247 # pickle over __reduce__, and it handles all pickle protocols.
248 if '__reduce_ex__' not in classdict:
Ethan Furmandc870522014-02-18 12:37:12 -0800249 if member_type is not object:
250 methods = ('__getnewargs_ex__', '__getnewargs__',
251 '__reduce_ex__', '__reduce__')
Ethan Furman2da95042014-03-03 12:42:52 -0800252 if not any(m in member_type.__dict__ for m in methods):
Ethan Furmandc870522014-02-18 12:37:12 -0800253 _make_class_unpicklable(enum_class)
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700254
255 # instantiate them, checking for duplicates as we go
256 # we instantiate first instead of checking for duplicates first in case
257 # a custom __new__ is doing something funky with the values -- such as
258 # auto-numbering ;)
259 for member_name in classdict._member_names:
Ethan Furmanee47e5c2016-08-31 00:12:15 -0700260 value = enum_members[member_name]
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700261 if not isinstance(value, tuple):
262 args = (value, )
263 else:
264 args = value
265 if member_type is tuple: # special case for tuple enums
266 args = (args, ) # wrap it one more time
267 if not use_args:
268 enum_member = __new__(enum_class)
Ethan Furmanb41803e2013-07-25 13:50:45 -0700269 if not hasattr(enum_member, '_value_'):
270 enum_member._value_ = value
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700271 else:
272 enum_member = __new__(enum_class, *args)
Ethan Furmanb41803e2013-07-25 13:50:45 -0700273 if not hasattr(enum_member, '_value_'):
Ethan Furmanee47e5c2016-08-31 00:12:15 -0700274 if member_type is object:
275 enum_member._value_ = value
276 else:
277 enum_member._value_ = member_type(*args)
Ethan Furman520ad572013-07-19 19:47:21 -0700278 value = enum_member._value_
Ethan Furman520ad572013-07-19 19:47:21 -0700279 enum_member._name_ = member_name
Ethan Furmanc850f342013-09-15 16:59:35 -0700280 enum_member.__objclass__ = enum_class
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700281 enum_member.__init__(*args)
282 # If another member with the same value was already defined, the
283 # new member becomes an alias to the existing one.
Ethan Furman520ad572013-07-19 19:47:21 -0700284 for name, canonical_member in enum_class._member_map_.items():
Ethan Furman0081f232014-09-16 17:31:23 -0700285 if canonical_member._value_ == enum_member._value_:
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700286 enum_member = canonical_member
287 break
288 else:
289 # Aliases don't appear in member names (only in __members__).
Ethan Furman520ad572013-07-19 19:47:21 -0700290 enum_class._member_names_.append(member_name)
Ethan Furman354ecf12015-03-11 08:43:12 -0700291 # performance boost for any member that would not shadow
292 # a DynamicClassAttribute
orlnub1230fb9fad2018-09-12 20:28:53 +0300293 if member_name not in dynamic_attributes:
Ethan Furman354ecf12015-03-11 08:43:12 -0700294 setattr(enum_class, member_name, enum_member)
295 # now add to _member_map_
Ethan Furman520ad572013-07-19 19:47:21 -0700296 enum_class._member_map_[member_name] = enum_member
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700297 try:
298 # This may fail if value is not hashable. We can't add the value
299 # to the map, and by-value lookups for this value will be
300 # linear.
Ethan Furman520ad572013-07-19 19:47:21 -0700301 enum_class._value2member_map_[value] = enum_member
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700302 except TypeError:
303 pass
304
305 # double check that repr and friends are not the mixin's or various
306 # things break (such as pickle)
Ethan Furman22415ad2020-09-15 16:28:25 -0700307 # however, if the method is defined in the Enum itself, don't replace
308 # it
Ethan Furmandc870522014-02-18 12:37:12 -0800309 for name in ('__repr__', '__str__', '__format__', '__reduce_ex__'):
Ethan Furman22415ad2020-09-15 16:28:25 -0700310 if name in classdict:
311 continue
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700312 class_method = getattr(enum_class, name)
313 obj_method = getattr(member_type, name, None)
314 enum_method = getattr(first_enum, name, None)
315 if obj_method is not None and obj_method is class_method:
316 setattr(enum_class, name, enum_method)
317
318 # replace any other __new__ with our own (as long as Enum is not None,
319 # anyway) -- again, this is to support pickle
320 if Enum is not None:
321 # if the user defined their own __new__, save it before it gets
322 # clobbered in case they subclass later
323 if save_new:
324 enum_class.__new_member__ = __new__
325 enum_class.__new__ = Enum.__new__
Ethan Furmane8e61272016-08-20 07:19:31 -0700326
327 # py3 support for definition order (helps keep py2/py3 code in sync)
328 if _order_ is not None:
329 if isinstance(_order_, str):
330 _order_ = _order_.replace(',', ' ').split()
331 if _order_ != enum_class._member_names_:
332 raise TypeError('member order does not match _order_')
333
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700334 return enum_class
335
Ethan Furman5de67b12016-04-13 23:52:09 -0700336 def __bool__(self):
337 """
338 classes/types should always be True.
339 """
340 return True
341
Ethan Furmand9925a12014-09-16 20:35:55 -0700342 def __call__(cls, value, names=None, *, module=None, qualname=None, type=None, start=1):
Ethan Furman6d3dfee2020-12-08 12:26:56 -0800343 """
344 Either returns an existing member, or creates a new enum class.
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700345
346 This method is used both when an enum class is given a value to match
347 to an enumeration member (i.e. Color(3)) and for the functional API
Ethan Furman23bb6f42016-11-21 09:22:05 -0800348 (i.e. Color = Enum('Color', names='RED GREEN BLUE')).
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700349
Ethan Furman2da95042014-03-03 12:42:52 -0800350 When used for the functional API:
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700351
Ethan Furman2da95042014-03-03 12:42:52 -0800352 `value` will be the name of the new class.
353
354 `names` should be either a string of white-space/comma delimited names
Ethan Furmand9925a12014-09-16 20:35:55 -0700355 (values will start at `start`), or an iterator/mapping of name, value pairs.
Ethan Furman2da95042014-03-03 12:42:52 -0800356
357 `module` should be set to the module this class is being created in;
358 if it is not set, an attempt to find that module will be made, but if
359 it fails the class will not be picklable.
360
361 `qualname` should be set to the actual location this class can be found
362 at in its module; by default it is set to the global scope. If this is
363 not correct, unpickling will fail in some circumstances.
364
365 `type`, if set, will be mixed in as the first base class.
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700366 """
367 if names is None: # simple value lookup
368 return cls.__new__(cls, value)
369 # otherwise, functional API: we're creating a new Enum type
Ethan Furman6d3dfee2020-12-08 12:26:56 -0800370 return cls._create_(
371 value,
372 names,
373 module=module,
374 qualname=qualname,
375 type=type,
376 start=start,
377 )
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700378
379 def __contains__(cls, member):
Rahul Jha94306522018-09-10 23:51:04 +0530380 if not isinstance(member, Enum):
381 raise TypeError(
382 "unsupported operand type(s) for 'in': '%s' and '%s'" % (
383 type(member).__qualname__, cls.__class__.__qualname__))
Ethan Furman0081f232014-09-16 17:31:23 -0700384 return isinstance(member, cls) and member._name_ in cls._member_map_
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700385
Ethan Furman64a99722013-09-22 16:18:19 -0700386 def __delattr__(cls, attr):
387 # nicer error message when someone tries to delete an attribute
388 # (see issue19025).
389 if attr in cls._member_map_:
Ethan Furman6d3dfee2020-12-08 12:26:56 -0800390 raise AttributeError("%s: cannot delete Enum member %r." % (cls.__name__, attr))
Ethan Furman64a99722013-09-22 16:18:19 -0700391 super().__delattr__(attr)
392
Ethan Furman388a3922013-08-12 06:51:41 -0700393 def __dir__(self):
Ethan Furman6d3dfee2020-12-08 12:26:56 -0800394 return (
395 ['__class__', '__doc__', '__members__', '__module__']
396 + self._member_names_
397 )
Ethan Furman388a3922013-08-12 06:51:41 -0700398
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700399 def __getattr__(cls, name):
Ethan Furman6d3dfee2020-12-08 12:26:56 -0800400 """
401 Return the enum member matching `name`
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700402
403 We use __getattr__ instead of descriptors or inserting into the enum
404 class' __dict__ in order to support `name` and `value` being both
405 properties for enum members (which live in the class' __dict__) and
406 enum members themselves.
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700407 """
408 if _is_dunder(name):
409 raise AttributeError(name)
410 try:
Ethan Furman520ad572013-07-19 19:47:21 -0700411 return cls._member_map_[name]
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700412 except KeyError:
413 raise AttributeError(name) from None
414
415 def __getitem__(cls, name):
Ethan Furman520ad572013-07-19 19:47:21 -0700416 return cls._member_map_[name]
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700417
418 def __iter__(cls):
Ethan Furman6d3dfee2020-12-08 12:26:56 -0800419 """
420 Returns members in definition order.
421 """
Ethan Furman520ad572013-07-19 19:47:21 -0700422 return (cls._member_map_[name] for name in cls._member_names_)
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700423
424 def __len__(cls):
Ethan Furman520ad572013-07-19 19:47:21 -0700425 return len(cls._member_names_)
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700426
Ethan Furman2131a4a2013-09-14 18:11:24 -0700427 @property
428 def __members__(cls):
Ethan Furman6d3dfee2020-12-08 12:26:56 -0800429 """
430 Returns a mapping of member name->value.
Ethan Furman2131a4a2013-09-14 18:11:24 -0700431
432 This mapping lists all enum members, including aliases. Note that this
433 is a read-only view of the internal mapping.
Ethan Furman2131a4a2013-09-14 18:11:24 -0700434 """
435 return MappingProxyType(cls._member_map_)
436
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700437 def __repr__(cls):
438 return "<enum %r>" % cls.__name__
439
Ethan Furman2131a4a2013-09-14 18:11:24 -0700440 def __reversed__(cls):
Ethan Furman6d3dfee2020-12-08 12:26:56 -0800441 """
442 Returns members in reverse definition order.
443 """
Ethan Furman2131a4a2013-09-14 18:11:24 -0700444 return (cls._member_map_[name] for name in reversed(cls._member_names_))
445
Ethan Furmanf203f2d2013-09-06 07:16:48 -0700446 def __setattr__(cls, name, value):
Ethan Furman6d3dfee2020-12-08 12:26:56 -0800447 """
448 Block attempts to reassign Enum members.
Ethan Furmanf203f2d2013-09-06 07:16:48 -0700449
450 A simple assignment to the class namespace only changes one of the
451 several possible ways to get an Enum member from the Enum class,
452 resulting in an inconsistent Enumeration.
Ethan Furmanf203f2d2013-09-06 07:16:48 -0700453 """
454 member_map = cls.__dict__.get('_member_map_', {})
455 if name in member_map:
456 raise AttributeError('Cannot reassign members.')
457 super().__setattr__(name, value)
458
anentropicb8e21f12018-04-16 04:40:35 +0100459 def _create_(cls, class_name, names, *, module=None, qualname=None, type=None, start=1):
Ethan Furman6d3dfee2020-12-08 12:26:56 -0800460 """
461 Convenience method to create a new Enum class.
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700462
463 `names` can be:
464
465 * A string containing member names, separated either with spaces or
Ethan Furmand9925a12014-09-16 20:35:55 -0700466 commas. Values are incremented by 1 from `start`.
467 * An iterable of member names. Values are incremented by 1 from `start`.
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700468 * An iterable of (member name, value) pairs.
Ethan Furmand9925a12014-09-16 20:35:55 -0700469 * A mapping of member name -> value pairs.
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700470 """
471 metacls = cls.__class__
472 bases = (cls, ) if type is None else (type, cls)
Ethan Furman3064dbf2020-09-16 07:11:57 -0700473 _, first_enum = cls._get_mixins_(cls, bases)
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700474 classdict = metacls.__prepare__(class_name, bases)
475
476 # special processing needed for names?
477 if isinstance(names, str):
478 names = names.replace(',', ' ').split()
Dong-hee Nadcc8ce42017-06-22 01:52:32 +0900479 if isinstance(names, (tuple, list)) and names and isinstance(names[0], str):
Ethan Furmanee47e5c2016-08-31 00:12:15 -0700480 original_names, names = names, []
Ethan Furmanc16595e2016-09-10 23:36:59 -0700481 last_values = []
Ethan Furmanee47e5c2016-08-31 00:12:15 -0700482 for count, name in enumerate(original_names):
Ethan Furmanc16595e2016-09-10 23:36:59 -0700483 value = first_enum._generate_next_value_(name, start, count, last_values[:])
484 last_values.append(value)
485 names.append((name, value))
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700486
487 # Here, names is either an iterable of (name, value) or a mapping.
488 for item in names:
489 if isinstance(item, str):
490 member_name, member_value = item, names[item]
491 else:
492 member_name, member_value = item
493 classdict[member_name] = member_value
494 enum_class = metacls.__new__(metacls, class_name, bases, classdict)
495
496 # TODO: replace the frame hack if a blessed way to know the calling
497 # module is ever developed
498 if module is None:
499 try:
500 module = sys._getframe(2).f_globals['__name__']
Pablo Galindo293dd232019-11-19 21:34:03 +0000501 except (AttributeError, ValueError, KeyError):
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700502 pass
503 if module is None:
504 _make_class_unpicklable(enum_class)
505 else:
506 enum_class.__module__ = module
Ethan Furmanca1b7942014-02-08 11:36:27 -0800507 if qualname is not None:
508 enum_class.__qualname__ = qualname
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700509
510 return enum_class
511
orlnub1230fb9fad2018-09-12 20:28:53 +0300512 def _convert_(cls, name, module, filter, source=None):
513 """
514 Create a new Enum subclass that replaces a collection of global constants
515 """
516 # convert all constants from source (or module) that pass filter() to
517 # a new Enum called name, and export the enum and its members back to
518 # module;
519 # also, replace the __reduce_ex__ method so unpickling works in
520 # previous Python versions
521 module_globals = vars(sys.modules[module])
522 if source:
523 source = vars(source)
524 else:
525 source = module_globals
526 # _value2member_map_ is populated in the same order every time
527 # for a consistent reverse mapping of number to name when there
528 # are multiple names for the same number.
529 members = [
530 (name, value)
531 for name, value in source.items()
532 if filter(name)]
533 try:
534 # sort by value
535 members.sort(key=lambda t: (t[1], t[0]))
536 except TypeError:
537 # unless some values aren't comparable, in which case sort by name
538 members.sort(key=lambda t: t[0])
539 cls = cls(name, members, module=module)
540 cls.__reduce_ex__ = _reduce_ex_by_name
541 module_globals.update(cls.__members__)
542 module_globals[name] = cls
543 return cls
544
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700545 @staticmethod
Ethan Furman3064dbf2020-09-16 07:11:57 -0700546 def _check_for_existing_members(class_name, bases):
547 for chain in bases:
548 for base in chain.__mro__:
549 if issubclass(base, Enum) and base._member_names_:
Ethan Furman6d3dfee2020-12-08 12:26:56 -0800550 raise TypeError(
551 "%s: cannot extend enumeration %r"
552 % (class_name, base.__name__)
553 )
Ethan Furman3064dbf2020-09-16 07:11:57 -0700554
555 @staticmethod
556 def _get_mixins_(class_name, bases):
Ethan Furman6d3dfee2020-12-08 12:26:56 -0800557 """
558 Returns the type for creating enum members, and the first inherited
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700559 enum class.
560
561 bases: the tuple of bases that was given to __new__
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700562 """
563 if not bases:
564 return object, Enum
565
Ethan Furman5bdab642018-09-21 19:03:09 -0700566 def _find_data_type(bases):
Ethan Furmanbff01f32020-09-15 15:56:26 -0700567 data_types = []
Ethan Furman5bdab642018-09-21 19:03:09 -0700568 for chain in bases:
Ethan Furmanbff01f32020-09-15 15:56:26 -0700569 candidate = None
Ethan Furman5bdab642018-09-21 19:03:09 -0700570 for base in chain.__mro__:
571 if base is object:
572 continue
Ethan Furmanc2667362020-12-07 00:17:31 -0800573 elif issubclass(base, Enum):
574 if base._member_type_ is not object:
575 data_types.append(base._member_type_)
576 break
Ethan Furman5bdab642018-09-21 19:03:09 -0700577 elif '__new__' in base.__dict__:
Ethan Furmancd453852018-10-05 23:29:36 -0700578 if issubclass(base, Enum):
Ethan Furman5bdab642018-09-21 19:03:09 -0700579 continue
Ethan Furmanbff01f32020-09-15 15:56:26 -0700580 data_types.append(candidate or base)
581 break
Ethan Furmanc2667362020-12-07 00:17:31 -0800582 else:
Ethan Furmanbff01f32020-09-15 15:56:26 -0700583 candidate = base
584 if len(data_types) > 1:
Ethan Furman3064dbf2020-09-16 07:11:57 -0700585 raise TypeError('%r: too many data types: %r' % (class_name, data_types))
Ethan Furmanbff01f32020-09-15 15:56:26 -0700586 elif data_types:
587 return data_types[0]
588 else:
589 return None
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700590
Ethan Furman5bdab642018-09-21 19:03:09 -0700591 # ensure final parent class is an Enum derivative, find any concrete
592 # data type, and check that Enum has no members
593 first_enum = bases[-1]
594 if not issubclass(first_enum, Enum):
595 raise TypeError("new enumerations should be created as "
596 "`EnumName([mixin_type, ...] [data_type,] enum_type)`")
597 member_type = _find_data_type(bases) or object
598 if first_enum._member_names_:
599 raise TypeError("Cannot extend enumerations")
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700600 return member_type, first_enum
601
602 @staticmethod
603 def _find_new_(classdict, member_type, first_enum):
Ethan Furman6d3dfee2020-12-08 12:26:56 -0800604 """
605 Returns the __new__ to be used for creating the enum members.
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700606
607 classdict: the class dictionary given to __new__
608 member_type: the data type whose __new__ will be used by default
609 first_enum: enumeration to check for an overriding __new__
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700610 """
611 # now find the correct __new__, checking to see of one was defined
612 # by the user; also check earlier enum classes in case a __new__ was
613 # saved as __new_member__
614 __new__ = classdict.get('__new__', None)
615
616 # should __new__ be saved as __new_member__ later?
617 save_new = __new__ is not None
618
619 if __new__ is None:
620 # check all possibles for __new_member__ before falling back to
621 # __new__
622 for method in ('__new_member__', '__new__'):
623 for possible in (member_type, first_enum):
624 target = getattr(possible, method, None)
625 if target not in {
626 None,
627 None.__new__,
628 object.__new__,
629 Enum.__new__,
630 }:
631 __new__ = target
632 break
633 if __new__ is not None:
634 break
635 else:
636 __new__ = object.__new__
637
638 # if a non-object.__new__ is used then whatever value/tuple was
639 # assigned to the enum member name will be passed to __new__ and to the
640 # new enum member's __init__
641 if __new__ is object.__new__:
642 use_args = False
643 else:
644 use_args = True
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700645 return __new__, save_new, use_args
646
647
648class Enum(metaclass=EnumMeta):
Ethan Furman6d3dfee2020-12-08 12:26:56 -0800649 """
650 Generic enumeration.
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700651
652 Derive from this class to define new enumerations.
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700653 """
654 def __new__(cls, value):
655 # all enum instances are actually created during class construction
656 # without calling this method; this method is called by the metaclass'
657 # __call__ (i.e. Color(3) ), and by pickle
658 if type(value) is cls:
Ethan Furman23bb6f42016-11-21 09:22:05 -0800659 # For lookups like Color(Color.RED)
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700660 return value
661 # by-value search for a matching enum member
662 # see if it's in the reverse mapping (for hashable values)
Ethan Furman2aa27322013-07-19 19:35:56 -0700663 try:
Andrew Svetlov34ae04f2018-12-26 20:45:33 +0200664 return cls._value2member_map_[value]
665 except KeyError:
666 # Not found, no need to do long O(n) search
667 pass
Ethan Furman2aa27322013-07-19 19:35:56 -0700668 except TypeError:
669 # not there, now do long search -- O(n) behavior
Ethan Furman520ad572013-07-19 19:47:21 -0700670 for member in cls._member_map_.values():
Ethan Furman0081f232014-09-16 17:31:23 -0700671 if member._value_ == value:
Ethan Furman2aa27322013-07-19 19:35:56 -0700672 return member
Ethan Furmanee47e5c2016-08-31 00:12:15 -0700673 # still not found -- try _missing_ hook
Ethan Furman019f0a02018-09-12 11:43:34 -0700674 try:
675 exc = None
676 result = cls._missing_(value)
677 except Exception as e:
678 exc = e
679 result = None
680 if isinstance(result, cls):
681 return result
682 else:
Walter Dörwald323842c2019-07-18 20:37:13 +0200683 ve_exc = ValueError("%r is not a valid %s" % (value, cls.__qualname__))
Ethan Furman019f0a02018-09-12 11:43:34 -0700684 if result is None and exc is None:
685 raise ve_exc
686 elif exc is None:
687 exc = TypeError(
688 'error in %s._missing_: returned %r instead of None or a valid member'
689 % (cls.__name__, result)
690 )
691 exc.__context__ = ve_exc
692 raise exc
Ethan Furmanee47e5c2016-08-31 00:12:15 -0700693
Ethan Furmanc16595e2016-09-10 23:36:59 -0700694 def _generate_next_value_(name, start, count, last_values):
Ethan Furman6d3dfee2020-12-08 12:26:56 -0800695 """
696 Generate the next value when not given.
697
698 name: the name of the member
699 start: the initial start value or None
700 count: the number of existing members
701 last_value: the last value assigned or None
702 """
Ethan Furmanc16595e2016-09-10 23:36:59 -0700703 for last_value in reversed(last_values):
704 try:
705 return last_value + 1
706 except TypeError:
707 pass
708 else:
Ethan Furmanee47e5c2016-08-31 00:12:15 -0700709 return start
Ethan Furmanc16595e2016-09-10 23:36:59 -0700710
Ethan Furmanee47e5c2016-08-31 00:12:15 -0700711 @classmethod
712 def _missing_(cls, value):
Ethan Furmanc95ad7a2020-09-16 10:26:50 -0700713 return None
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700714
715 def __repr__(self):
716 return "<%s.%s: %r>" % (
Ethan Furman520ad572013-07-19 19:47:21 -0700717 self.__class__.__name__, self._name_, self._value_)
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700718
719 def __str__(self):
Ethan Furman520ad572013-07-19 19:47:21 -0700720 return "%s.%s" % (self.__class__.__name__, self._name_)
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700721
Ethan Furman388a3922013-08-12 06:51:41 -0700722 def __dir__(self):
Ethan Furman6d3dfee2020-12-08 12:26:56 -0800723 """
724 Returns all members and all public methods
725 """
Ethan Furman0ae550b2014-10-14 08:58:32 -0700726 added_behavior = [
727 m
728 for cls in self.__class__.mro()
729 for m in cls.__dict__
Ethan Furman354ecf12015-03-11 08:43:12 -0700730 if m[0] != '_' and m not in self._member_map_
Angelin BOOZ68526fe2020-09-21 15:11:06 +0200731 ] + [m for m in self.__dict__ if m[0] != '_']
Ethan Furmanec5f8eb2014-10-21 13:40:35 -0700732 return (['__class__', '__doc__', '__module__'] + added_behavior)
Ethan Furman388a3922013-08-12 06:51:41 -0700733
Ethan Furmanec15a822013-08-31 19:17:41 -0700734 def __format__(self, format_spec):
Ethan Furman6d3dfee2020-12-08 12:26:56 -0800735 """
736 Returns format using actual value type unless __str__ has been overridden.
737 """
Ethan Furmanec15a822013-08-31 19:17:41 -0700738 # mixed-in Enums should use the mixed-in type's __format__, otherwise
739 # we can get strange results with the Enum name showing up instead of
740 # the value
741
thatneat2f19e822019-07-04 11:28:37 -0700742 # pure Enum branch, or branch with __str__ explicitly overridden
Ethan Furman37440ee2020-12-08 11:14:10 -0800743 str_overridden = type(self).__str__ not in (Enum.__str__, Flag.__str__)
thatneat2f19e822019-07-04 11:28:37 -0700744 if self._member_type_ is object or str_overridden:
Ethan Furmanec15a822013-08-31 19:17:41 -0700745 cls = str
746 val = str(self)
747 # mix-in branch
748 else:
749 cls = self._member_type_
Ethan Furman0081f232014-09-16 17:31:23 -0700750 val = self._value_
Ethan Furmanec15a822013-08-31 19:17:41 -0700751 return cls.__format__(val, format_spec)
752
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700753 def __hash__(self):
Ethan Furman520ad572013-07-19 19:47:21 -0700754 return hash(self._name_)
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700755
Ethan Furmanca1b7942014-02-08 11:36:27 -0800756 def __reduce_ex__(self, proto):
Ethan Furmandc870522014-02-18 12:37:12 -0800757 return self.__class__, (self._value_, )
Ethan Furmanca1b7942014-02-08 11:36:27 -0800758
Ethan Furman33918c12013-09-27 23:02:02 -0700759 # DynamicClassAttribute is used to provide access to the `name` and
760 # `value` properties of enum members while keeping some measure of
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700761 # protection from modification, while still allowing for an enumeration
762 # to have members named `name` and `value`. This works because enumeration
763 # members are not set directly on the enum class -- __getattr__ is
764 # used to look them up.
765
Ethan Furmane03ea372013-09-25 07:14:41 -0700766 @DynamicClassAttribute
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700767 def name(self):
Ethan Furmanc850f342013-09-15 16:59:35 -0700768 """The name of the Enum member."""
Ethan Furman520ad572013-07-19 19:47:21 -0700769 return self._name_
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700770
Ethan Furmane03ea372013-09-25 07:14:41 -0700771 @DynamicClassAttribute
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700772 def value(self):
Ethan Furmanc850f342013-09-15 16:59:35 -0700773 """The value of the Enum member."""
Ethan Furman520ad572013-07-19 19:47:21 -0700774 return self._value_
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700775
776
777class IntEnum(int, Enum):
Ethan Furman0063ff42020-09-21 17:23:13 -0700778 """
779 Enum where members are also (and must be) ints
780 """
781
782
783class StrEnum(str, Enum):
784 """
785 Enum where members are also (and must be) strings
786 """
787
788 def __new__(cls, *values):
789 if len(values) > 3:
790 raise TypeError('too many arguments for str(): %r' % (values, ))
791 if len(values) == 1:
792 # it must be a string
793 if not isinstance(values[0], str):
794 raise TypeError('%r is not a string' % (values[0], ))
795 if len(values) > 1:
796 # check that encoding argument is a string
797 if not isinstance(values[1], str):
798 raise TypeError('encoding must be a string, not %r' % (values[1], ))
799 if len(values) > 2:
800 # check that errors argument is a string
801 if not isinstance(values[2], str):
802 raise TypeError('errors must be a string, not %r' % (values[2], ))
803 value = str(*values)
804 member = str.__new__(cls, value)
805 member._value_ = value
806 return member
Ethan Furmanf24bb352013-07-18 17:05:39 -0700807
Ethan Furmand986d162020-09-22 13:00:07 -0700808 __str__ = str.__str__
809
Ethan Furmanefb13be2020-12-10 12:20:06 -0800810 def _generate_next_value_(name, start, count, last_values):
811 """
812 Return the lower-cased version of the member name.
813 """
814 return name.lower()
815
Ethan Furmanf24bb352013-07-18 17:05:39 -0700816
Ethan Furman24e837f2015-03-18 17:27:57 -0700817def _reduce_ex_by_name(self, proto):
818 return self.name
819
Ethan Furman65a5a472016-09-01 23:55:19 -0700820class Flag(Enum):
Ethan Furman6d3dfee2020-12-08 12:26:56 -0800821 """
822 Support for flags
823 """
Ethan Furmanc16595e2016-09-10 23:36:59 -0700824
825 def _generate_next_value_(name, start, count, last_values):
Ethan Furmanee47e5c2016-08-31 00:12:15 -0700826 """
827 Generate the next value when not given.
828
829 name: the name of the member
HongWeipengbb16fb22019-09-21 13:22:54 +0800830 start: the initial start value or None
Ethan Furmanee47e5c2016-08-31 00:12:15 -0700831 count: the number of existing members
832 last_value: the last value assigned or None
833 """
834 if not count:
835 return start if start is not None else 1
Ethan Furmanc16595e2016-09-10 23:36:59 -0700836 for last_value in reversed(last_values):
837 try:
838 high_bit = _high_bit(last_value)
839 break
Ethan Furman3515dcc2016-09-18 13:15:41 -0700840 except Exception:
Ethan Furmanc16595e2016-09-10 23:36:59 -0700841 raise TypeError('Invalid Flag value: %r' % last_value) from None
Ethan Furmanee47e5c2016-08-31 00:12:15 -0700842 return 2 ** (high_bit+1)
843
844 @classmethod
845 def _missing_(cls, value):
Ethan Furman6d3dfee2020-12-08 12:26:56 -0800846 """
847 Returns member (possibly creating it) if one can be found for value.
848 """
Ethan Furmanee47e5c2016-08-31 00:12:15 -0700849 original_value = value
850 if value < 0:
851 value = ~value
852 possible_member = cls._create_pseudo_member_(value)
Ethan Furmanee47e5c2016-08-31 00:12:15 -0700853 if original_value < 0:
854 possible_member = ~possible_member
855 return possible_member
856
857 @classmethod
858 def _create_pseudo_member_(cls, value):
Ethan Furman3515dcc2016-09-18 13:15:41 -0700859 """
860 Create a composite member iff value contains only members.
861 """
Ethan Furmanee47e5c2016-08-31 00:12:15 -0700862 pseudo_member = cls._value2member_map_.get(value, None)
863 if pseudo_member is None:
Ethan Furman3515dcc2016-09-18 13:15:41 -0700864 # verify all bits are accounted for
865 _, extra_flags = _decompose(cls, value)
866 if extra_flags:
Walter Dörwald323842c2019-07-18 20:37:13 +0200867 raise ValueError("%r is not a valid %s" % (value, cls.__qualname__))
Ethan Furman3515dcc2016-09-18 13:15:41 -0700868 # construct a singleton enum pseudo-member
Ethan Furmanee47e5c2016-08-31 00:12:15 -0700869 pseudo_member = object.__new__(cls)
870 pseudo_member._name_ = None
871 pseudo_member._value_ = value
Ethan Furman28cf6632017-01-24 12:12:06 -0800872 # use setdefault in case another thread already created a composite
873 # with this value
874 pseudo_member = cls._value2member_map_.setdefault(value, pseudo_member)
Ethan Furmanee47e5c2016-08-31 00:12:15 -0700875 return pseudo_member
876
Ethan Furmanee47e5c2016-08-31 00:12:15 -0700877 def __contains__(self, other):
Ethan Furman6d3dfee2020-12-08 12:26:56 -0800878 """
879 Returns True if self has at least the same flags set as other.
880 """
Ethan Furmanee47e5c2016-08-31 00:12:15 -0700881 if not isinstance(other, self.__class__):
Rahul Jha94306522018-09-10 23:51:04 +0530882 raise TypeError(
883 "unsupported operand type(s) for 'in': '%s' and '%s'" % (
884 type(other).__qualname__, self.__class__.__qualname__))
Ethan Furmanee47e5c2016-08-31 00:12:15 -0700885 return other._value_ & self._value_ == other._value_
886
Ethan Furman7219e272020-09-16 13:01:00 -0700887 def __iter__(self):
Ethan Furman6d3dfee2020-12-08 12:26:56 -0800888 """
889 Returns flags in decreasing value order.
890 """
Ethan Furman7219e272020-09-16 13:01:00 -0700891 members, extra_flags = _decompose(self.__class__, self.value)
892 return (m for m in members if m._value_ != 0)
893
Ethan Furmanee47e5c2016-08-31 00:12:15 -0700894 def __repr__(self):
895 cls = self.__class__
896 if self._name_ is not None:
897 return '<%s.%s: %r>' % (cls.__name__, self._name_, self._value_)
Ethan Furman3515dcc2016-09-18 13:15:41 -0700898 members, uncovered = _decompose(cls, self._value_)
Ethan Furman27682d22016-09-04 11:39:01 -0700899 return '<%s.%s: %r>' % (
900 cls.__name__,
901 '|'.join([str(m._name_ or m._value_) for m in members]),
902 self._value_,
903 )
Ethan Furmanee47e5c2016-08-31 00:12:15 -0700904
905 def __str__(self):
906 cls = self.__class__
907 if self._name_ is not None:
908 return '%s.%s' % (cls.__name__, self._name_)
Ethan Furman3515dcc2016-09-18 13:15:41 -0700909 members, uncovered = _decompose(cls, self._value_)
Ethan Furmanee47e5c2016-08-31 00:12:15 -0700910 if len(members) == 1 and members[0]._name_ is None:
911 return '%s.%r' % (cls.__name__, members[0]._value_)
912 else:
913 return '%s.%s' % (
914 cls.__name__,
915 '|'.join([str(m._name_ or m._value_) for m in members]),
916 )
917
Ethan Furman25d94bb2016-09-02 16:32:32 -0700918 def __bool__(self):
919 return bool(self._value_)
920
Ethan Furmanee47e5c2016-08-31 00:12:15 -0700921 def __or__(self, other):
922 if not isinstance(other, self.__class__):
923 return NotImplemented
924 return self.__class__(self._value_ | other._value_)
925
926 def __and__(self, other):
927 if not isinstance(other, self.__class__):
928 return NotImplemented
929 return self.__class__(self._value_ & other._value_)
930
931 def __xor__(self, other):
932 if not isinstance(other, self.__class__):
933 return NotImplemented
934 return self.__class__(self._value_ ^ other._value_)
935
936 def __invert__(self):
Ethan Furman3515dcc2016-09-18 13:15:41 -0700937 members, uncovered = _decompose(self.__class__, self._value_)
Serhiy Storchaka81108372017-09-26 00:55:55 +0300938 inverted = self.__class__(0)
939 for m in self.__class__:
940 if m not in members and not (m._value_ & self._value_):
941 inverted = inverted | m
Ethan Furmanee47e5c2016-08-31 00:12:15 -0700942 return self.__class__(inverted)
943
944
Ethan Furman65a5a472016-09-01 23:55:19 -0700945class IntFlag(int, Flag):
Ethan Furman6d3dfee2020-12-08 12:26:56 -0800946 """
947 Support for integer-based Flags
948 """
Ethan Furmanee47e5c2016-08-31 00:12:15 -0700949
950 @classmethod
Ethan Furman3515dcc2016-09-18 13:15:41 -0700951 def _missing_(cls, value):
Ethan Furman6d3dfee2020-12-08 12:26:56 -0800952 """
953 Returns member (possibly creating it) if one can be found for value.
954 """
Ethan Furman3515dcc2016-09-18 13:15:41 -0700955 if not isinstance(value, int):
Walter Dörwald323842c2019-07-18 20:37:13 +0200956 raise ValueError("%r is not a valid %s" % (value, cls.__qualname__))
Ethan Furman3515dcc2016-09-18 13:15:41 -0700957 new_member = cls._create_pseudo_member_(value)
958 return new_member
959
960 @classmethod
Ethan Furmanee47e5c2016-08-31 00:12:15 -0700961 def _create_pseudo_member_(cls, value):
Ethan Furman6d3dfee2020-12-08 12:26:56 -0800962 """
963 Create a composite member iff value contains only members.
964 """
Ethan Furmanee47e5c2016-08-31 00:12:15 -0700965 pseudo_member = cls._value2member_map_.get(value, None)
966 if pseudo_member is None:
Ethan Furman3515dcc2016-09-18 13:15:41 -0700967 need_to_create = [value]
968 # get unaccounted for bits
969 _, extra_flags = _decompose(cls, value)
970 # timer = 10
971 while extra_flags:
972 # timer -= 1
973 bit = _high_bit(extra_flags)
974 flag_value = 2 ** bit
975 if (flag_value not in cls._value2member_map_ and
976 flag_value not in need_to_create
977 ):
978 need_to_create.append(flag_value)
979 if extra_flags == -flag_value:
980 extra_flags = 0
981 else:
982 extra_flags ^= flag_value
983 for value in reversed(need_to_create):
984 # construct singleton pseudo-members
985 pseudo_member = int.__new__(cls, value)
986 pseudo_member._name_ = None
987 pseudo_member._value_ = value
Ethan Furman28cf6632017-01-24 12:12:06 -0800988 # use setdefault in case another thread already created a composite
989 # with this value
990 pseudo_member = cls._value2member_map_.setdefault(value, pseudo_member)
Ethan Furmanee47e5c2016-08-31 00:12:15 -0700991 return pseudo_member
992
Ethan Furmanee47e5c2016-08-31 00:12:15 -0700993 def __or__(self, other):
994 if not isinstance(other, (self.__class__, int)):
995 return NotImplemented
Ethan Furman3515dcc2016-09-18 13:15:41 -0700996 result = self.__class__(self._value_ | self.__class__(other)._value_)
997 return result
Ethan Furmanee47e5c2016-08-31 00:12:15 -0700998
999 def __and__(self, other):
1000 if not isinstance(other, (self.__class__, int)):
1001 return NotImplemented
1002 return self.__class__(self._value_ & self.__class__(other)._value_)
1003
1004 def __xor__(self, other):
1005 if not isinstance(other, (self.__class__, int)):
1006 return NotImplemented
1007 return self.__class__(self._value_ ^ self.__class__(other)._value_)
1008
1009 __ror__ = __or__
1010 __rand__ = __and__
1011 __rxor__ = __xor__
1012
1013 def __invert__(self):
Ethan Furman3515dcc2016-09-18 13:15:41 -07001014 result = self.__class__(~self._value_)
1015 return result
Ethan Furmanee47e5c2016-08-31 00:12:15 -07001016
1017
1018def _high_bit(value):
Ethan Furman6d3dfee2020-12-08 12:26:56 -08001019 """
1020 returns index of highest bit, or -1 if value is zero or negative
1021 """
Ethan Furman3515dcc2016-09-18 13:15:41 -07001022 return value.bit_length() - 1
Ethan Furmanee47e5c2016-08-31 00:12:15 -07001023
Ethan Furmanf24bb352013-07-18 17:05:39 -07001024def unique(enumeration):
Ethan Furman6d3dfee2020-12-08 12:26:56 -08001025 """
1026 Class decorator for enumerations ensuring unique member values.
1027 """
Ethan Furmanf24bb352013-07-18 17:05:39 -07001028 duplicates = []
1029 for name, member in enumeration.__members__.items():
1030 if name != member.name:
1031 duplicates.append((name, member.name))
1032 if duplicates:
1033 alias_details = ', '.join(
1034 ["%s -> %s" % (alias, name) for (alias, name) in duplicates])
1035 raise ValueError('duplicate values found in %r: %s' %
1036 (enumeration, alias_details))
1037 return enumeration
Ethan Furman3515dcc2016-09-18 13:15:41 -07001038
1039def _decompose(flag, value):
Ethan Furman6d3dfee2020-12-08 12:26:56 -08001040 """
1041 Extract all members from the value.
1042 """
Ethan Furman3515dcc2016-09-18 13:15:41 -07001043 # _decompose is only called if the value is not named
1044 not_covered = value
1045 negative = value < 0
Ethan Furman3515dcc2016-09-18 13:15:41 -07001046 members = []
HongWeipeng0b41a922019-11-27 06:36:02 +08001047 for member in flag:
1048 member_value = member.value
Ethan Furman3515dcc2016-09-18 13:15:41 -07001049 if member_value and member_value & value == member_value:
1050 members.append(member)
1051 not_covered &= ~member_value
HongWeipeng0b41a922019-11-27 06:36:02 +08001052 if not negative:
1053 tmp = not_covered
1054 while tmp:
1055 flag_value = 2 ** _high_bit(tmp)
1056 if flag_value in flag._value2member_map_:
1057 members.append(flag._value2member_map_[flag_value])
1058 not_covered &= ~flag_value
1059 tmp &= ~flag_value
Ethan Furman3515dcc2016-09-18 13:15:41 -07001060 if not members and value in flag._value2member_map_:
1061 members.append(flag._value2member_map_[value])
1062 members.sort(key=lambda m: m._value_, reverse=True)
1063 if len(members) > 1 and members[0].value == value:
1064 # we have the breakdown, don't need the value member itself
1065 members.pop(0)
1066 return members, not_covered