blob: 8711558448107fe117841e844e1b27d26215c338 [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',
7 'Enum', 'IntEnum', 'Flag', 'IntFlag',
8 'auto', 'unique',
9 ]
Ethan Furman6b3d64a2013-06-14 16:55:46 -070010
11
Ethan Furman101e0742013-09-15 12:34:36 -070012def _is_descriptor(obj):
13 """Returns True if obj is a descriptor, False otherwise."""
14 return (
15 hasattr(obj, '__get__') or
16 hasattr(obj, '__set__') or
17 hasattr(obj, '__delete__'))
18
19
Ethan Furman6b3d64a2013-06-14 16:55:46 -070020def _is_dunder(name):
21 """Returns True if a __dunder__ name, False otherwise."""
Brennan D Baraban8b914d22019-03-03 14:09:11 -080022 return (len(name) > 4 and
23 name[:2] == name[-2:] == '__' and
24 name[2] != '_' and
25 name[-3] != '_')
Ethan Furman6b3d64a2013-06-14 16:55:46 -070026
27
28def _is_sunder(name):
29 """Returns True if a _sunder_ name, False otherwise."""
Brennan D Baraban8b914d22019-03-03 14:09:11 -080030 return (len(name) > 2 and
31 name[0] == name[-1] == '_' and
Ethan Furman6b3d64a2013-06-14 16:55:46 -070032 name[1:2] != '_' and
Brennan D Baraban8b914d22019-03-03 14:09:11 -080033 name[-2:-1] != '_')
34
Ethan Furman6b3d64a2013-06-14 16:55:46 -070035
Ethan Furman6b3d64a2013-06-14 16:55:46 -070036def _make_class_unpicklable(cls):
37 """Make the given class un-picklable."""
Ethan Furmanca1b7942014-02-08 11:36:27 -080038 def _break_on_call_reduce(self, proto):
Ethan Furman6b3d64a2013-06-14 16:55:46 -070039 raise TypeError('%r cannot be pickled' % self)
Ethan Furmanca1b7942014-02-08 11:36:27 -080040 cls.__reduce_ex__ = _break_on_call_reduce
Ethan Furman6b3d64a2013-06-14 16:55:46 -070041 cls.__module__ = '<unknown>'
42
Ethan Furman3515dcc2016-09-18 13:15:41 -070043_auto_null = object()
Ethan Furmanc16595e2016-09-10 23:36:59 -070044class auto:
45 """
46 Instances are replaced with an appropriate value in Enum class suites.
47 """
Ethan Furman3515dcc2016-09-18 13:15:41 -070048 value = _auto_null
Ethan Furmanc16595e2016-09-10 23:36:59 -070049
Ethan Furman101e0742013-09-15 12:34:36 -070050
Ethan Furman6b3d64a2013-06-14 16:55:46 -070051class _EnumDict(dict):
Ethan Furman101e0742013-09-15 12:34:36 -070052 """Track enum member order and ensure member names are not reused.
Ethan Furman6b3d64a2013-06-14 16:55:46 -070053
54 EnumMeta will use the names found in self._member_names as the
55 enumeration member names.
56
57 """
58 def __init__(self):
59 super().__init__()
60 self._member_names = []
Ethan Furmanc16595e2016-09-10 23:36:59 -070061 self._last_values = []
Ethan Furmana4b1bb42018-01-22 07:56:37 -080062 self._ignore = []
Ethan Onstottd9a43e22020-04-28 13:20:55 -040063 self._auto_called = False
Ethan Furman6b3d64a2013-06-14 16:55:46 -070064
65 def __setitem__(self, key, value):
Ethan Furman101e0742013-09-15 12:34:36 -070066 """Changes anything not dundered or not a descriptor.
Ethan Furman6b3d64a2013-06-14 16:55:46 -070067
68 If an enum member name is used twice, an error is raised; duplicate
69 values are not checked for.
70
71 Single underscore (sunder) names are reserved.
72
73 """
74 if _is_sunder(key):
Ethan Furmanee47e5c2016-08-31 00:12:15 -070075 if key not in (
Ethan Furman3515dcc2016-09-18 13:15:41 -070076 '_order_', '_create_pseudo_member_',
Ethan Furmana4b1bb42018-01-22 07:56:37 -080077 '_generate_next_value_', '_missing_', '_ignore_',
Ethan Furmanee47e5c2016-08-31 00:12:15 -070078 ):
Ethan Furmane8e61272016-08-20 07:19:31 -070079 raise ValueError('_names_ are reserved for future Enum use')
Ethan Furmanc16595e2016-09-10 23:36:59 -070080 if key == '_generate_next_value_':
Ethan Onstottd9a43e22020-04-28 13:20:55 -040081 # check if members already defined as auto()
82 if self._auto_called:
83 raise TypeError("_generate_next_value_ must be defined before members")
Ethan Furmanc16595e2016-09-10 23:36:59 -070084 setattr(self, '_generate_next_value', value)
Ethan Furmana4b1bb42018-01-22 07:56:37 -080085 elif key == '_ignore_':
86 if isinstance(value, str):
87 value = value.replace(',',' ').split()
88 else:
89 value = list(value)
90 self._ignore = value
91 already = set(value) & set(self._member_names)
92 if already:
93 raise ValueError('_ignore_ cannot specify already set names: %r' % (already, ))
Ethan Furman101e0742013-09-15 12:34:36 -070094 elif _is_dunder(key):
Ethan Furmane8e61272016-08-20 07:19:31 -070095 if key == '__order__':
96 key = '_order_'
Ethan Furman101e0742013-09-15 12:34:36 -070097 elif key in self._member_names:
98 # descriptor overwriting an enum?
99 raise TypeError('Attempted to reuse key: %r' % key)
Ethan Furmana4b1bb42018-01-22 07:56:37 -0800100 elif key in self._ignore:
101 pass
Ethan Furman101e0742013-09-15 12:34:36 -0700102 elif not _is_descriptor(value):
103 if key in self:
104 # enum overwriting a descriptor?
Ethan Furmanee47e5c2016-08-31 00:12:15 -0700105 raise TypeError('%r already defined as: %r' % (key, self[key]))
Ethan Furmanc16595e2016-09-10 23:36:59 -0700106 if isinstance(value, auto):
Ethan Onstottd9a43e22020-04-28 13:20:55 -0400107 self._auto_called = True
Ethan Furman3515dcc2016-09-18 13:15:41 -0700108 if value.value == _auto_null:
109 value.value = self._generate_next_value(key, 1, len(self._member_names), self._last_values[:])
110 value = value.value
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700111 self._member_names.append(key)
Ethan Furmanc16595e2016-09-10 23:36:59 -0700112 self._last_values.append(value)
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700113 super().__setitem__(key, value)
114
115
Ezio Melotti9a3777e2013-08-17 15:53:55 +0300116# Dummy value for Enum as EnumMeta explicitly checks for it, but of course
117# until EnumMeta finishes running the first time the Enum class doesn't exist.
118# This is also why there are checks in EnumMeta like `if Enum is not None`
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700119Enum = None
120
Ethan Furman332dbc72016-08-20 00:00:52 -0700121
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700122class EnumMeta(type):
123 """Metaclass for Enum"""
124 @classmethod
Ethan Furman332dbc72016-08-20 00:00:52 -0700125 def __prepare__(metacls, cls, bases):
Ethan Furmanee47e5c2016-08-31 00:12:15 -0700126 # create the namespace dict
127 enum_dict = _EnumDict()
128 # inherit previous flags and _generate_next_value_ function
129 member_type, first_enum = metacls._get_mixins_(bases)
130 if first_enum is not None:
131 enum_dict['_generate_next_value_'] = getattr(first_enum, '_generate_next_value_', None)
132 return enum_dict
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700133
Ethan Furman65a5a472016-09-01 23:55:19 -0700134 def __new__(metacls, cls, bases, classdict):
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700135 # an Enum class is final once enumeration items have been defined; it
136 # cannot be mixed with other types (int, float, etc.) if it has an
137 # inherited __new__ unless a new __new__ is defined (or the resulting
138 # class will fail).
Ethan Furmana4b1bb42018-01-22 07:56:37 -0800139 #
140 # remove any keys listed in _ignore_
141 classdict.setdefault('_ignore_', []).append('_ignore_')
142 ignore = classdict['_ignore_']
143 for key in ignore:
144 classdict.pop(key, None)
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700145 member_type, first_enum = metacls._get_mixins_(bases)
146 __new__, save_new, use_args = metacls._find_new_(classdict, member_type,
147 first_enum)
148
149 # save enum items into separate mapping so they don't get baked into
150 # the new class
Ethan Furmanee47e5c2016-08-31 00:12:15 -0700151 enum_members = {k: classdict[k] for k in classdict._member_names}
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700152 for name in classdict._member_names:
153 del classdict[name]
154
Ethan Furmane8e61272016-08-20 07:19:31 -0700155 # adjust the sunders
156 _order_ = classdict.pop('_order_', None)
157
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700158 # check for illegal enum names (any others?)
Brennan D Baraban8b914d22019-03-03 14:09:11 -0800159 invalid_names = set(enum_members) & {'mro', ''}
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700160 if invalid_names:
161 raise ValueError('Invalid enum member name: {0}'.format(
162 ','.join(invalid_names)))
163
Ethan Furman48a724f2015-04-11 23:23:06 -0700164 # create a default docstring if one has not been provided
165 if '__doc__' not in classdict:
166 classdict['__doc__'] = 'An enumeration.'
167
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700168 # create our new Enum type
169 enum_class = super().__new__(metacls, cls, bases, classdict)
Ethan Furman520ad572013-07-19 19:47:21 -0700170 enum_class._member_names_ = [] # names in definition order
INADA Naokie57f91a2018-06-19 01:14:26 +0900171 enum_class._member_map_ = {} # name->value map
Ethan Furman5e5a8232013-08-04 08:42:23 -0700172 enum_class._member_type_ = member_type
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700173
orlnub1230fb9fad2018-09-12 20:28:53 +0300174 # save DynamicClassAttribute attributes from super classes so we know
175 # if we can take the shortcut of storing members in the class dict
176 dynamic_attributes = {k for c in enum_class.mro()
177 for k, v in c.__dict__.items()
178 if isinstance(v, DynamicClassAttribute)}
Ethan Furman354ecf12015-03-11 08:43:12 -0700179
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700180 # Reverse value->name map for hashable values.
Ethan Furman520ad572013-07-19 19:47:21 -0700181 enum_class._value2member_map_ = {}
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700182
Ethan Furman2da95042014-03-03 12:42:52 -0800183 # If a custom type is mixed into the Enum, and it does not know how
184 # to pickle itself, pickle.dumps will succeed but pickle.loads will
185 # fail. Rather than have the error show up later and possibly far
186 # from the source, sabotage the pickle protocol for this class so
187 # that pickle.dumps also fails.
188 #
189 # However, if the new class implements its own __reduce_ex__, do not
190 # sabotage -- it's on them to make sure it works correctly. We use
191 # __reduce_ex__ instead of any of the others as it is preferred by
192 # pickle over __reduce__, and it handles all pickle protocols.
193 if '__reduce_ex__' not in classdict:
Ethan Furmandc870522014-02-18 12:37:12 -0800194 if member_type is not object:
195 methods = ('__getnewargs_ex__', '__getnewargs__',
196 '__reduce_ex__', '__reduce__')
Ethan Furman2da95042014-03-03 12:42:52 -0800197 if not any(m in member_type.__dict__ for m in methods):
Ethan Furmandc870522014-02-18 12:37:12 -0800198 _make_class_unpicklable(enum_class)
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700199
200 # instantiate them, checking for duplicates as we go
201 # we instantiate first instead of checking for duplicates first in case
202 # a custom __new__ is doing something funky with the values -- such as
203 # auto-numbering ;)
204 for member_name in classdict._member_names:
Ethan Furmanee47e5c2016-08-31 00:12:15 -0700205 value = enum_members[member_name]
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700206 if not isinstance(value, tuple):
207 args = (value, )
208 else:
209 args = value
210 if member_type is tuple: # special case for tuple enums
211 args = (args, ) # wrap it one more time
212 if not use_args:
213 enum_member = __new__(enum_class)
Ethan Furmanb41803e2013-07-25 13:50:45 -0700214 if not hasattr(enum_member, '_value_'):
215 enum_member._value_ = value
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700216 else:
217 enum_member = __new__(enum_class, *args)
Ethan Furmanb41803e2013-07-25 13:50:45 -0700218 if not hasattr(enum_member, '_value_'):
Ethan Furmanee47e5c2016-08-31 00:12:15 -0700219 if member_type is object:
220 enum_member._value_ = value
221 else:
222 enum_member._value_ = member_type(*args)
Ethan Furman520ad572013-07-19 19:47:21 -0700223 value = enum_member._value_
Ethan Furman520ad572013-07-19 19:47:21 -0700224 enum_member._name_ = member_name
Ethan Furmanc850f342013-09-15 16:59:35 -0700225 enum_member.__objclass__ = enum_class
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700226 enum_member.__init__(*args)
227 # If another member with the same value was already defined, the
228 # new member becomes an alias to the existing one.
Ethan Furman520ad572013-07-19 19:47:21 -0700229 for name, canonical_member in enum_class._member_map_.items():
Ethan Furman0081f232014-09-16 17:31:23 -0700230 if canonical_member._value_ == enum_member._value_:
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700231 enum_member = canonical_member
232 break
233 else:
234 # Aliases don't appear in member names (only in __members__).
Ethan Furman520ad572013-07-19 19:47:21 -0700235 enum_class._member_names_.append(member_name)
Ethan Furman354ecf12015-03-11 08:43:12 -0700236 # performance boost for any member that would not shadow
237 # a DynamicClassAttribute
orlnub1230fb9fad2018-09-12 20:28:53 +0300238 if member_name not in dynamic_attributes:
Ethan Furman354ecf12015-03-11 08:43:12 -0700239 setattr(enum_class, member_name, enum_member)
240 # now add to _member_map_
Ethan Furman520ad572013-07-19 19:47:21 -0700241 enum_class._member_map_[member_name] = enum_member
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700242 try:
243 # This may fail if value is not hashable. We can't add the value
244 # to the map, and by-value lookups for this value will be
245 # linear.
Ethan Furman520ad572013-07-19 19:47:21 -0700246 enum_class._value2member_map_[value] = enum_member
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700247 except TypeError:
248 pass
249
250 # double check that repr and friends are not the mixin's or various
251 # things break (such as pickle)
Ethan Furmana4677062020-09-16 03:58:33 -0700252 # however, if the method is defined in the Enum itself, don't replace
253 # it
Ethan Furmandc870522014-02-18 12:37:12 -0800254 for name in ('__repr__', '__str__', '__format__', '__reduce_ex__'):
Ethan Furmana4677062020-09-16 03:58:33 -0700255 if name in classdict:
256 continue
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700257 class_method = getattr(enum_class, name)
258 obj_method = getattr(member_type, name, None)
259 enum_method = getattr(first_enum, name, None)
260 if obj_method is not None and obj_method is class_method:
261 setattr(enum_class, name, enum_method)
262
263 # replace any other __new__ with our own (as long as Enum is not None,
264 # anyway) -- again, this is to support pickle
265 if Enum is not None:
266 # if the user defined their own __new__, save it before it gets
267 # clobbered in case they subclass later
268 if save_new:
269 enum_class.__new_member__ = __new__
270 enum_class.__new__ = Enum.__new__
Ethan Furmane8e61272016-08-20 07:19:31 -0700271
272 # py3 support for definition order (helps keep py2/py3 code in sync)
273 if _order_ is not None:
274 if isinstance(_order_, str):
275 _order_ = _order_.replace(',', ' ').split()
276 if _order_ != enum_class._member_names_:
277 raise TypeError('member order does not match _order_')
278
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700279 return enum_class
280
Ethan Furman5de67b12016-04-13 23:52:09 -0700281 def __bool__(self):
282 """
283 classes/types should always be True.
284 """
285 return True
286
Ethan Furmand9925a12014-09-16 20:35:55 -0700287 def __call__(cls, value, names=None, *, module=None, qualname=None, type=None, start=1):
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700288 """Either returns an existing member, or creates a new enum class.
289
290 This method is used both when an enum class is given a value to match
291 to an enumeration member (i.e. Color(3)) and for the functional API
Ethan Furman23bb6f42016-11-21 09:22:05 -0800292 (i.e. Color = Enum('Color', names='RED GREEN BLUE')).
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700293
Ethan Furman2da95042014-03-03 12:42:52 -0800294 When used for the functional API:
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700295
Ethan Furman2da95042014-03-03 12:42:52 -0800296 `value` will be the name of the new class.
297
298 `names` should be either a string of white-space/comma delimited names
Ethan Furmand9925a12014-09-16 20:35:55 -0700299 (values will start at `start`), or an iterator/mapping of name, value pairs.
Ethan Furman2da95042014-03-03 12:42:52 -0800300
301 `module` should be set to the module this class is being created in;
302 if it is not set, an attempt to find that module will be made, but if
303 it fails the class will not be picklable.
304
305 `qualname` should be set to the actual location this class can be found
306 at in its module; by default it is set to the global scope. If this is
307 not correct, unpickling will fail in some circumstances.
308
309 `type`, if set, will be mixed in as the first base class.
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700310
311 """
312 if names is None: # simple value lookup
313 return cls.__new__(cls, value)
314 # otherwise, functional API: we're creating a new Enum type
Ethan Furmand9925a12014-09-16 20:35:55 -0700315 return cls._create_(value, names, module=module, qualname=qualname, type=type, start=start)
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700316
317 def __contains__(cls, member):
Rahul Jha94306522018-09-10 23:51:04 +0530318 if not isinstance(member, Enum):
319 raise TypeError(
320 "unsupported operand type(s) for 'in': '%s' and '%s'" % (
321 type(member).__qualname__, cls.__class__.__qualname__))
Ethan Furman0081f232014-09-16 17:31:23 -0700322 return isinstance(member, cls) and member._name_ in cls._member_map_
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700323
Ethan Furman64a99722013-09-22 16:18:19 -0700324 def __delattr__(cls, attr):
325 # nicer error message when someone tries to delete an attribute
326 # (see issue19025).
327 if attr in cls._member_map_:
328 raise AttributeError(
329 "%s: cannot delete Enum member." % cls.__name__)
330 super().__delattr__(attr)
331
Ethan Furman388a3922013-08-12 06:51:41 -0700332 def __dir__(self):
Ethan Furman64a99722013-09-22 16:18:19 -0700333 return (['__class__', '__doc__', '__members__', '__module__'] +
334 self._member_names_)
Ethan Furman388a3922013-08-12 06:51:41 -0700335
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700336 def __getattr__(cls, name):
337 """Return the enum member matching `name`
338
339 We use __getattr__ instead of descriptors or inserting into the enum
340 class' __dict__ in order to support `name` and `value` being both
341 properties for enum members (which live in the class' __dict__) and
342 enum members themselves.
343
344 """
345 if _is_dunder(name):
346 raise AttributeError(name)
347 try:
Ethan Furman520ad572013-07-19 19:47:21 -0700348 return cls._member_map_[name]
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700349 except KeyError:
350 raise AttributeError(name) from None
351
352 def __getitem__(cls, name):
Ethan Furman520ad572013-07-19 19:47:21 -0700353 return cls._member_map_[name]
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700354
355 def __iter__(cls):
Ethan Furman520ad572013-07-19 19:47:21 -0700356 return (cls._member_map_[name] for name in cls._member_names_)
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700357
358 def __len__(cls):
Ethan Furman520ad572013-07-19 19:47:21 -0700359 return len(cls._member_names_)
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700360
Ethan Furman2131a4a2013-09-14 18:11:24 -0700361 @property
362 def __members__(cls):
363 """Returns a mapping of member name->value.
364
365 This mapping lists all enum members, including aliases. Note that this
366 is a read-only view of the internal mapping.
367
368 """
369 return MappingProxyType(cls._member_map_)
370
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700371 def __repr__(cls):
372 return "<enum %r>" % cls.__name__
373
Ethan Furman2131a4a2013-09-14 18:11:24 -0700374 def __reversed__(cls):
375 return (cls._member_map_[name] for name in reversed(cls._member_names_))
376
Ethan Furmanf203f2d2013-09-06 07:16:48 -0700377 def __setattr__(cls, name, value):
378 """Block attempts to reassign Enum members.
379
380 A simple assignment to the class namespace only changes one of the
381 several possible ways to get an Enum member from the Enum class,
382 resulting in an inconsistent Enumeration.
383
384 """
385 member_map = cls.__dict__.get('_member_map_', {})
386 if name in member_map:
387 raise AttributeError('Cannot reassign members.')
388 super().__setattr__(name, value)
389
anentropicb8e21f12018-04-16 04:40:35 +0100390 def _create_(cls, class_name, names, *, module=None, qualname=None, type=None, start=1):
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700391 """Convenience method to create a new Enum class.
392
393 `names` can be:
394
395 * A string containing member names, separated either with spaces or
Ethan Furmand9925a12014-09-16 20:35:55 -0700396 commas. Values are incremented by 1 from `start`.
397 * An iterable of member names. Values are incremented by 1 from `start`.
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700398 * An iterable of (member name, value) pairs.
Ethan Furmand9925a12014-09-16 20:35:55 -0700399 * A mapping of member name -> value pairs.
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700400
401 """
402 metacls = cls.__class__
403 bases = (cls, ) if type is None else (type, cls)
Ethan Furmanee47e5c2016-08-31 00:12:15 -0700404 _, first_enum = cls._get_mixins_(bases)
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700405 classdict = metacls.__prepare__(class_name, bases)
406
407 # special processing needed for names?
408 if isinstance(names, str):
409 names = names.replace(',', ' ').split()
Dong-hee Nadcc8ce42017-06-22 01:52:32 +0900410 if isinstance(names, (tuple, list)) and names and isinstance(names[0], str):
Ethan Furmanee47e5c2016-08-31 00:12:15 -0700411 original_names, names = names, []
Ethan Furmanc16595e2016-09-10 23:36:59 -0700412 last_values = []
Ethan Furmanee47e5c2016-08-31 00:12:15 -0700413 for count, name in enumerate(original_names):
Ethan Furmanc16595e2016-09-10 23:36:59 -0700414 value = first_enum._generate_next_value_(name, start, count, last_values[:])
415 last_values.append(value)
416 names.append((name, value))
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700417
418 # Here, names is either an iterable of (name, value) or a mapping.
419 for item in names:
420 if isinstance(item, str):
421 member_name, member_value = item, names[item]
422 else:
423 member_name, member_value = item
424 classdict[member_name] = member_value
425 enum_class = metacls.__new__(metacls, class_name, bases, classdict)
426
427 # TODO: replace the frame hack if a blessed way to know the calling
428 # module is ever developed
429 if module is None:
430 try:
431 module = sys._getframe(2).f_globals['__name__']
Pablo Galindo293dd232019-11-19 21:34:03 +0000432 except (AttributeError, ValueError, KeyError):
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700433 pass
434 if module is None:
435 _make_class_unpicklable(enum_class)
436 else:
437 enum_class.__module__ = module
Ethan Furmanca1b7942014-02-08 11:36:27 -0800438 if qualname is not None:
439 enum_class.__qualname__ = qualname
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700440
441 return enum_class
442
orlnub1230fb9fad2018-09-12 20:28:53 +0300443 def _convert_(cls, name, module, filter, source=None):
444 """
445 Create a new Enum subclass that replaces a collection of global constants
446 """
447 # convert all constants from source (or module) that pass filter() to
448 # a new Enum called name, and export the enum and its members back to
449 # module;
450 # also, replace the __reduce_ex__ method so unpickling works in
451 # previous Python versions
452 module_globals = vars(sys.modules[module])
453 if source:
454 source = vars(source)
455 else:
456 source = module_globals
457 # _value2member_map_ is populated in the same order every time
458 # for a consistent reverse mapping of number to name when there
459 # are multiple names for the same number.
460 members = [
461 (name, value)
462 for name, value in source.items()
463 if filter(name)]
464 try:
465 # sort by value
466 members.sort(key=lambda t: (t[1], t[0]))
467 except TypeError:
468 # unless some values aren't comparable, in which case sort by name
469 members.sort(key=lambda t: t[0])
470 cls = cls(name, members, module=module)
471 cls.__reduce_ex__ = _reduce_ex_by_name
472 module_globals.update(cls.__members__)
473 module_globals[name] = cls
474 return cls
475
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700476 @staticmethod
477 def _get_mixins_(bases):
478 """Returns the type for creating enum members, and the first inherited
479 enum class.
480
481 bases: the tuple of bases that was given to __new__
482
483 """
484 if not bases:
485 return object, Enum
486
Ethan Furman5bdab642018-09-21 19:03:09 -0700487 def _find_data_type(bases):
Miss Islington (bot)95b81e22020-09-15 16:59:48 -0700488 data_types = []
Ethan Furman5bdab642018-09-21 19:03:09 -0700489 for chain in bases:
Miss Islington (bot)95b81e22020-09-15 16:59:48 -0700490 candidate = None
Ethan Furman5bdab642018-09-21 19:03:09 -0700491 for base in chain.__mro__:
492 if base is object:
493 continue
494 elif '__new__' in base.__dict__:
Ethan Furmancd453852018-10-05 23:29:36 -0700495 if issubclass(base, Enum):
Ethan Furman5bdab642018-09-21 19:03:09 -0700496 continue
Miss Islington (bot)95b81e22020-09-15 16:59:48 -0700497 data_types.append(candidate or base)
498 break
499 elif not issubclass(base, Enum):
500 candidate = base
501 if len(data_types) > 1:
502 raise TypeError('too many data types: %r' % data_types)
503 elif data_types:
504 return data_types[0]
505 else:
506 return None
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700507
Ethan Furman5bdab642018-09-21 19:03:09 -0700508 # ensure final parent class is an Enum derivative, find any concrete
509 # data type, and check that Enum has no members
510 first_enum = bases[-1]
511 if not issubclass(first_enum, Enum):
512 raise TypeError("new enumerations should be created as "
513 "`EnumName([mixin_type, ...] [data_type,] enum_type)`")
514 member_type = _find_data_type(bases) or object
515 if first_enum._member_names_:
516 raise TypeError("Cannot extend enumerations")
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700517 return member_type, first_enum
518
519 @staticmethod
520 def _find_new_(classdict, member_type, first_enum):
521 """Returns the __new__ to be used for creating the enum members.
522
523 classdict: the class dictionary given to __new__
524 member_type: the data type whose __new__ will be used by default
525 first_enum: enumeration to check for an overriding __new__
526
527 """
528 # now find the correct __new__, checking to see of one was defined
529 # by the user; also check earlier enum classes in case a __new__ was
530 # saved as __new_member__
531 __new__ = classdict.get('__new__', None)
532
533 # should __new__ be saved as __new_member__ later?
534 save_new = __new__ is not None
535
536 if __new__ is None:
537 # check all possibles for __new_member__ before falling back to
538 # __new__
539 for method in ('__new_member__', '__new__'):
540 for possible in (member_type, first_enum):
541 target = getattr(possible, method, None)
542 if target not in {
543 None,
544 None.__new__,
545 object.__new__,
546 Enum.__new__,
547 }:
548 __new__ = target
549 break
550 if __new__ is not None:
551 break
552 else:
553 __new__ = object.__new__
554
555 # if a non-object.__new__ is used then whatever value/tuple was
556 # assigned to the enum member name will be passed to __new__ and to the
557 # new enum member's __init__
558 if __new__ is object.__new__:
559 use_args = False
560 else:
561 use_args = True
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700562 return __new__, save_new, use_args
563
564
565class Enum(metaclass=EnumMeta):
566 """Generic enumeration.
567
568 Derive from this class to define new enumerations.
569
570 """
571 def __new__(cls, value):
572 # all enum instances are actually created during class construction
573 # without calling this method; this method is called by the metaclass'
574 # __call__ (i.e. Color(3) ), and by pickle
575 if type(value) is cls:
Ethan Furman23bb6f42016-11-21 09:22:05 -0800576 # For lookups like Color(Color.RED)
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700577 return value
578 # by-value search for a matching enum member
579 # see if it's in the reverse mapping (for hashable values)
Ethan Furman2aa27322013-07-19 19:35:56 -0700580 try:
Andrew Svetlov34ae04f2018-12-26 20:45:33 +0200581 return cls._value2member_map_[value]
582 except KeyError:
583 # Not found, no need to do long O(n) search
584 pass
Ethan Furman2aa27322013-07-19 19:35:56 -0700585 except TypeError:
586 # not there, now do long search -- O(n) behavior
Ethan Furman520ad572013-07-19 19:47:21 -0700587 for member in cls._member_map_.values():
Ethan Furman0081f232014-09-16 17:31:23 -0700588 if member._value_ == value:
Ethan Furman2aa27322013-07-19 19:35:56 -0700589 return member
Ethan Furmanee47e5c2016-08-31 00:12:15 -0700590 # still not found -- try _missing_ hook
Ethan Furman019f0a02018-09-12 11:43:34 -0700591 try:
592 exc = None
593 result = cls._missing_(value)
594 except Exception as e:
595 exc = e
596 result = None
597 if isinstance(result, cls):
598 return result
599 else:
Walter Dörwald323842c2019-07-18 20:37:13 +0200600 ve_exc = ValueError("%r is not a valid %s" % (value, cls.__qualname__))
Ethan Furman019f0a02018-09-12 11:43:34 -0700601 if result is None and exc is None:
602 raise ve_exc
603 elif exc is None:
604 exc = TypeError(
605 'error in %s._missing_: returned %r instead of None or a valid member'
606 % (cls.__name__, result)
607 )
608 exc.__context__ = ve_exc
609 raise exc
Ethan Furmanee47e5c2016-08-31 00:12:15 -0700610
Ethan Furmanc16595e2016-09-10 23:36:59 -0700611 def _generate_next_value_(name, start, count, last_values):
612 for last_value in reversed(last_values):
613 try:
614 return last_value + 1
615 except TypeError:
616 pass
617 else:
Ethan Furmanee47e5c2016-08-31 00:12:15 -0700618 return start
Ethan Furmanc16595e2016-09-10 23:36:59 -0700619
Ethan Furmanee47e5c2016-08-31 00:12:15 -0700620 @classmethod
621 def _missing_(cls, value):
Walter Dörwald323842c2019-07-18 20:37:13 +0200622 raise ValueError("%r is not a valid %s" % (value, cls.__qualname__))
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700623
624 def __repr__(self):
625 return "<%s.%s: %r>" % (
Ethan Furman520ad572013-07-19 19:47:21 -0700626 self.__class__.__name__, self._name_, self._value_)
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700627
628 def __str__(self):
Ethan Furman520ad572013-07-19 19:47:21 -0700629 return "%s.%s" % (self.__class__.__name__, self._name_)
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700630
Ethan Furman388a3922013-08-12 06:51:41 -0700631 def __dir__(self):
Ethan Furman0ae550b2014-10-14 08:58:32 -0700632 added_behavior = [
633 m
634 for cls in self.__class__.mro()
635 for m in cls.__dict__
Ethan Furman354ecf12015-03-11 08:43:12 -0700636 if m[0] != '_' and m not in self._member_map_
Ethan Furman0ae550b2014-10-14 08:58:32 -0700637 ]
Ethan Furmanec5f8eb2014-10-21 13:40:35 -0700638 return (['__class__', '__doc__', '__module__'] + added_behavior)
Ethan Furman388a3922013-08-12 06:51:41 -0700639
Ethan Furmanec15a822013-08-31 19:17:41 -0700640 def __format__(self, format_spec):
641 # mixed-in Enums should use the mixed-in type's __format__, otherwise
642 # we can get strange results with the Enum name showing up instead of
643 # the value
644
thatneat2f19e822019-07-04 11:28:37 -0700645 # pure Enum branch, or branch with __str__ explicitly overridden
646 str_overridden = type(self).__str__ != Enum.__str__
647 if self._member_type_ is object or str_overridden:
Ethan Furmanec15a822013-08-31 19:17:41 -0700648 cls = str
649 val = str(self)
650 # mix-in branch
651 else:
652 cls = self._member_type_
Ethan Furman0081f232014-09-16 17:31:23 -0700653 val = self._value_
Ethan Furmanec15a822013-08-31 19:17:41 -0700654 return cls.__format__(val, format_spec)
655
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700656 def __hash__(self):
Ethan Furman520ad572013-07-19 19:47:21 -0700657 return hash(self._name_)
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700658
Ethan Furmanca1b7942014-02-08 11:36:27 -0800659 def __reduce_ex__(self, proto):
Ethan Furmandc870522014-02-18 12:37:12 -0800660 return self.__class__, (self._value_, )
Ethan Furmanca1b7942014-02-08 11:36:27 -0800661
Ethan Furman33918c12013-09-27 23:02:02 -0700662 # DynamicClassAttribute is used to provide access to the `name` and
663 # `value` properties of enum members while keeping some measure of
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700664 # protection from modification, while still allowing for an enumeration
665 # to have members named `name` and `value`. This works because enumeration
666 # members are not set directly on the enum class -- __getattr__ is
667 # used to look them up.
668
Ethan Furmane03ea372013-09-25 07:14:41 -0700669 @DynamicClassAttribute
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700670 def name(self):
Ethan Furmanc850f342013-09-15 16:59:35 -0700671 """The name of the Enum member."""
Ethan Furman520ad572013-07-19 19:47:21 -0700672 return self._name_
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700673
Ethan Furmane03ea372013-09-25 07:14:41 -0700674 @DynamicClassAttribute
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700675 def value(self):
Ethan Furmanc850f342013-09-15 16:59:35 -0700676 """The value of the Enum member."""
Ethan Furman520ad572013-07-19 19:47:21 -0700677 return self._value_
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700678
679
680class IntEnum(int, Enum):
681 """Enum where members are also (and must be) ints"""
Ethan Furmanf24bb352013-07-18 17:05:39 -0700682
683
Ethan Furman24e837f2015-03-18 17:27:57 -0700684def _reduce_ex_by_name(self, proto):
685 return self.name
686
Ethan Furman65a5a472016-09-01 23:55:19 -0700687class Flag(Enum):
Ethan Furmanee47e5c2016-08-31 00:12:15 -0700688 """Support for flags"""
Ethan Furmanc16595e2016-09-10 23:36:59 -0700689
690 def _generate_next_value_(name, start, count, last_values):
Ethan Furmanee47e5c2016-08-31 00:12:15 -0700691 """
692 Generate the next value when not given.
693
694 name: the name of the member
HongWeipengbb16fb22019-09-21 13:22:54 +0800695 start: the initial start value or None
Ethan Furmanee47e5c2016-08-31 00:12:15 -0700696 count: the number of existing members
697 last_value: the last value assigned or None
698 """
699 if not count:
700 return start if start is not None else 1
Ethan Furmanc16595e2016-09-10 23:36:59 -0700701 for last_value in reversed(last_values):
702 try:
703 high_bit = _high_bit(last_value)
704 break
Ethan Furman3515dcc2016-09-18 13:15:41 -0700705 except Exception:
Ethan Furmanc16595e2016-09-10 23:36:59 -0700706 raise TypeError('Invalid Flag value: %r' % last_value) from None
Ethan Furmanee47e5c2016-08-31 00:12:15 -0700707 return 2 ** (high_bit+1)
708
709 @classmethod
710 def _missing_(cls, value):
711 original_value = value
712 if value < 0:
713 value = ~value
714 possible_member = cls._create_pseudo_member_(value)
Ethan Furmanee47e5c2016-08-31 00:12:15 -0700715 if original_value < 0:
716 possible_member = ~possible_member
717 return possible_member
718
719 @classmethod
720 def _create_pseudo_member_(cls, value):
Ethan Furman3515dcc2016-09-18 13:15:41 -0700721 """
722 Create a composite member iff value contains only members.
723 """
Ethan Furmanee47e5c2016-08-31 00:12:15 -0700724 pseudo_member = cls._value2member_map_.get(value, None)
725 if pseudo_member is None:
Ethan Furman3515dcc2016-09-18 13:15:41 -0700726 # verify all bits are accounted for
727 _, extra_flags = _decompose(cls, value)
728 if extra_flags:
Walter Dörwald323842c2019-07-18 20:37:13 +0200729 raise ValueError("%r is not a valid %s" % (value, cls.__qualname__))
Ethan Furman3515dcc2016-09-18 13:15:41 -0700730 # construct a singleton enum pseudo-member
Ethan Furmanee47e5c2016-08-31 00:12:15 -0700731 pseudo_member = object.__new__(cls)
732 pseudo_member._name_ = None
733 pseudo_member._value_ = value
Ethan Furman28cf6632017-01-24 12:12:06 -0800734 # use setdefault in case another thread already created a composite
735 # with this value
736 pseudo_member = cls._value2member_map_.setdefault(value, pseudo_member)
Ethan Furmanee47e5c2016-08-31 00:12:15 -0700737 return pseudo_member
738
Ethan Furmanee47e5c2016-08-31 00:12:15 -0700739 def __contains__(self, other):
740 if not isinstance(other, self.__class__):
Rahul Jha94306522018-09-10 23:51:04 +0530741 raise TypeError(
742 "unsupported operand type(s) for 'in': '%s' and '%s'" % (
743 type(other).__qualname__, self.__class__.__qualname__))
Ethan Furmanee47e5c2016-08-31 00:12:15 -0700744 return other._value_ & self._value_ == other._value_
745
Ethan Furmanee47e5c2016-08-31 00:12:15 -0700746 def __repr__(self):
747 cls = self.__class__
748 if self._name_ is not None:
749 return '<%s.%s: %r>' % (cls.__name__, self._name_, self._value_)
Ethan Furman3515dcc2016-09-18 13:15:41 -0700750 members, uncovered = _decompose(cls, self._value_)
Ethan Furman27682d22016-09-04 11:39:01 -0700751 return '<%s.%s: %r>' % (
752 cls.__name__,
753 '|'.join([str(m._name_ or m._value_) for m in members]),
754 self._value_,
755 )
Ethan Furmanee47e5c2016-08-31 00:12:15 -0700756
757 def __str__(self):
758 cls = self.__class__
759 if self._name_ is not None:
760 return '%s.%s' % (cls.__name__, self._name_)
Ethan Furman3515dcc2016-09-18 13:15:41 -0700761 members, uncovered = _decompose(cls, self._value_)
Ethan Furmanee47e5c2016-08-31 00:12:15 -0700762 if len(members) == 1 and members[0]._name_ is None:
763 return '%s.%r' % (cls.__name__, members[0]._value_)
764 else:
765 return '%s.%s' % (
766 cls.__name__,
767 '|'.join([str(m._name_ or m._value_) for m in members]),
768 )
769
Ethan Furman25d94bb2016-09-02 16:32:32 -0700770 def __bool__(self):
771 return bool(self._value_)
772
Ethan Furmanee47e5c2016-08-31 00:12:15 -0700773 def __or__(self, other):
774 if not isinstance(other, self.__class__):
775 return NotImplemented
776 return self.__class__(self._value_ | other._value_)
777
778 def __and__(self, other):
779 if not isinstance(other, self.__class__):
780 return NotImplemented
781 return self.__class__(self._value_ & other._value_)
782
783 def __xor__(self, other):
784 if not isinstance(other, self.__class__):
785 return NotImplemented
786 return self.__class__(self._value_ ^ other._value_)
787
788 def __invert__(self):
Ethan Furman3515dcc2016-09-18 13:15:41 -0700789 members, uncovered = _decompose(self.__class__, self._value_)
Serhiy Storchaka81108372017-09-26 00:55:55 +0300790 inverted = self.__class__(0)
791 for m in self.__class__:
792 if m not in members and not (m._value_ & self._value_):
793 inverted = inverted | m
Ethan Furmanee47e5c2016-08-31 00:12:15 -0700794 return self.__class__(inverted)
795
796
Ethan Furman65a5a472016-09-01 23:55:19 -0700797class IntFlag(int, Flag):
Ethan Furmanee47e5c2016-08-31 00:12:15 -0700798 """Support for integer-based Flags"""
799
800 @classmethod
Ethan Furman3515dcc2016-09-18 13:15:41 -0700801 def _missing_(cls, value):
802 if not isinstance(value, int):
Walter Dörwald323842c2019-07-18 20:37:13 +0200803 raise ValueError("%r is not a valid %s" % (value, cls.__qualname__))
Ethan Furman3515dcc2016-09-18 13:15:41 -0700804 new_member = cls._create_pseudo_member_(value)
805 return new_member
806
807 @classmethod
Ethan Furmanee47e5c2016-08-31 00:12:15 -0700808 def _create_pseudo_member_(cls, value):
809 pseudo_member = cls._value2member_map_.get(value, None)
810 if pseudo_member is None:
Ethan Furman3515dcc2016-09-18 13:15:41 -0700811 need_to_create = [value]
812 # get unaccounted for bits
813 _, extra_flags = _decompose(cls, value)
814 # timer = 10
815 while extra_flags:
816 # timer -= 1
817 bit = _high_bit(extra_flags)
818 flag_value = 2 ** bit
819 if (flag_value not in cls._value2member_map_ and
820 flag_value not in need_to_create
821 ):
822 need_to_create.append(flag_value)
823 if extra_flags == -flag_value:
824 extra_flags = 0
825 else:
826 extra_flags ^= flag_value
827 for value in reversed(need_to_create):
828 # construct singleton pseudo-members
829 pseudo_member = int.__new__(cls, value)
830 pseudo_member._name_ = None
831 pseudo_member._value_ = value
Ethan Furman28cf6632017-01-24 12:12:06 -0800832 # use setdefault in case another thread already created a composite
833 # with this value
834 pseudo_member = cls._value2member_map_.setdefault(value, pseudo_member)
Ethan Furmanee47e5c2016-08-31 00:12:15 -0700835 return pseudo_member
836
Ethan Furmanee47e5c2016-08-31 00:12:15 -0700837 def __or__(self, other):
838 if not isinstance(other, (self.__class__, int)):
839 return NotImplemented
Ethan Furman3515dcc2016-09-18 13:15:41 -0700840 result = self.__class__(self._value_ | self.__class__(other)._value_)
841 return result
Ethan Furmanee47e5c2016-08-31 00:12:15 -0700842
843 def __and__(self, other):
844 if not isinstance(other, (self.__class__, int)):
845 return NotImplemented
846 return self.__class__(self._value_ & self.__class__(other)._value_)
847
848 def __xor__(self, other):
849 if not isinstance(other, (self.__class__, int)):
850 return NotImplemented
851 return self.__class__(self._value_ ^ self.__class__(other)._value_)
852
853 __ror__ = __or__
854 __rand__ = __and__
855 __rxor__ = __xor__
856
857 def __invert__(self):
Ethan Furman3515dcc2016-09-18 13:15:41 -0700858 result = self.__class__(~self._value_)
859 return result
Ethan Furmanee47e5c2016-08-31 00:12:15 -0700860
861
862def _high_bit(value):
Ethan Furman04439532016-09-02 15:50:21 -0700863 """returns index of highest bit, or -1 if value is zero or negative"""
Ethan Furman3515dcc2016-09-18 13:15:41 -0700864 return value.bit_length() - 1
Ethan Furmanee47e5c2016-08-31 00:12:15 -0700865
Ethan Furmanf24bb352013-07-18 17:05:39 -0700866def unique(enumeration):
867 """Class decorator for enumerations ensuring unique member values."""
868 duplicates = []
869 for name, member in enumeration.__members__.items():
870 if name != member.name:
871 duplicates.append((name, member.name))
872 if duplicates:
873 alias_details = ', '.join(
874 ["%s -> %s" % (alias, name) for (alias, name) in duplicates])
875 raise ValueError('duplicate values found in %r: %s' %
876 (enumeration, alias_details))
877 return enumeration
Ethan Furman3515dcc2016-09-18 13:15:41 -0700878
879def _decompose(flag, value):
880 """Extract all members from the value."""
881 # _decompose is only called if the value is not named
882 not_covered = value
883 negative = value < 0
Ethan Furman3515dcc2016-09-18 13:15:41 -0700884 members = []
HongWeipeng0b41a922019-11-27 06:36:02 +0800885 for member in flag:
886 member_value = member.value
Ethan Furman3515dcc2016-09-18 13:15:41 -0700887 if member_value and member_value & value == member_value:
888 members.append(member)
889 not_covered &= ~member_value
HongWeipeng0b41a922019-11-27 06:36:02 +0800890 if not negative:
891 tmp = not_covered
892 while tmp:
893 flag_value = 2 ** _high_bit(tmp)
894 if flag_value in flag._value2member_map_:
895 members.append(flag._value2member_map_[flag_value])
896 not_covered &= ~flag_value
897 tmp &= ~flag_value
Ethan Furman3515dcc2016-09-18 13:15:41 -0700898 if not members and value in flag._value2member_map_:
899 members.append(flag._value2member_map_[value])
900 members.sort(key=lambda m: m._value_, reverse=True)
901 if len(members) > 1 and members[0].value == value:
902 # we have the breakdown, don't need the value member itself
903 members.pop(0)
904 return members, not_covered