blob: f6c7e8b2334139571ae3103197b4eaa7f23e787c [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 Furman6b3d64a2013-06-14 16:55:46 -070044def _make_class_unpicklable(cls):
Ethan Furman6d3dfee2020-12-08 12:26:56 -080045 """
46 Make the given class un-picklable.
47 """
Ethan Furmanca1b7942014-02-08 11:36:27 -080048 def _break_on_call_reduce(self, proto):
Ethan Furman6b3d64a2013-06-14 16:55:46 -070049 raise TypeError('%r cannot be pickled' % self)
Ethan Furmanca1b7942014-02-08 11:36:27 -080050 cls.__reduce_ex__ = _break_on_call_reduce
Ethan Furman6b3d64a2013-06-14 16:55:46 -070051 cls.__module__ = '<unknown>'
52
Ethan Furman3515dcc2016-09-18 13:15:41 -070053_auto_null = object()
Ethan Furmanc16595e2016-09-10 23:36:59 -070054class auto:
55 """
56 Instances are replaced with an appropriate value in Enum class suites.
57 """
Ethan Furman3515dcc2016-09-18 13:15:41 -070058 value = _auto_null
Ethan Furmanc16595e2016-09-10 23:36:59 -070059
Ethan Furman101e0742013-09-15 12:34:36 -070060
Ethan Furman6b3d64a2013-06-14 16:55:46 -070061class _EnumDict(dict):
Ethan Furman6d3dfee2020-12-08 12:26:56 -080062 """
63 Track enum member order and ensure member names are not reused.
Ethan Furman6b3d64a2013-06-14 16:55:46 -070064
65 EnumMeta will use the names found in self._member_names as the
66 enumeration member names.
Ethan Furman6b3d64a2013-06-14 16:55:46 -070067 """
68 def __init__(self):
69 super().__init__()
70 self._member_names = []
Ethan Furmanc16595e2016-09-10 23:36:59 -070071 self._last_values = []
Ethan Furmana4b1bb42018-01-22 07:56:37 -080072 self._ignore = []
Ethan Onstottd9a43e22020-04-28 13:20:55 -040073 self._auto_called = False
Ethan Furman6b3d64a2013-06-14 16:55:46 -070074
75 def __setitem__(self, key, value):
Ethan Furman6d3dfee2020-12-08 12:26:56 -080076 """
77 Changes anything not dundered or not a descriptor.
Ethan Furman6b3d64a2013-06-14 16:55:46 -070078
79 If an enum member name is used twice, an error is raised; duplicate
80 values are not checked for.
81
82 Single underscore (sunder) names are reserved.
Ethan Furman6b3d64a2013-06-14 16:55:46 -070083 """
84 if _is_sunder(key):
Ethan Furmanee47e5c2016-08-31 00:12:15 -070085 if key not in (
Ethan Furman3515dcc2016-09-18 13:15:41 -070086 '_order_', '_create_pseudo_member_',
Ethan Furmana4b1bb42018-01-22 07:56:37 -080087 '_generate_next_value_', '_missing_', '_ignore_',
Ethan Furmanee47e5c2016-08-31 00:12:15 -070088 ):
Ethan Furman6d3dfee2020-12-08 12:26:56 -080089 raise ValueError(
90 '_sunder_ names, such as %r, are reserved for future Enum use'
91 % (key, )
92 )
Ethan Furmanc16595e2016-09-10 23:36:59 -070093 if key == '_generate_next_value_':
Ethan Onstottd9a43e22020-04-28 13:20:55 -040094 # check if members already defined as auto()
95 if self._auto_called:
96 raise TypeError("_generate_next_value_ must be defined before members")
Ethan Furmanc16595e2016-09-10 23:36:59 -070097 setattr(self, '_generate_next_value', value)
Ethan Furmana4b1bb42018-01-22 07:56:37 -080098 elif key == '_ignore_':
99 if isinstance(value, str):
100 value = value.replace(',',' ').split()
101 else:
102 value = list(value)
103 self._ignore = value
104 already = set(value) & set(self._member_names)
105 if already:
Ethan Furman6d3dfee2020-12-08 12:26:56 -0800106 raise ValueError(
107 '_ignore_ cannot specify already set names: %r'
108 % (already, )
109 )
Ethan Furman101e0742013-09-15 12:34:36 -0700110 elif _is_dunder(key):
Ethan Furmane8e61272016-08-20 07:19:31 -0700111 if key == '__order__':
112 key = '_order_'
Ethan Furman101e0742013-09-15 12:34:36 -0700113 elif key in self._member_names:
114 # descriptor overwriting an enum?
115 raise TypeError('Attempted to reuse key: %r' % key)
Ethan Furmana4b1bb42018-01-22 07:56:37 -0800116 elif key in self._ignore:
117 pass
Ethan Furman101e0742013-09-15 12:34:36 -0700118 elif not _is_descriptor(value):
119 if key in self:
120 # enum overwriting a descriptor?
Ethan Furmanee47e5c2016-08-31 00:12:15 -0700121 raise TypeError('%r already defined as: %r' % (key, self[key]))
Ethan Furmanc16595e2016-09-10 23:36:59 -0700122 if isinstance(value, auto):
Ethan Furman3515dcc2016-09-18 13:15:41 -0700123 if value.value == _auto_null:
Ethan Furman6d3dfee2020-12-08 12:26:56 -0800124 value.value = self._generate_next_value(
125 key,
126 1,
127 len(self._member_names),
128 self._last_values[:],
129 )
Ethan Furmanfc23a942020-09-16 12:37:54 -0700130 self._auto_called = True
Ethan Furman3515dcc2016-09-18 13:15:41 -0700131 value = value.value
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700132 self._member_names.append(key)
Ethan Furmanc16595e2016-09-10 23:36:59 -0700133 self._last_values.append(value)
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700134 super().__setitem__(key, value)
135
136
Ezio Melotti9a3777e2013-08-17 15:53:55 +0300137# Dummy value for Enum as EnumMeta explicitly checks for it, but of course
138# until EnumMeta finishes running the first time the Enum class doesn't exist.
139# This is also why there are checks in EnumMeta like `if Enum is not None`
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700140Enum = None
141
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700142class EnumMeta(type):
Ethan Furman6d3dfee2020-12-08 12:26:56 -0800143 """
144 Metaclass for Enum
145 """
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700146 @classmethod
Ethan Furman332dbc72016-08-20 00:00:52 -0700147 def __prepare__(metacls, cls, bases):
Ethan Furman3064dbf2020-09-16 07:11:57 -0700148 # check that previous enum members do not exist
149 metacls._check_for_existing_members(cls, bases)
Ethan Furmanee47e5c2016-08-31 00:12:15 -0700150 # create the namespace dict
151 enum_dict = _EnumDict()
152 # inherit previous flags and _generate_next_value_ function
Ethan Furman3064dbf2020-09-16 07:11:57 -0700153 member_type, first_enum = metacls._get_mixins_(cls, bases)
Ethan Furmanee47e5c2016-08-31 00:12:15 -0700154 if first_enum is not None:
Ethan Furman6d3dfee2020-12-08 12:26:56 -0800155 enum_dict['_generate_next_value_'] = getattr(
156 first_enum, '_generate_next_value_', None,
157 )
Ethan Furmanee47e5c2016-08-31 00:12:15 -0700158 return enum_dict
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700159
Ethan Furman65a5a472016-09-01 23:55:19 -0700160 def __new__(metacls, cls, bases, classdict):
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700161 # an Enum class is final once enumeration items have been defined; it
162 # cannot be mixed with other types (int, float, etc.) if it has an
163 # inherited __new__ unless a new __new__ is defined (or the resulting
164 # class will fail).
Ethan Furmana4b1bb42018-01-22 07:56:37 -0800165 #
166 # remove any keys listed in _ignore_
167 classdict.setdefault('_ignore_', []).append('_ignore_')
168 ignore = classdict['_ignore_']
169 for key in ignore:
170 classdict.pop(key, None)
Ethan Furman3064dbf2020-09-16 07:11:57 -0700171 member_type, first_enum = metacls._get_mixins_(cls, bases)
Ethan Furmanc2667362020-12-07 00:17:31 -0800172 __new__, save_new, use_args = metacls._find_new_(
173 classdict, member_type, first_enum,
174 )
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700175
176 # save enum items into separate mapping so they don't get baked into
177 # the new class
Ethan Furmanee47e5c2016-08-31 00:12:15 -0700178 enum_members = {k: classdict[k] for k in classdict._member_names}
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700179 for name in classdict._member_names:
180 del classdict[name]
181
Ethan Furmane8e61272016-08-20 07:19:31 -0700182 # adjust the sunders
183 _order_ = classdict.pop('_order_', None)
184
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700185 # check for illegal enum names (any others?)
Brennan D Baraban8b914d22019-03-03 14:09:11 -0800186 invalid_names = set(enum_members) & {'mro', ''}
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700187 if invalid_names:
188 raise ValueError('Invalid enum member name: {0}'.format(
189 ','.join(invalid_names)))
190
Ethan Furman48a724f2015-04-11 23:23:06 -0700191 # create a default docstring if one has not been provided
192 if '__doc__' not in classdict:
193 classdict['__doc__'] = 'An enumeration.'
194
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700195 # create our new Enum type
196 enum_class = super().__new__(metacls, cls, bases, classdict)
Ethan Furman520ad572013-07-19 19:47:21 -0700197 enum_class._member_names_ = [] # names in definition order
INADA Naokie57f91a2018-06-19 01:14:26 +0900198 enum_class._member_map_ = {} # name->value map
Ethan Furman5e5a8232013-08-04 08:42:23 -0700199 enum_class._member_type_ = member_type
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700200
orlnub1230fb9fad2018-09-12 20:28:53 +0300201 # save DynamicClassAttribute attributes from super classes so we know
202 # if we can take the shortcut of storing members in the class dict
Ethan Furman6d3dfee2020-12-08 12:26:56 -0800203 dynamic_attributes = {
204 k for c in enum_class.mro()
205 for k, v in c.__dict__.items()
206 if isinstance(v, DynamicClassAttribute)
207 }
Ethan Furman354ecf12015-03-11 08:43:12 -0700208
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700209 # Reverse value->name map for hashable values.
Ethan Furman520ad572013-07-19 19:47:21 -0700210 enum_class._value2member_map_ = {}
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700211
Ethan Furman2da95042014-03-03 12:42:52 -0800212 # If a custom type is mixed into the Enum, and it does not know how
213 # to pickle itself, pickle.dumps will succeed but pickle.loads will
214 # fail. Rather than have the error show up later and possibly far
215 # from the source, sabotage the pickle protocol for this class so
216 # that pickle.dumps also fails.
217 #
218 # However, if the new class implements its own __reduce_ex__, do not
219 # sabotage -- it's on them to make sure it works correctly. We use
220 # __reduce_ex__ instead of any of the others as it is preferred by
221 # pickle over __reduce__, and it handles all pickle protocols.
222 if '__reduce_ex__' not in classdict:
Ethan Furmandc870522014-02-18 12:37:12 -0800223 if member_type is not object:
224 methods = ('__getnewargs_ex__', '__getnewargs__',
225 '__reduce_ex__', '__reduce__')
Ethan Furman2da95042014-03-03 12:42:52 -0800226 if not any(m in member_type.__dict__ for m in methods):
Ethan Furmandc870522014-02-18 12:37:12 -0800227 _make_class_unpicklable(enum_class)
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700228
229 # instantiate them, checking for duplicates as we go
230 # we instantiate first instead of checking for duplicates first in case
231 # a custom __new__ is doing something funky with the values -- such as
232 # auto-numbering ;)
233 for member_name in classdict._member_names:
Ethan Furmanee47e5c2016-08-31 00:12:15 -0700234 value = enum_members[member_name]
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700235 if not isinstance(value, tuple):
236 args = (value, )
237 else:
238 args = value
239 if member_type is tuple: # special case for tuple enums
240 args = (args, ) # wrap it one more time
241 if not use_args:
242 enum_member = __new__(enum_class)
Ethan Furmanb41803e2013-07-25 13:50:45 -0700243 if not hasattr(enum_member, '_value_'):
244 enum_member._value_ = value
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700245 else:
246 enum_member = __new__(enum_class, *args)
Ethan Furmanb41803e2013-07-25 13:50:45 -0700247 if not hasattr(enum_member, '_value_'):
Ethan Furmanee47e5c2016-08-31 00:12:15 -0700248 if member_type is object:
249 enum_member._value_ = value
250 else:
251 enum_member._value_ = member_type(*args)
Ethan Furman520ad572013-07-19 19:47:21 -0700252 value = enum_member._value_
Ethan Furman520ad572013-07-19 19:47:21 -0700253 enum_member._name_ = member_name
Ethan Furmanc850f342013-09-15 16:59:35 -0700254 enum_member.__objclass__ = enum_class
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700255 enum_member.__init__(*args)
256 # If another member with the same value was already defined, the
257 # new member becomes an alias to the existing one.
Ethan Furman520ad572013-07-19 19:47:21 -0700258 for name, canonical_member in enum_class._member_map_.items():
Ethan Furman0081f232014-09-16 17:31:23 -0700259 if canonical_member._value_ == enum_member._value_:
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700260 enum_member = canonical_member
261 break
262 else:
263 # Aliases don't appear in member names (only in __members__).
Ethan Furman520ad572013-07-19 19:47:21 -0700264 enum_class._member_names_.append(member_name)
Ethan Furman354ecf12015-03-11 08:43:12 -0700265 # performance boost for any member that would not shadow
266 # a DynamicClassAttribute
orlnub1230fb9fad2018-09-12 20:28:53 +0300267 if member_name not in dynamic_attributes:
Ethan Furman354ecf12015-03-11 08:43:12 -0700268 setattr(enum_class, member_name, enum_member)
269 # now add to _member_map_
Ethan Furman520ad572013-07-19 19:47:21 -0700270 enum_class._member_map_[member_name] = enum_member
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700271 try:
272 # This may fail if value is not hashable. We can't add the value
273 # to the map, and by-value lookups for this value will be
274 # linear.
Ethan Furman520ad572013-07-19 19:47:21 -0700275 enum_class._value2member_map_[value] = enum_member
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700276 except TypeError:
277 pass
278
279 # double check that repr and friends are not the mixin's or various
280 # things break (such as pickle)
Ethan Furman22415ad2020-09-15 16:28:25 -0700281 # however, if the method is defined in the Enum itself, don't replace
282 # it
Ethan Furmandc870522014-02-18 12:37:12 -0800283 for name in ('__repr__', '__str__', '__format__', '__reduce_ex__'):
Ethan Furman22415ad2020-09-15 16:28:25 -0700284 if name in classdict:
285 continue
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700286 class_method = getattr(enum_class, name)
287 obj_method = getattr(member_type, name, None)
288 enum_method = getattr(first_enum, name, None)
289 if obj_method is not None and obj_method is class_method:
290 setattr(enum_class, name, enum_method)
291
292 # replace any other __new__ with our own (as long as Enum is not None,
293 # anyway) -- again, this is to support pickle
294 if Enum is not None:
295 # if the user defined their own __new__, save it before it gets
296 # clobbered in case they subclass later
297 if save_new:
298 enum_class.__new_member__ = __new__
299 enum_class.__new__ = Enum.__new__
Ethan Furmane8e61272016-08-20 07:19:31 -0700300
301 # py3 support for definition order (helps keep py2/py3 code in sync)
302 if _order_ is not None:
303 if isinstance(_order_, str):
304 _order_ = _order_.replace(',', ' ').split()
305 if _order_ != enum_class._member_names_:
306 raise TypeError('member order does not match _order_')
307
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700308 return enum_class
309
Ethan Furman5de67b12016-04-13 23:52:09 -0700310 def __bool__(self):
311 """
312 classes/types should always be True.
313 """
314 return True
315
Ethan Furmand9925a12014-09-16 20:35:55 -0700316 def __call__(cls, value, names=None, *, module=None, qualname=None, type=None, start=1):
Ethan Furman6d3dfee2020-12-08 12:26:56 -0800317 """
318 Either returns an existing member, or creates a new enum class.
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700319
320 This method is used both when an enum class is given a value to match
321 to an enumeration member (i.e. Color(3)) and for the functional API
Ethan Furman23bb6f42016-11-21 09:22:05 -0800322 (i.e. Color = Enum('Color', names='RED GREEN BLUE')).
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700323
Ethan Furman2da95042014-03-03 12:42:52 -0800324 When used for the functional API:
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700325
Ethan Furman2da95042014-03-03 12:42:52 -0800326 `value` will be the name of the new class.
327
328 `names` should be either a string of white-space/comma delimited names
Ethan Furmand9925a12014-09-16 20:35:55 -0700329 (values will start at `start`), or an iterator/mapping of name, value pairs.
Ethan Furman2da95042014-03-03 12:42:52 -0800330
331 `module` should be set to the module this class is being created in;
332 if it is not set, an attempt to find that module will be made, but if
333 it fails the class will not be picklable.
334
335 `qualname` should be set to the actual location this class can be found
336 at in its module; by default it is set to the global scope. If this is
337 not correct, unpickling will fail in some circumstances.
338
339 `type`, if set, will be mixed in as the first base class.
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700340 """
341 if names is None: # simple value lookup
342 return cls.__new__(cls, value)
343 # otherwise, functional API: we're creating a new Enum type
Ethan Furman6d3dfee2020-12-08 12:26:56 -0800344 return cls._create_(
345 value,
346 names,
347 module=module,
348 qualname=qualname,
349 type=type,
350 start=start,
351 )
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700352
353 def __contains__(cls, member):
Rahul Jha94306522018-09-10 23:51:04 +0530354 if not isinstance(member, Enum):
355 raise TypeError(
356 "unsupported operand type(s) for 'in': '%s' and '%s'" % (
357 type(member).__qualname__, cls.__class__.__qualname__))
Ethan Furman0081f232014-09-16 17:31:23 -0700358 return isinstance(member, cls) and member._name_ in cls._member_map_
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700359
Ethan Furman64a99722013-09-22 16:18:19 -0700360 def __delattr__(cls, attr):
361 # nicer error message when someone tries to delete an attribute
362 # (see issue19025).
363 if attr in cls._member_map_:
Ethan Furman6d3dfee2020-12-08 12:26:56 -0800364 raise AttributeError("%s: cannot delete Enum member %r." % (cls.__name__, attr))
Ethan Furman64a99722013-09-22 16:18:19 -0700365 super().__delattr__(attr)
366
Ethan Furman388a3922013-08-12 06:51:41 -0700367 def __dir__(self):
Ethan Furman6d3dfee2020-12-08 12:26:56 -0800368 return (
369 ['__class__', '__doc__', '__members__', '__module__']
370 + self._member_names_
371 )
Ethan Furman388a3922013-08-12 06:51:41 -0700372
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700373 def __getattr__(cls, name):
Ethan Furman6d3dfee2020-12-08 12:26:56 -0800374 """
375 Return the enum member matching `name`
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700376
377 We use __getattr__ instead of descriptors or inserting into the enum
378 class' __dict__ in order to support `name` and `value` being both
379 properties for enum members (which live in the class' __dict__) and
380 enum members themselves.
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700381 """
382 if _is_dunder(name):
383 raise AttributeError(name)
384 try:
Ethan Furman520ad572013-07-19 19:47:21 -0700385 return cls._member_map_[name]
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700386 except KeyError:
387 raise AttributeError(name) from None
388
389 def __getitem__(cls, name):
Ethan Furman520ad572013-07-19 19:47:21 -0700390 return cls._member_map_[name]
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700391
392 def __iter__(cls):
Ethan Furman6d3dfee2020-12-08 12:26:56 -0800393 """
394 Returns members in definition order.
395 """
Ethan Furman520ad572013-07-19 19:47:21 -0700396 return (cls._member_map_[name] for name in cls._member_names_)
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700397
398 def __len__(cls):
Ethan Furman520ad572013-07-19 19:47:21 -0700399 return len(cls._member_names_)
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700400
Ethan Furman2131a4a2013-09-14 18:11:24 -0700401 @property
402 def __members__(cls):
Ethan Furman6d3dfee2020-12-08 12:26:56 -0800403 """
404 Returns a mapping of member name->value.
Ethan Furman2131a4a2013-09-14 18:11:24 -0700405
406 This mapping lists all enum members, including aliases. Note that this
407 is a read-only view of the internal mapping.
Ethan Furman2131a4a2013-09-14 18:11:24 -0700408 """
409 return MappingProxyType(cls._member_map_)
410
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700411 def __repr__(cls):
412 return "<enum %r>" % cls.__name__
413
Ethan Furman2131a4a2013-09-14 18:11:24 -0700414 def __reversed__(cls):
Ethan Furman6d3dfee2020-12-08 12:26:56 -0800415 """
416 Returns members in reverse definition order.
417 """
Ethan Furman2131a4a2013-09-14 18:11:24 -0700418 return (cls._member_map_[name] for name in reversed(cls._member_names_))
419
Ethan Furmanf203f2d2013-09-06 07:16:48 -0700420 def __setattr__(cls, name, value):
Ethan Furman6d3dfee2020-12-08 12:26:56 -0800421 """
422 Block attempts to reassign Enum members.
Ethan Furmanf203f2d2013-09-06 07:16:48 -0700423
424 A simple assignment to the class namespace only changes one of the
425 several possible ways to get an Enum member from the Enum class,
426 resulting in an inconsistent Enumeration.
Ethan Furmanf203f2d2013-09-06 07:16:48 -0700427 """
428 member_map = cls.__dict__.get('_member_map_', {})
429 if name in member_map:
430 raise AttributeError('Cannot reassign members.')
431 super().__setattr__(name, value)
432
anentropicb8e21f12018-04-16 04:40:35 +0100433 def _create_(cls, class_name, names, *, module=None, qualname=None, type=None, start=1):
Ethan Furman6d3dfee2020-12-08 12:26:56 -0800434 """
435 Convenience method to create a new Enum class.
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700436
437 `names` can be:
438
439 * A string containing member names, separated either with spaces or
Ethan Furmand9925a12014-09-16 20:35:55 -0700440 commas. Values are incremented by 1 from `start`.
441 * An iterable of member names. Values are incremented by 1 from `start`.
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700442 * An iterable of (member name, value) pairs.
Ethan Furmand9925a12014-09-16 20:35:55 -0700443 * A mapping of member name -> value pairs.
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700444 """
445 metacls = cls.__class__
446 bases = (cls, ) if type is None else (type, cls)
Ethan Furman3064dbf2020-09-16 07:11:57 -0700447 _, first_enum = cls._get_mixins_(cls, bases)
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700448 classdict = metacls.__prepare__(class_name, bases)
449
450 # special processing needed for names?
451 if isinstance(names, str):
452 names = names.replace(',', ' ').split()
Dong-hee Nadcc8ce42017-06-22 01:52:32 +0900453 if isinstance(names, (tuple, list)) and names and isinstance(names[0], str):
Ethan Furmanee47e5c2016-08-31 00:12:15 -0700454 original_names, names = names, []
Ethan Furmanc16595e2016-09-10 23:36:59 -0700455 last_values = []
Ethan Furmanee47e5c2016-08-31 00:12:15 -0700456 for count, name in enumerate(original_names):
Ethan Furmanc16595e2016-09-10 23:36:59 -0700457 value = first_enum._generate_next_value_(name, start, count, last_values[:])
458 last_values.append(value)
459 names.append((name, value))
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700460
461 # Here, names is either an iterable of (name, value) or a mapping.
462 for item in names:
463 if isinstance(item, str):
464 member_name, member_value = item, names[item]
465 else:
466 member_name, member_value = item
467 classdict[member_name] = member_value
468 enum_class = metacls.__new__(metacls, class_name, bases, classdict)
469
470 # TODO: replace the frame hack if a blessed way to know the calling
471 # module is ever developed
472 if module is None:
473 try:
474 module = sys._getframe(2).f_globals['__name__']
Pablo Galindo293dd232019-11-19 21:34:03 +0000475 except (AttributeError, ValueError, KeyError):
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700476 pass
477 if module is None:
478 _make_class_unpicklable(enum_class)
479 else:
480 enum_class.__module__ = module
Ethan Furmanca1b7942014-02-08 11:36:27 -0800481 if qualname is not None:
482 enum_class.__qualname__ = qualname
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700483
484 return enum_class
485
orlnub1230fb9fad2018-09-12 20:28:53 +0300486 def _convert_(cls, name, module, filter, source=None):
487 """
488 Create a new Enum subclass that replaces a collection of global constants
489 """
490 # convert all constants from source (or module) that pass filter() to
491 # a new Enum called name, and export the enum and its members back to
492 # module;
493 # also, replace the __reduce_ex__ method so unpickling works in
494 # previous Python versions
495 module_globals = vars(sys.modules[module])
496 if source:
497 source = vars(source)
498 else:
499 source = module_globals
500 # _value2member_map_ is populated in the same order every time
501 # for a consistent reverse mapping of number to name when there
502 # are multiple names for the same number.
503 members = [
504 (name, value)
505 for name, value in source.items()
506 if filter(name)]
507 try:
508 # sort by value
509 members.sort(key=lambda t: (t[1], t[0]))
510 except TypeError:
511 # unless some values aren't comparable, in which case sort by name
512 members.sort(key=lambda t: t[0])
513 cls = cls(name, members, module=module)
514 cls.__reduce_ex__ = _reduce_ex_by_name
515 module_globals.update(cls.__members__)
516 module_globals[name] = cls
517 return cls
518
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700519 @staticmethod
Ethan Furman3064dbf2020-09-16 07:11:57 -0700520 def _check_for_existing_members(class_name, bases):
521 for chain in bases:
522 for base in chain.__mro__:
523 if issubclass(base, Enum) and base._member_names_:
Ethan Furman6d3dfee2020-12-08 12:26:56 -0800524 raise TypeError(
525 "%s: cannot extend enumeration %r"
526 % (class_name, base.__name__)
527 )
Ethan Furman3064dbf2020-09-16 07:11:57 -0700528
529 @staticmethod
530 def _get_mixins_(class_name, bases):
Ethan Furman6d3dfee2020-12-08 12:26:56 -0800531 """
532 Returns the type for creating enum members, and the first inherited
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700533 enum class.
534
535 bases: the tuple of bases that was given to __new__
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700536 """
537 if not bases:
538 return object, Enum
539
Ethan Furman5bdab642018-09-21 19:03:09 -0700540 def _find_data_type(bases):
Ethan Furmanbff01f32020-09-15 15:56:26 -0700541 data_types = []
Ethan Furman5bdab642018-09-21 19:03:09 -0700542 for chain in bases:
Ethan Furmanbff01f32020-09-15 15:56:26 -0700543 candidate = None
Ethan Furman5bdab642018-09-21 19:03:09 -0700544 for base in chain.__mro__:
545 if base is object:
546 continue
Ethan Furmanc2667362020-12-07 00:17:31 -0800547 elif issubclass(base, Enum):
548 if base._member_type_ is not object:
549 data_types.append(base._member_type_)
550 break
Ethan Furman5bdab642018-09-21 19:03:09 -0700551 elif '__new__' in base.__dict__:
Ethan Furmancd453852018-10-05 23:29:36 -0700552 if issubclass(base, Enum):
Ethan Furman5bdab642018-09-21 19:03:09 -0700553 continue
Ethan Furmanbff01f32020-09-15 15:56:26 -0700554 data_types.append(candidate or base)
555 break
Ethan Furmanc2667362020-12-07 00:17:31 -0800556 else:
Ethan Furmanbff01f32020-09-15 15:56:26 -0700557 candidate = base
558 if len(data_types) > 1:
Ethan Furman3064dbf2020-09-16 07:11:57 -0700559 raise TypeError('%r: too many data types: %r' % (class_name, data_types))
Ethan Furmanbff01f32020-09-15 15:56:26 -0700560 elif data_types:
561 return data_types[0]
562 else:
563 return None
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700564
Ethan Furman5bdab642018-09-21 19:03:09 -0700565 # ensure final parent class is an Enum derivative, find any concrete
566 # data type, and check that Enum has no members
567 first_enum = bases[-1]
568 if not issubclass(first_enum, Enum):
569 raise TypeError("new enumerations should be created as "
570 "`EnumName([mixin_type, ...] [data_type,] enum_type)`")
571 member_type = _find_data_type(bases) or object
572 if first_enum._member_names_:
573 raise TypeError("Cannot extend enumerations")
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700574 return member_type, first_enum
575
576 @staticmethod
577 def _find_new_(classdict, member_type, first_enum):
Ethan Furman6d3dfee2020-12-08 12:26:56 -0800578 """
579 Returns the __new__ to be used for creating the enum members.
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700580
581 classdict: the class dictionary given to __new__
582 member_type: the data type whose __new__ will be used by default
583 first_enum: enumeration to check for an overriding __new__
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700584 """
585 # now find the correct __new__, checking to see of one was defined
586 # by the user; also check earlier enum classes in case a __new__ was
587 # saved as __new_member__
588 __new__ = classdict.get('__new__', None)
589
590 # should __new__ be saved as __new_member__ later?
591 save_new = __new__ is not None
592
593 if __new__ is None:
594 # check all possibles for __new_member__ before falling back to
595 # __new__
596 for method in ('__new_member__', '__new__'):
597 for possible in (member_type, first_enum):
598 target = getattr(possible, method, None)
599 if target not in {
600 None,
601 None.__new__,
602 object.__new__,
603 Enum.__new__,
604 }:
605 __new__ = target
606 break
607 if __new__ is not None:
608 break
609 else:
610 __new__ = object.__new__
611
612 # if a non-object.__new__ is used then whatever value/tuple was
613 # assigned to the enum member name will be passed to __new__ and to the
614 # new enum member's __init__
615 if __new__ is object.__new__:
616 use_args = False
617 else:
618 use_args = True
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700619 return __new__, save_new, use_args
620
621
622class Enum(metaclass=EnumMeta):
Ethan Furman6d3dfee2020-12-08 12:26:56 -0800623 """
624 Generic enumeration.
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700625
626 Derive from this class to define new enumerations.
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700627 """
628 def __new__(cls, value):
629 # all enum instances are actually created during class construction
630 # without calling this method; this method is called by the metaclass'
631 # __call__ (i.e. Color(3) ), and by pickle
632 if type(value) is cls:
Ethan Furman23bb6f42016-11-21 09:22:05 -0800633 # For lookups like Color(Color.RED)
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700634 return value
635 # by-value search for a matching enum member
636 # see if it's in the reverse mapping (for hashable values)
Ethan Furman2aa27322013-07-19 19:35:56 -0700637 try:
Andrew Svetlov34ae04f2018-12-26 20:45:33 +0200638 return cls._value2member_map_[value]
639 except KeyError:
640 # Not found, no need to do long O(n) search
641 pass
Ethan Furman2aa27322013-07-19 19:35:56 -0700642 except TypeError:
643 # not there, now do long search -- O(n) behavior
Ethan Furman520ad572013-07-19 19:47:21 -0700644 for member in cls._member_map_.values():
Ethan Furman0081f232014-09-16 17:31:23 -0700645 if member._value_ == value:
Ethan Furman2aa27322013-07-19 19:35:56 -0700646 return member
Ethan Furmanee47e5c2016-08-31 00:12:15 -0700647 # still not found -- try _missing_ hook
Ethan Furman019f0a02018-09-12 11:43:34 -0700648 try:
649 exc = None
650 result = cls._missing_(value)
651 except Exception as e:
652 exc = e
653 result = None
654 if isinstance(result, cls):
655 return result
656 else:
Walter Dörwald323842c2019-07-18 20:37:13 +0200657 ve_exc = ValueError("%r is not a valid %s" % (value, cls.__qualname__))
Ethan Furman019f0a02018-09-12 11:43:34 -0700658 if result is None and exc is None:
659 raise ve_exc
660 elif exc is None:
661 exc = TypeError(
662 'error in %s._missing_: returned %r instead of None or a valid member'
663 % (cls.__name__, result)
664 )
665 exc.__context__ = ve_exc
666 raise exc
Ethan Furmanee47e5c2016-08-31 00:12:15 -0700667
Ethan Furmanc16595e2016-09-10 23:36:59 -0700668 def _generate_next_value_(name, start, count, last_values):
Ethan Furman6d3dfee2020-12-08 12:26:56 -0800669 """
670 Generate the next value when not given.
671
672 name: the name of the member
673 start: the initial start value or None
674 count: the number of existing members
675 last_value: the last value assigned or None
676 """
Ethan Furmanc16595e2016-09-10 23:36:59 -0700677 for last_value in reversed(last_values):
678 try:
679 return last_value + 1
680 except TypeError:
681 pass
682 else:
Ethan Furmanee47e5c2016-08-31 00:12:15 -0700683 return start
Ethan Furmanc16595e2016-09-10 23:36:59 -0700684
Ethan Furmanee47e5c2016-08-31 00:12:15 -0700685 @classmethod
686 def _missing_(cls, value):
Ethan Furmanc95ad7a2020-09-16 10:26:50 -0700687 return None
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700688
689 def __repr__(self):
690 return "<%s.%s: %r>" % (
Ethan Furman520ad572013-07-19 19:47:21 -0700691 self.__class__.__name__, self._name_, self._value_)
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700692
693 def __str__(self):
Ethan Furman520ad572013-07-19 19:47:21 -0700694 return "%s.%s" % (self.__class__.__name__, self._name_)
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700695
Ethan Furman388a3922013-08-12 06:51:41 -0700696 def __dir__(self):
Ethan Furman6d3dfee2020-12-08 12:26:56 -0800697 """
698 Returns all members and all public methods
699 """
Ethan Furman0ae550b2014-10-14 08:58:32 -0700700 added_behavior = [
701 m
702 for cls in self.__class__.mro()
703 for m in cls.__dict__
Ethan Furman354ecf12015-03-11 08:43:12 -0700704 if m[0] != '_' and m not in self._member_map_
Angelin BOOZ68526fe2020-09-21 15:11:06 +0200705 ] + [m for m in self.__dict__ if m[0] != '_']
Ethan Furmanec5f8eb2014-10-21 13:40:35 -0700706 return (['__class__', '__doc__', '__module__'] + added_behavior)
Ethan Furman388a3922013-08-12 06:51:41 -0700707
Ethan Furmanec15a822013-08-31 19:17:41 -0700708 def __format__(self, format_spec):
Ethan Furman6d3dfee2020-12-08 12:26:56 -0800709 """
710 Returns format using actual value type unless __str__ has been overridden.
711 """
Ethan Furmanec15a822013-08-31 19:17:41 -0700712 # mixed-in Enums should use the mixed-in type's __format__, otherwise
713 # we can get strange results with the Enum name showing up instead of
714 # the value
715
thatneat2f19e822019-07-04 11:28:37 -0700716 # pure Enum branch, or branch with __str__ explicitly overridden
Ethan Furman37440ee2020-12-08 11:14:10 -0800717 str_overridden = type(self).__str__ not in (Enum.__str__, Flag.__str__)
thatneat2f19e822019-07-04 11:28:37 -0700718 if self._member_type_ is object or str_overridden:
Ethan Furmanec15a822013-08-31 19:17:41 -0700719 cls = str
720 val = str(self)
721 # mix-in branch
722 else:
723 cls = self._member_type_
Ethan Furman0081f232014-09-16 17:31:23 -0700724 val = self._value_
Ethan Furmanec15a822013-08-31 19:17:41 -0700725 return cls.__format__(val, format_spec)
726
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700727 def __hash__(self):
Ethan Furman520ad572013-07-19 19:47:21 -0700728 return hash(self._name_)
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700729
Ethan Furmanca1b7942014-02-08 11:36:27 -0800730 def __reduce_ex__(self, proto):
Ethan Furmandc870522014-02-18 12:37:12 -0800731 return self.__class__, (self._value_, )
Ethan Furmanca1b7942014-02-08 11:36:27 -0800732
Ethan Furman33918c12013-09-27 23:02:02 -0700733 # DynamicClassAttribute is used to provide access to the `name` and
734 # `value` properties of enum members while keeping some measure of
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700735 # protection from modification, while still allowing for an enumeration
736 # to have members named `name` and `value`. This works because enumeration
737 # members are not set directly on the enum class -- __getattr__ is
738 # used to look them up.
739
Ethan Furmane03ea372013-09-25 07:14:41 -0700740 @DynamicClassAttribute
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700741 def name(self):
Ethan Furmanc850f342013-09-15 16:59:35 -0700742 """The name of the Enum member."""
Ethan Furman520ad572013-07-19 19:47:21 -0700743 return self._name_
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700744
Ethan Furmane03ea372013-09-25 07:14:41 -0700745 @DynamicClassAttribute
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700746 def value(self):
Ethan Furmanc850f342013-09-15 16:59:35 -0700747 """The value of the Enum member."""
Ethan Furman520ad572013-07-19 19:47:21 -0700748 return self._value_
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700749
750
751class IntEnum(int, Enum):
Ethan Furman0063ff42020-09-21 17:23:13 -0700752 """
753 Enum where members are also (and must be) ints
754 """
755
756
757class StrEnum(str, Enum):
758 """
759 Enum where members are also (and must be) strings
760 """
761
762 def __new__(cls, *values):
763 if len(values) > 3:
764 raise TypeError('too many arguments for str(): %r' % (values, ))
765 if len(values) == 1:
766 # it must be a string
767 if not isinstance(values[0], str):
768 raise TypeError('%r is not a string' % (values[0], ))
769 if len(values) > 1:
770 # check that encoding argument is a string
771 if not isinstance(values[1], str):
772 raise TypeError('encoding must be a string, not %r' % (values[1], ))
773 if len(values) > 2:
774 # check that errors argument is a string
775 if not isinstance(values[2], str):
776 raise TypeError('errors must be a string, not %r' % (values[2], ))
777 value = str(*values)
778 member = str.__new__(cls, value)
779 member._value_ = value
780 return member
Ethan Furmanf24bb352013-07-18 17:05:39 -0700781
Ethan Furmand986d162020-09-22 13:00:07 -0700782 __str__ = str.__str__
783
Ethan Furmanf24bb352013-07-18 17:05:39 -0700784
Ethan Furman24e837f2015-03-18 17:27:57 -0700785def _reduce_ex_by_name(self, proto):
786 return self.name
787
Ethan Furman65a5a472016-09-01 23:55:19 -0700788class Flag(Enum):
Ethan Furman6d3dfee2020-12-08 12:26:56 -0800789 """
790 Support for flags
791 """
Ethan Furmanc16595e2016-09-10 23:36:59 -0700792
793 def _generate_next_value_(name, start, count, last_values):
Ethan Furmanee47e5c2016-08-31 00:12:15 -0700794 """
795 Generate the next value when not given.
796
797 name: the name of the member
HongWeipengbb16fb22019-09-21 13:22:54 +0800798 start: the initial start value or None
Ethan Furmanee47e5c2016-08-31 00:12:15 -0700799 count: the number of existing members
800 last_value: the last value assigned or None
801 """
802 if not count:
803 return start if start is not None else 1
Ethan Furmanc16595e2016-09-10 23:36:59 -0700804 for last_value in reversed(last_values):
805 try:
806 high_bit = _high_bit(last_value)
807 break
Ethan Furman3515dcc2016-09-18 13:15:41 -0700808 except Exception:
Ethan Furmanc16595e2016-09-10 23:36:59 -0700809 raise TypeError('Invalid Flag value: %r' % last_value) from None
Ethan Furmanee47e5c2016-08-31 00:12:15 -0700810 return 2 ** (high_bit+1)
811
812 @classmethod
813 def _missing_(cls, value):
Ethan Furman6d3dfee2020-12-08 12:26:56 -0800814 """
815 Returns member (possibly creating it) if one can be found for value.
816 """
Ethan Furmanee47e5c2016-08-31 00:12:15 -0700817 original_value = value
818 if value < 0:
819 value = ~value
820 possible_member = cls._create_pseudo_member_(value)
Ethan Furmanee47e5c2016-08-31 00:12:15 -0700821 if original_value < 0:
822 possible_member = ~possible_member
823 return possible_member
824
825 @classmethod
826 def _create_pseudo_member_(cls, value):
Ethan Furman3515dcc2016-09-18 13:15:41 -0700827 """
828 Create a composite member iff value contains only members.
829 """
Ethan Furmanee47e5c2016-08-31 00:12:15 -0700830 pseudo_member = cls._value2member_map_.get(value, None)
831 if pseudo_member is None:
Ethan Furman3515dcc2016-09-18 13:15:41 -0700832 # verify all bits are accounted for
833 _, extra_flags = _decompose(cls, value)
834 if extra_flags:
Walter Dörwald323842c2019-07-18 20:37:13 +0200835 raise ValueError("%r is not a valid %s" % (value, cls.__qualname__))
Ethan Furman3515dcc2016-09-18 13:15:41 -0700836 # construct a singleton enum pseudo-member
Ethan Furmanee47e5c2016-08-31 00:12:15 -0700837 pseudo_member = object.__new__(cls)
838 pseudo_member._name_ = None
839 pseudo_member._value_ = value
Ethan Furman28cf6632017-01-24 12:12:06 -0800840 # use setdefault in case another thread already created a composite
841 # with this value
842 pseudo_member = cls._value2member_map_.setdefault(value, pseudo_member)
Ethan Furmanee47e5c2016-08-31 00:12:15 -0700843 return pseudo_member
844
Ethan Furmanee47e5c2016-08-31 00:12:15 -0700845 def __contains__(self, other):
Ethan Furman6d3dfee2020-12-08 12:26:56 -0800846 """
847 Returns True if self has at least the same flags set as other.
848 """
Ethan Furmanee47e5c2016-08-31 00:12:15 -0700849 if not isinstance(other, self.__class__):
Rahul Jha94306522018-09-10 23:51:04 +0530850 raise TypeError(
851 "unsupported operand type(s) for 'in': '%s' and '%s'" % (
852 type(other).__qualname__, self.__class__.__qualname__))
Ethan Furmanee47e5c2016-08-31 00:12:15 -0700853 return other._value_ & self._value_ == other._value_
854
Ethan Furman7219e272020-09-16 13:01:00 -0700855 def __iter__(self):
Ethan Furman6d3dfee2020-12-08 12:26:56 -0800856 """
857 Returns flags in decreasing value order.
858 """
Ethan Furman7219e272020-09-16 13:01:00 -0700859 members, extra_flags = _decompose(self.__class__, self.value)
860 return (m for m in members if m._value_ != 0)
861
Ethan Furmanee47e5c2016-08-31 00:12:15 -0700862 def __repr__(self):
863 cls = self.__class__
864 if self._name_ is not None:
865 return '<%s.%s: %r>' % (cls.__name__, self._name_, self._value_)
Ethan Furman3515dcc2016-09-18 13:15:41 -0700866 members, uncovered = _decompose(cls, self._value_)
Ethan Furman27682d22016-09-04 11:39:01 -0700867 return '<%s.%s: %r>' % (
868 cls.__name__,
869 '|'.join([str(m._name_ or m._value_) for m in members]),
870 self._value_,
871 )
Ethan Furmanee47e5c2016-08-31 00:12:15 -0700872
873 def __str__(self):
874 cls = self.__class__
875 if self._name_ is not None:
876 return '%s.%s' % (cls.__name__, self._name_)
Ethan Furman3515dcc2016-09-18 13:15:41 -0700877 members, uncovered = _decompose(cls, self._value_)
Ethan Furmanee47e5c2016-08-31 00:12:15 -0700878 if len(members) == 1 and members[0]._name_ is None:
879 return '%s.%r' % (cls.__name__, members[0]._value_)
880 else:
881 return '%s.%s' % (
882 cls.__name__,
883 '|'.join([str(m._name_ or m._value_) for m in members]),
884 )
885
Ethan Furman25d94bb2016-09-02 16:32:32 -0700886 def __bool__(self):
887 return bool(self._value_)
888
Ethan Furmanee47e5c2016-08-31 00:12:15 -0700889 def __or__(self, other):
890 if not isinstance(other, self.__class__):
891 return NotImplemented
892 return self.__class__(self._value_ | other._value_)
893
894 def __and__(self, other):
895 if not isinstance(other, self.__class__):
896 return NotImplemented
897 return self.__class__(self._value_ & other._value_)
898
899 def __xor__(self, other):
900 if not isinstance(other, self.__class__):
901 return NotImplemented
902 return self.__class__(self._value_ ^ other._value_)
903
904 def __invert__(self):
Ethan Furman3515dcc2016-09-18 13:15:41 -0700905 members, uncovered = _decompose(self.__class__, self._value_)
Serhiy Storchaka81108372017-09-26 00:55:55 +0300906 inverted = self.__class__(0)
907 for m in self.__class__:
908 if m not in members and not (m._value_ & self._value_):
909 inverted = inverted | m
Ethan Furmanee47e5c2016-08-31 00:12:15 -0700910 return self.__class__(inverted)
911
912
Ethan Furman65a5a472016-09-01 23:55:19 -0700913class IntFlag(int, Flag):
Ethan Furman6d3dfee2020-12-08 12:26:56 -0800914 """
915 Support for integer-based Flags
916 """
Ethan Furmanee47e5c2016-08-31 00:12:15 -0700917
918 @classmethod
Ethan Furman3515dcc2016-09-18 13:15:41 -0700919 def _missing_(cls, value):
Ethan Furman6d3dfee2020-12-08 12:26:56 -0800920 """
921 Returns member (possibly creating it) if one can be found for value.
922 """
Ethan Furman3515dcc2016-09-18 13:15:41 -0700923 if not isinstance(value, int):
Walter Dörwald323842c2019-07-18 20:37:13 +0200924 raise ValueError("%r is not a valid %s" % (value, cls.__qualname__))
Ethan Furman3515dcc2016-09-18 13:15:41 -0700925 new_member = cls._create_pseudo_member_(value)
926 return new_member
927
928 @classmethod
Ethan Furmanee47e5c2016-08-31 00:12:15 -0700929 def _create_pseudo_member_(cls, value):
Ethan Furman6d3dfee2020-12-08 12:26:56 -0800930 """
931 Create a composite member iff value contains only members.
932 """
Ethan Furmanee47e5c2016-08-31 00:12:15 -0700933 pseudo_member = cls._value2member_map_.get(value, None)
934 if pseudo_member is None:
Ethan Furman3515dcc2016-09-18 13:15:41 -0700935 need_to_create = [value]
936 # get unaccounted for bits
937 _, extra_flags = _decompose(cls, value)
938 # timer = 10
939 while extra_flags:
940 # timer -= 1
941 bit = _high_bit(extra_flags)
942 flag_value = 2 ** bit
943 if (flag_value not in cls._value2member_map_ and
944 flag_value not in need_to_create
945 ):
946 need_to_create.append(flag_value)
947 if extra_flags == -flag_value:
948 extra_flags = 0
949 else:
950 extra_flags ^= flag_value
951 for value in reversed(need_to_create):
952 # construct singleton pseudo-members
953 pseudo_member = int.__new__(cls, value)
954 pseudo_member._name_ = None
955 pseudo_member._value_ = value
Ethan Furman28cf6632017-01-24 12:12:06 -0800956 # use setdefault in case another thread already created a composite
957 # with this value
958 pseudo_member = cls._value2member_map_.setdefault(value, pseudo_member)
Ethan Furmanee47e5c2016-08-31 00:12:15 -0700959 return pseudo_member
960
Ethan Furmanee47e5c2016-08-31 00:12:15 -0700961 def __or__(self, other):
962 if not isinstance(other, (self.__class__, int)):
963 return NotImplemented
Ethan Furman3515dcc2016-09-18 13:15:41 -0700964 result = self.__class__(self._value_ | self.__class__(other)._value_)
965 return result
Ethan Furmanee47e5c2016-08-31 00:12:15 -0700966
967 def __and__(self, other):
968 if not isinstance(other, (self.__class__, int)):
969 return NotImplemented
970 return self.__class__(self._value_ & self.__class__(other)._value_)
971
972 def __xor__(self, other):
973 if not isinstance(other, (self.__class__, int)):
974 return NotImplemented
975 return self.__class__(self._value_ ^ self.__class__(other)._value_)
976
977 __ror__ = __or__
978 __rand__ = __and__
979 __rxor__ = __xor__
980
981 def __invert__(self):
Ethan Furman3515dcc2016-09-18 13:15:41 -0700982 result = self.__class__(~self._value_)
983 return result
Ethan Furmanee47e5c2016-08-31 00:12:15 -0700984
985
986def _high_bit(value):
Ethan Furman6d3dfee2020-12-08 12:26:56 -0800987 """
988 returns index of highest bit, or -1 if value is zero or negative
989 """
Ethan Furman3515dcc2016-09-18 13:15:41 -0700990 return value.bit_length() - 1
Ethan Furmanee47e5c2016-08-31 00:12:15 -0700991
Ethan Furmanf24bb352013-07-18 17:05:39 -0700992def unique(enumeration):
Ethan Furman6d3dfee2020-12-08 12:26:56 -0800993 """
994 Class decorator for enumerations ensuring unique member values.
995 """
Ethan Furmanf24bb352013-07-18 17:05:39 -0700996 duplicates = []
997 for name, member in enumeration.__members__.items():
998 if name != member.name:
999 duplicates.append((name, member.name))
1000 if duplicates:
1001 alias_details = ', '.join(
1002 ["%s -> %s" % (alias, name) for (alias, name) in duplicates])
1003 raise ValueError('duplicate values found in %r: %s' %
1004 (enumeration, alias_details))
1005 return enumeration
Ethan Furman3515dcc2016-09-18 13:15:41 -07001006
1007def _decompose(flag, value):
Ethan Furman6d3dfee2020-12-08 12:26:56 -08001008 """
1009 Extract all members from the value.
1010 """
Ethan Furman3515dcc2016-09-18 13:15:41 -07001011 # _decompose is only called if the value is not named
1012 not_covered = value
1013 negative = value < 0
Ethan Furman3515dcc2016-09-18 13:15:41 -07001014 members = []
HongWeipeng0b41a922019-11-27 06:36:02 +08001015 for member in flag:
1016 member_value = member.value
Ethan Furman3515dcc2016-09-18 13:15:41 -07001017 if member_value and member_value & value == member_value:
1018 members.append(member)
1019 not_covered &= ~member_value
HongWeipeng0b41a922019-11-27 06:36:02 +08001020 if not negative:
1021 tmp = not_covered
1022 while tmp:
1023 flag_value = 2 ** _high_bit(tmp)
1024 if flag_value in flag._value2member_map_:
1025 members.append(flag._value2member_map_[flag_value])
1026 not_covered &= ~flag_value
1027 tmp &= ~flag_value
Ethan Furman3515dcc2016-09-18 13:15:41 -07001028 if not members and value in flag._value2member_map_:
1029 members.append(flag._value2member_map_[value])
1030 members.sort(key=lambda m: m._value_, reverse=True)
1031 if len(members) > 1 and members[0].value == value:
1032 # we have the breakdown, don't need the value member itself
1033 members.pop(0)
1034 return members, not_covered