blob: e72d306267461842a1250189655983677b4b3264 [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 ):
Zackery Spytz2ec67522020-09-13 14:27:51 -060079 raise ValueError(f'_sunder_ names, such as "{key}", are '
80 'reserved for future Enum use')
Ethan Furmanc16595e2016-09-10 23:36:59 -070081 if key == '_generate_next_value_':
Ethan Onstottd9a43e22020-04-28 13:20:55 -040082 # check if members already defined as auto()
83 if self._auto_called:
84 raise TypeError("_generate_next_value_ must be defined before members")
Ethan Furmanc16595e2016-09-10 23:36:59 -070085 setattr(self, '_generate_next_value', value)
Ethan Furmana4b1bb42018-01-22 07:56:37 -080086 elif key == '_ignore_':
87 if isinstance(value, str):
88 value = value.replace(',',' ').split()
89 else:
90 value = list(value)
91 self._ignore = value
92 already = set(value) & set(self._member_names)
93 if already:
94 raise ValueError('_ignore_ cannot specify already set names: %r' % (already, ))
Ethan Furman101e0742013-09-15 12:34:36 -070095 elif _is_dunder(key):
Ethan Furmane8e61272016-08-20 07:19:31 -070096 if key == '__order__':
97 key = '_order_'
Ethan Furman101e0742013-09-15 12:34:36 -070098 elif key in self._member_names:
99 # descriptor overwriting an enum?
100 raise TypeError('Attempted to reuse key: %r' % key)
Ethan Furmana4b1bb42018-01-22 07:56:37 -0800101 elif key in self._ignore:
102 pass
Ethan Furman101e0742013-09-15 12:34:36 -0700103 elif not _is_descriptor(value):
104 if key in self:
105 # enum overwriting a descriptor?
Ethan Furmanee47e5c2016-08-31 00:12:15 -0700106 raise TypeError('%r already defined as: %r' % (key, self[key]))
Ethan Furmanc16595e2016-09-10 23:36:59 -0700107 if isinstance(value, auto):
Ethan Onstottd9a43e22020-04-28 13:20:55 -0400108 self._auto_called = True
Ethan Furman3515dcc2016-09-18 13:15:41 -0700109 if value.value == _auto_null:
110 value.value = self._generate_next_value(key, 1, len(self._member_names), self._last_values[:])
111 value = value.value
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700112 self._member_names.append(key)
Ethan Furmanc16595e2016-09-10 23:36:59 -0700113 self._last_values.append(value)
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700114 super().__setitem__(key, value)
115
116
Ezio Melotti9a3777e2013-08-17 15:53:55 +0300117# Dummy value for Enum as EnumMeta explicitly checks for it, but of course
118# until EnumMeta finishes running the first time the Enum class doesn't exist.
119# This is also why there are checks in EnumMeta like `if Enum is not None`
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700120Enum = None
121
Ethan Furman332dbc72016-08-20 00:00:52 -0700122
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700123class EnumMeta(type):
124 """Metaclass for Enum"""
125 @classmethod
Ethan Furman332dbc72016-08-20 00:00:52 -0700126 def __prepare__(metacls, cls, bases):
Ethan Furmanee47e5c2016-08-31 00:12:15 -0700127 # create the namespace dict
128 enum_dict = _EnumDict()
129 # inherit previous flags and _generate_next_value_ function
130 member_type, first_enum = metacls._get_mixins_(bases)
131 if first_enum is not None:
132 enum_dict['_generate_next_value_'] = getattr(first_enum, '_generate_next_value_', None)
133 return enum_dict
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700134
Ethan Furman65a5a472016-09-01 23:55:19 -0700135 def __new__(metacls, cls, bases, classdict):
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700136 # an Enum class is final once enumeration items have been defined; it
137 # cannot be mixed with other types (int, float, etc.) if it has an
138 # inherited __new__ unless a new __new__ is defined (or the resulting
139 # class will fail).
Ethan Furmana4b1bb42018-01-22 07:56:37 -0800140 #
141 # remove any keys listed in _ignore_
142 classdict.setdefault('_ignore_', []).append('_ignore_')
143 ignore = classdict['_ignore_']
144 for key in ignore:
145 classdict.pop(key, None)
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700146 member_type, first_enum = metacls._get_mixins_(bases)
147 __new__, save_new, use_args = metacls._find_new_(classdict, member_type,
148 first_enum)
149
150 # save enum items into separate mapping so they don't get baked into
151 # the new class
Ethan Furmanee47e5c2016-08-31 00:12:15 -0700152 enum_members = {k: classdict[k] for k in classdict._member_names}
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700153 for name in classdict._member_names:
154 del classdict[name]
155
Ethan Furmane8e61272016-08-20 07:19:31 -0700156 # adjust the sunders
157 _order_ = classdict.pop('_order_', None)
158
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700159 # check for illegal enum names (any others?)
Brennan D Baraban8b914d22019-03-03 14:09:11 -0800160 invalid_names = set(enum_members) & {'mro', ''}
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700161 if invalid_names:
162 raise ValueError('Invalid enum member name: {0}'.format(
163 ','.join(invalid_names)))
164
Ethan Furman48a724f2015-04-11 23:23:06 -0700165 # create a default docstring if one has not been provided
166 if '__doc__' not in classdict:
167 classdict['__doc__'] = 'An enumeration.'
168
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700169 # create our new Enum type
170 enum_class = super().__new__(metacls, cls, bases, classdict)
Ethan Furman520ad572013-07-19 19:47:21 -0700171 enum_class._member_names_ = [] # names in definition order
INADA Naokie57f91a2018-06-19 01:14:26 +0900172 enum_class._member_map_ = {} # name->value map
Ethan Furman5e5a8232013-08-04 08:42:23 -0700173 enum_class._member_type_ = member_type
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700174
orlnub1230fb9fad2018-09-12 20:28:53 +0300175 # save DynamicClassAttribute attributes from super classes so we know
176 # if we can take the shortcut of storing members in the class dict
177 dynamic_attributes = {k for c in enum_class.mro()
178 for k, v in c.__dict__.items()
179 if isinstance(v, DynamicClassAttribute)}
Ethan Furman354ecf12015-03-11 08:43:12 -0700180
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700181 # Reverse value->name map for hashable values.
Ethan Furman520ad572013-07-19 19:47:21 -0700182 enum_class._value2member_map_ = {}
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700183
Ethan Furman2da95042014-03-03 12:42:52 -0800184 # If a custom type is mixed into the Enum, and it does not know how
185 # to pickle itself, pickle.dumps will succeed but pickle.loads will
186 # fail. Rather than have the error show up later and possibly far
187 # from the source, sabotage the pickle protocol for this class so
188 # that pickle.dumps also fails.
189 #
190 # However, if the new class implements its own __reduce_ex__, do not
191 # sabotage -- it's on them to make sure it works correctly. We use
192 # __reduce_ex__ instead of any of the others as it is preferred by
193 # pickle over __reduce__, and it handles all pickle protocols.
194 if '__reduce_ex__' not in classdict:
Ethan Furmandc870522014-02-18 12:37:12 -0800195 if member_type is not object:
196 methods = ('__getnewargs_ex__', '__getnewargs__',
197 '__reduce_ex__', '__reduce__')
Ethan Furman2da95042014-03-03 12:42:52 -0800198 if not any(m in member_type.__dict__ for m in methods):
Ethan Furmandc870522014-02-18 12:37:12 -0800199 _make_class_unpicklable(enum_class)
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700200
201 # instantiate them, checking for duplicates as we go
202 # we instantiate first instead of checking for duplicates first in case
203 # a custom __new__ is doing something funky with the values -- such as
204 # auto-numbering ;)
205 for member_name in classdict._member_names:
Ethan Furmanee47e5c2016-08-31 00:12:15 -0700206 value = enum_members[member_name]
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700207 if not isinstance(value, tuple):
208 args = (value, )
209 else:
210 args = value
211 if member_type is tuple: # special case for tuple enums
212 args = (args, ) # wrap it one more time
213 if not use_args:
214 enum_member = __new__(enum_class)
Ethan Furmanb41803e2013-07-25 13:50:45 -0700215 if not hasattr(enum_member, '_value_'):
216 enum_member._value_ = value
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700217 else:
218 enum_member = __new__(enum_class, *args)
Ethan Furmanb41803e2013-07-25 13:50:45 -0700219 if not hasattr(enum_member, '_value_'):
Ethan Furmanee47e5c2016-08-31 00:12:15 -0700220 if member_type is object:
221 enum_member._value_ = value
222 else:
223 enum_member._value_ = member_type(*args)
Ethan Furman520ad572013-07-19 19:47:21 -0700224 value = enum_member._value_
Ethan Furman520ad572013-07-19 19:47:21 -0700225 enum_member._name_ = member_name
Ethan Furmanc850f342013-09-15 16:59:35 -0700226 enum_member.__objclass__ = enum_class
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700227 enum_member.__init__(*args)
228 # If another member with the same value was already defined, the
229 # new member becomes an alias to the existing one.
Ethan Furman520ad572013-07-19 19:47:21 -0700230 for name, canonical_member in enum_class._member_map_.items():
Ethan Furman0081f232014-09-16 17:31:23 -0700231 if canonical_member._value_ == enum_member._value_:
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700232 enum_member = canonical_member
233 break
234 else:
235 # Aliases don't appear in member names (only in __members__).
Ethan Furman520ad572013-07-19 19:47:21 -0700236 enum_class._member_names_.append(member_name)
Ethan Furman354ecf12015-03-11 08:43:12 -0700237 # performance boost for any member that would not shadow
238 # a DynamicClassAttribute
orlnub1230fb9fad2018-09-12 20:28:53 +0300239 if member_name not in dynamic_attributes:
Ethan Furman354ecf12015-03-11 08:43:12 -0700240 setattr(enum_class, member_name, enum_member)
241 # now add to _member_map_
Ethan Furman520ad572013-07-19 19:47:21 -0700242 enum_class._member_map_[member_name] = enum_member
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700243 try:
244 # This may fail if value is not hashable. We can't add the value
245 # to the map, and by-value lookups for this value will be
246 # linear.
Ethan Furman520ad572013-07-19 19:47:21 -0700247 enum_class._value2member_map_[value] = enum_member
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700248 except TypeError:
249 pass
250
251 # double check that repr and friends are not the mixin's or various
252 # things break (such as pickle)
Ethan Furman22415ad2020-09-15 16:28:25 -0700253 # however, if the method is defined in the Enum itself, don't replace
254 # it
Ethan Furmandc870522014-02-18 12:37:12 -0800255 for name in ('__repr__', '__str__', '__format__', '__reduce_ex__'):
Ethan Furman22415ad2020-09-15 16:28:25 -0700256 if name in classdict:
257 continue
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700258 class_method = getattr(enum_class, name)
259 obj_method = getattr(member_type, name, None)
260 enum_method = getattr(first_enum, name, None)
261 if obj_method is not None and obj_method is class_method:
262 setattr(enum_class, name, enum_method)
263
264 # replace any other __new__ with our own (as long as Enum is not None,
265 # anyway) -- again, this is to support pickle
266 if Enum is not None:
267 # if the user defined their own __new__, save it before it gets
268 # clobbered in case they subclass later
269 if save_new:
270 enum_class.__new_member__ = __new__
271 enum_class.__new__ = Enum.__new__
Ethan Furmane8e61272016-08-20 07:19:31 -0700272
273 # py3 support for definition order (helps keep py2/py3 code in sync)
274 if _order_ is not None:
275 if isinstance(_order_, str):
276 _order_ = _order_.replace(',', ' ').split()
277 if _order_ != enum_class._member_names_:
278 raise TypeError('member order does not match _order_')
279
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700280 return enum_class
281
Ethan Furman5de67b12016-04-13 23:52:09 -0700282 def __bool__(self):
283 """
284 classes/types should always be True.
285 """
286 return True
287
Ethan Furmand9925a12014-09-16 20:35:55 -0700288 def __call__(cls, value, names=None, *, module=None, qualname=None, type=None, start=1):
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700289 """Either returns an existing member, or creates a new enum class.
290
291 This method is used both when an enum class is given a value to match
292 to an enumeration member (i.e. Color(3)) and for the functional API
Ethan Furman23bb6f42016-11-21 09:22:05 -0800293 (i.e. Color = Enum('Color', names='RED GREEN BLUE')).
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700294
Ethan Furman2da95042014-03-03 12:42:52 -0800295 When used for the functional API:
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700296
Ethan Furman2da95042014-03-03 12:42:52 -0800297 `value` will be the name of the new class.
298
299 `names` should be either a string of white-space/comma delimited names
Ethan Furmand9925a12014-09-16 20:35:55 -0700300 (values will start at `start`), or an iterator/mapping of name, value pairs.
Ethan Furman2da95042014-03-03 12:42:52 -0800301
302 `module` should be set to the module this class is being created in;
303 if it is not set, an attempt to find that module will be made, but if
304 it fails the class will not be picklable.
305
306 `qualname` should be set to the actual location this class can be found
307 at in its module; by default it is set to the global scope. If this is
308 not correct, unpickling will fail in some circumstances.
309
310 `type`, if set, will be mixed in as the first base class.
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700311
312 """
313 if names is None: # simple value lookup
314 return cls.__new__(cls, value)
315 # otherwise, functional API: we're creating a new Enum type
Ethan Furmand9925a12014-09-16 20:35:55 -0700316 return cls._create_(value, names, module=module, qualname=qualname, type=type, start=start)
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700317
318 def __contains__(cls, member):
Rahul Jha94306522018-09-10 23:51:04 +0530319 if not isinstance(member, Enum):
320 raise TypeError(
321 "unsupported operand type(s) for 'in': '%s' and '%s'" % (
322 type(member).__qualname__, cls.__class__.__qualname__))
Ethan Furman0081f232014-09-16 17:31:23 -0700323 return isinstance(member, cls) and member._name_ in cls._member_map_
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700324
Ethan Furman64a99722013-09-22 16:18:19 -0700325 def __delattr__(cls, attr):
326 # nicer error message when someone tries to delete an attribute
327 # (see issue19025).
328 if attr in cls._member_map_:
329 raise AttributeError(
330 "%s: cannot delete Enum member." % cls.__name__)
331 super().__delattr__(attr)
332
Ethan Furman388a3922013-08-12 06:51:41 -0700333 def __dir__(self):
Ethan Furman64a99722013-09-22 16:18:19 -0700334 return (['__class__', '__doc__', '__members__', '__module__'] +
335 self._member_names_)
Ethan Furman388a3922013-08-12 06:51:41 -0700336
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700337 def __getattr__(cls, name):
338 """Return the enum member matching `name`
339
340 We use __getattr__ instead of descriptors or inserting into the enum
341 class' __dict__ in order to support `name` and `value` being both
342 properties for enum members (which live in the class' __dict__) and
343 enum members themselves.
344
345 """
346 if _is_dunder(name):
347 raise AttributeError(name)
348 try:
Ethan Furman520ad572013-07-19 19:47:21 -0700349 return cls._member_map_[name]
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700350 except KeyError:
351 raise AttributeError(name) from None
352
353 def __getitem__(cls, name):
Ethan Furman520ad572013-07-19 19:47:21 -0700354 return cls._member_map_[name]
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700355
356 def __iter__(cls):
Ethan Furman520ad572013-07-19 19:47:21 -0700357 return (cls._member_map_[name] for name in cls._member_names_)
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700358
359 def __len__(cls):
Ethan Furman520ad572013-07-19 19:47:21 -0700360 return len(cls._member_names_)
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700361
Ethan Furman2131a4a2013-09-14 18:11:24 -0700362 @property
363 def __members__(cls):
364 """Returns a mapping of member name->value.
365
366 This mapping lists all enum members, including aliases. Note that this
367 is a read-only view of the internal mapping.
368
369 """
370 return MappingProxyType(cls._member_map_)
371
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700372 def __repr__(cls):
373 return "<enum %r>" % cls.__name__
374
Ethan Furman2131a4a2013-09-14 18:11:24 -0700375 def __reversed__(cls):
376 return (cls._member_map_[name] for name in reversed(cls._member_names_))
377
Ethan Furmanf203f2d2013-09-06 07:16:48 -0700378 def __setattr__(cls, name, value):
379 """Block attempts to reassign Enum members.
380
381 A simple assignment to the class namespace only changes one of the
382 several possible ways to get an Enum member from the Enum class,
383 resulting in an inconsistent Enumeration.
384
385 """
386 member_map = cls.__dict__.get('_member_map_', {})
387 if name in member_map:
388 raise AttributeError('Cannot reassign members.')
389 super().__setattr__(name, value)
390
anentropicb8e21f12018-04-16 04:40:35 +0100391 def _create_(cls, class_name, names, *, module=None, qualname=None, type=None, start=1):
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700392 """Convenience method to create a new Enum class.
393
394 `names` can be:
395
396 * A string containing member names, separated either with spaces or
Ethan Furmand9925a12014-09-16 20:35:55 -0700397 commas. Values are incremented by 1 from `start`.
398 * An iterable of member names. Values are incremented by 1 from `start`.
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700399 * An iterable of (member name, value) pairs.
Ethan Furmand9925a12014-09-16 20:35:55 -0700400 * A mapping of member name -> value pairs.
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700401
402 """
403 metacls = cls.__class__
404 bases = (cls, ) if type is None else (type, cls)
Ethan Furmanee47e5c2016-08-31 00:12:15 -0700405 _, first_enum = cls._get_mixins_(bases)
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700406 classdict = metacls.__prepare__(class_name, bases)
407
408 # special processing needed for names?
409 if isinstance(names, str):
410 names = names.replace(',', ' ').split()
Dong-hee Nadcc8ce42017-06-22 01:52:32 +0900411 if isinstance(names, (tuple, list)) and names and isinstance(names[0], str):
Ethan Furmanee47e5c2016-08-31 00:12:15 -0700412 original_names, names = names, []
Ethan Furmanc16595e2016-09-10 23:36:59 -0700413 last_values = []
Ethan Furmanee47e5c2016-08-31 00:12:15 -0700414 for count, name in enumerate(original_names):
Ethan Furmanc16595e2016-09-10 23:36:59 -0700415 value = first_enum._generate_next_value_(name, start, count, last_values[:])
416 last_values.append(value)
417 names.append((name, value))
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700418
419 # Here, names is either an iterable of (name, value) or a mapping.
420 for item in names:
421 if isinstance(item, str):
422 member_name, member_value = item, names[item]
423 else:
424 member_name, member_value = item
425 classdict[member_name] = member_value
426 enum_class = metacls.__new__(metacls, class_name, bases, classdict)
427
428 # TODO: replace the frame hack if a blessed way to know the calling
429 # module is ever developed
430 if module is None:
431 try:
432 module = sys._getframe(2).f_globals['__name__']
Pablo Galindo293dd232019-11-19 21:34:03 +0000433 except (AttributeError, ValueError, KeyError):
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700434 pass
435 if module is None:
436 _make_class_unpicklable(enum_class)
437 else:
438 enum_class.__module__ = module
Ethan Furmanca1b7942014-02-08 11:36:27 -0800439 if qualname is not None:
440 enum_class.__qualname__ = qualname
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700441
442 return enum_class
443
orlnub1230fb9fad2018-09-12 20:28:53 +0300444 def _convert_(cls, name, module, filter, source=None):
445 """
446 Create a new Enum subclass that replaces a collection of global constants
447 """
448 # convert all constants from source (or module) that pass filter() to
449 # a new Enum called name, and export the enum and its members back to
450 # module;
451 # also, replace the __reduce_ex__ method so unpickling works in
452 # previous Python versions
453 module_globals = vars(sys.modules[module])
454 if source:
455 source = vars(source)
456 else:
457 source = module_globals
458 # _value2member_map_ is populated in the same order every time
459 # for a consistent reverse mapping of number to name when there
460 # are multiple names for the same number.
461 members = [
462 (name, value)
463 for name, value in source.items()
464 if filter(name)]
465 try:
466 # sort by value
467 members.sort(key=lambda t: (t[1], t[0]))
468 except TypeError:
469 # unless some values aren't comparable, in which case sort by name
470 members.sort(key=lambda t: t[0])
471 cls = cls(name, members, module=module)
472 cls.__reduce_ex__ = _reduce_ex_by_name
473 module_globals.update(cls.__members__)
474 module_globals[name] = cls
475 return cls
476
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700477 @staticmethod
478 def _get_mixins_(bases):
479 """Returns the type for creating enum members, and the first inherited
480 enum class.
481
482 bases: the tuple of bases that was given to __new__
483
484 """
485 if not bases:
486 return object, Enum
487
Ethan Furman5bdab642018-09-21 19:03:09 -0700488 def _find_data_type(bases):
Ethan Furmanbff01f32020-09-15 15:56:26 -0700489 data_types = []
Ethan Furman5bdab642018-09-21 19:03:09 -0700490 for chain in bases:
Ethan Furmanbff01f32020-09-15 15:56:26 -0700491 candidate = None
Ethan Furman5bdab642018-09-21 19:03:09 -0700492 for base in chain.__mro__:
493 if base is object:
494 continue
495 elif '__new__' in base.__dict__:
Ethan Furmancd453852018-10-05 23:29:36 -0700496 if issubclass(base, Enum):
Ethan Furman5bdab642018-09-21 19:03:09 -0700497 continue
Ethan Furmanbff01f32020-09-15 15:56:26 -0700498 data_types.append(candidate or base)
499 break
500 elif not issubclass(base, Enum):
501 candidate = base
502 if len(data_types) > 1:
503 raise TypeError('too many data types: %r' % data_types)
504 elif data_types:
505 return data_types[0]
506 else:
507 return None
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700508
Ethan Furman5bdab642018-09-21 19:03:09 -0700509 # ensure final parent class is an Enum derivative, find any concrete
510 # data type, and check that Enum has no members
511 first_enum = bases[-1]
512 if not issubclass(first_enum, Enum):
513 raise TypeError("new enumerations should be created as "
514 "`EnumName([mixin_type, ...] [data_type,] enum_type)`")
515 member_type = _find_data_type(bases) or object
516 if first_enum._member_names_:
517 raise TypeError("Cannot extend enumerations")
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700518 return member_type, first_enum
519
520 @staticmethod
521 def _find_new_(classdict, member_type, first_enum):
522 """Returns the __new__ to be used for creating the enum members.
523
524 classdict: the class dictionary given to __new__
525 member_type: the data type whose __new__ will be used by default
526 first_enum: enumeration to check for an overriding __new__
527
528 """
529 # now find the correct __new__, checking to see of one was defined
530 # by the user; also check earlier enum classes in case a __new__ was
531 # saved as __new_member__
532 __new__ = classdict.get('__new__', None)
533
534 # should __new__ be saved as __new_member__ later?
535 save_new = __new__ is not None
536
537 if __new__ is None:
538 # check all possibles for __new_member__ before falling back to
539 # __new__
540 for method in ('__new_member__', '__new__'):
541 for possible in (member_type, first_enum):
542 target = getattr(possible, method, None)
543 if target not in {
544 None,
545 None.__new__,
546 object.__new__,
547 Enum.__new__,
548 }:
549 __new__ = target
550 break
551 if __new__ is not None:
552 break
553 else:
554 __new__ = object.__new__
555
556 # if a non-object.__new__ is used then whatever value/tuple was
557 # assigned to the enum member name will be passed to __new__ and to the
558 # new enum member's __init__
559 if __new__ is object.__new__:
560 use_args = False
561 else:
562 use_args = True
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700563 return __new__, save_new, use_args
564
565
566class Enum(metaclass=EnumMeta):
567 """Generic enumeration.
568
569 Derive from this class to define new enumerations.
570
571 """
572 def __new__(cls, value):
573 # all enum instances are actually created during class construction
574 # without calling this method; this method is called by the metaclass'
575 # __call__ (i.e. Color(3) ), and by pickle
576 if type(value) is cls:
Ethan Furman23bb6f42016-11-21 09:22:05 -0800577 # For lookups like Color(Color.RED)
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700578 return value
579 # by-value search for a matching enum member
580 # see if it's in the reverse mapping (for hashable values)
Ethan Furman2aa27322013-07-19 19:35:56 -0700581 try:
Andrew Svetlov34ae04f2018-12-26 20:45:33 +0200582 return cls._value2member_map_[value]
583 except KeyError:
584 # Not found, no need to do long O(n) search
585 pass
Ethan Furman2aa27322013-07-19 19:35:56 -0700586 except TypeError:
587 # not there, now do long search -- O(n) behavior
Ethan Furman520ad572013-07-19 19:47:21 -0700588 for member in cls._member_map_.values():
Ethan Furman0081f232014-09-16 17:31:23 -0700589 if member._value_ == value:
Ethan Furman2aa27322013-07-19 19:35:56 -0700590 return member
Ethan Furmanee47e5c2016-08-31 00:12:15 -0700591 # still not found -- try _missing_ hook
Ethan Furman019f0a02018-09-12 11:43:34 -0700592 try:
593 exc = None
594 result = cls._missing_(value)
595 except Exception as e:
596 exc = e
597 result = None
598 if isinstance(result, cls):
599 return result
600 else:
Walter Dörwald323842c2019-07-18 20:37:13 +0200601 ve_exc = ValueError("%r is not a valid %s" % (value, cls.__qualname__))
Ethan Furman019f0a02018-09-12 11:43:34 -0700602 if result is None and exc is None:
603 raise ve_exc
604 elif exc is None:
605 exc = TypeError(
606 'error in %s._missing_: returned %r instead of None or a valid member'
607 % (cls.__name__, result)
608 )
609 exc.__context__ = ve_exc
610 raise exc
Ethan Furmanee47e5c2016-08-31 00:12:15 -0700611
Ethan Furmanc16595e2016-09-10 23:36:59 -0700612 def _generate_next_value_(name, start, count, last_values):
613 for last_value in reversed(last_values):
614 try:
615 return last_value + 1
616 except TypeError:
617 pass
618 else:
Ethan Furmanee47e5c2016-08-31 00:12:15 -0700619 return start
Ethan Furmanc16595e2016-09-10 23:36:59 -0700620
Ethan Furmanee47e5c2016-08-31 00:12:15 -0700621 @classmethod
622 def _missing_(cls, value):
Walter Dörwald323842c2019-07-18 20:37:13 +0200623 raise ValueError("%r is not a valid %s" % (value, cls.__qualname__))
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700624
625 def __repr__(self):
626 return "<%s.%s: %r>" % (
Ethan Furman520ad572013-07-19 19:47:21 -0700627 self.__class__.__name__, self._name_, self._value_)
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700628
629 def __str__(self):
Ethan Furman520ad572013-07-19 19:47:21 -0700630 return "%s.%s" % (self.__class__.__name__, self._name_)
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700631
Ethan Furman388a3922013-08-12 06:51:41 -0700632 def __dir__(self):
Ethan Furman0ae550b2014-10-14 08:58:32 -0700633 added_behavior = [
634 m
635 for cls in self.__class__.mro()
636 for m in cls.__dict__
Ethan Furman354ecf12015-03-11 08:43:12 -0700637 if m[0] != '_' and m not in self._member_map_
Ethan Furman0ae550b2014-10-14 08:58:32 -0700638 ]
Ethan Furmanec5f8eb2014-10-21 13:40:35 -0700639 return (['__class__', '__doc__', '__module__'] + added_behavior)
Ethan Furman388a3922013-08-12 06:51:41 -0700640
Ethan Furmanec15a822013-08-31 19:17:41 -0700641 def __format__(self, format_spec):
642 # mixed-in Enums should use the mixed-in type's __format__, otherwise
643 # we can get strange results with the Enum name showing up instead of
644 # the value
645
thatneat2f19e822019-07-04 11:28:37 -0700646 # pure Enum branch, or branch with __str__ explicitly overridden
647 str_overridden = type(self).__str__ != Enum.__str__
648 if self._member_type_ is object or str_overridden:
Ethan Furmanec15a822013-08-31 19:17:41 -0700649 cls = str
650 val = str(self)
651 # mix-in branch
652 else:
653 cls = self._member_type_
Ethan Furman0081f232014-09-16 17:31:23 -0700654 val = self._value_
Ethan Furmanec15a822013-08-31 19:17:41 -0700655 return cls.__format__(val, format_spec)
656
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700657 def __hash__(self):
Ethan Furman520ad572013-07-19 19:47:21 -0700658 return hash(self._name_)
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700659
Ethan Furmanca1b7942014-02-08 11:36:27 -0800660 def __reduce_ex__(self, proto):
Ethan Furmandc870522014-02-18 12:37:12 -0800661 return self.__class__, (self._value_, )
Ethan Furmanca1b7942014-02-08 11:36:27 -0800662
Ethan Furman33918c12013-09-27 23:02:02 -0700663 # DynamicClassAttribute is used to provide access to the `name` and
664 # `value` properties of enum members while keeping some measure of
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700665 # protection from modification, while still allowing for an enumeration
666 # to have members named `name` and `value`. This works because enumeration
667 # members are not set directly on the enum class -- __getattr__ is
668 # used to look them up.
669
Ethan Furmane03ea372013-09-25 07:14:41 -0700670 @DynamicClassAttribute
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700671 def name(self):
Ethan Furmanc850f342013-09-15 16:59:35 -0700672 """The name of the Enum member."""
Ethan Furman520ad572013-07-19 19:47:21 -0700673 return self._name_
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700674
Ethan Furmane03ea372013-09-25 07:14:41 -0700675 @DynamicClassAttribute
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700676 def value(self):
Ethan Furmanc850f342013-09-15 16:59:35 -0700677 """The value of the Enum member."""
Ethan Furman520ad572013-07-19 19:47:21 -0700678 return self._value_
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700679
680
681class IntEnum(int, Enum):
682 """Enum where members are also (and must be) ints"""
Ethan Furmanf24bb352013-07-18 17:05:39 -0700683
684
Ethan Furman24e837f2015-03-18 17:27:57 -0700685def _reduce_ex_by_name(self, proto):
686 return self.name
687
Ethan Furman65a5a472016-09-01 23:55:19 -0700688class Flag(Enum):
Ethan Furmanee47e5c2016-08-31 00:12:15 -0700689 """Support for flags"""
Ethan Furmanc16595e2016-09-10 23:36:59 -0700690
691 def _generate_next_value_(name, start, count, last_values):
Ethan Furmanee47e5c2016-08-31 00:12:15 -0700692 """
693 Generate the next value when not given.
694
695 name: the name of the member
HongWeipengbb16fb22019-09-21 13:22:54 +0800696 start: the initial start value or None
Ethan Furmanee47e5c2016-08-31 00:12:15 -0700697 count: the number of existing members
698 last_value: the last value assigned or None
699 """
700 if not count:
701 return start if start is not None else 1
Ethan Furmanc16595e2016-09-10 23:36:59 -0700702 for last_value in reversed(last_values):
703 try:
704 high_bit = _high_bit(last_value)
705 break
Ethan Furman3515dcc2016-09-18 13:15:41 -0700706 except Exception:
Ethan Furmanc16595e2016-09-10 23:36:59 -0700707 raise TypeError('Invalid Flag value: %r' % last_value) from None
Ethan Furmanee47e5c2016-08-31 00:12:15 -0700708 return 2 ** (high_bit+1)
709
710 @classmethod
711 def _missing_(cls, value):
712 original_value = value
713 if value < 0:
714 value = ~value
715 possible_member = cls._create_pseudo_member_(value)
Ethan Furmanee47e5c2016-08-31 00:12:15 -0700716 if original_value < 0:
717 possible_member = ~possible_member
718 return possible_member
719
720 @classmethod
721 def _create_pseudo_member_(cls, value):
Ethan Furman3515dcc2016-09-18 13:15:41 -0700722 """
723 Create a composite member iff value contains only members.
724 """
Ethan Furmanee47e5c2016-08-31 00:12:15 -0700725 pseudo_member = cls._value2member_map_.get(value, None)
726 if pseudo_member is None:
Ethan Furman3515dcc2016-09-18 13:15:41 -0700727 # verify all bits are accounted for
728 _, extra_flags = _decompose(cls, value)
729 if extra_flags:
Walter Dörwald323842c2019-07-18 20:37:13 +0200730 raise ValueError("%r is not a valid %s" % (value, cls.__qualname__))
Ethan Furman3515dcc2016-09-18 13:15:41 -0700731 # construct a singleton enum pseudo-member
Ethan Furmanee47e5c2016-08-31 00:12:15 -0700732 pseudo_member = object.__new__(cls)
733 pseudo_member._name_ = None
734 pseudo_member._value_ = value
Ethan Furman28cf6632017-01-24 12:12:06 -0800735 # use setdefault in case another thread already created a composite
736 # with this value
737 pseudo_member = cls._value2member_map_.setdefault(value, pseudo_member)
Ethan Furmanee47e5c2016-08-31 00:12:15 -0700738 return pseudo_member
739
Ethan Furmanee47e5c2016-08-31 00:12:15 -0700740 def __contains__(self, other):
741 if not isinstance(other, self.__class__):
Rahul Jha94306522018-09-10 23:51:04 +0530742 raise TypeError(
743 "unsupported operand type(s) for 'in': '%s' and '%s'" % (
744 type(other).__qualname__, self.__class__.__qualname__))
Ethan Furmanee47e5c2016-08-31 00:12:15 -0700745 return other._value_ & self._value_ == other._value_
746
Ethan Furmanee47e5c2016-08-31 00:12:15 -0700747 def __repr__(self):
748 cls = self.__class__
749 if self._name_ is not None:
750 return '<%s.%s: %r>' % (cls.__name__, self._name_, self._value_)
Ethan Furman3515dcc2016-09-18 13:15:41 -0700751 members, uncovered = _decompose(cls, self._value_)
Ethan Furman27682d22016-09-04 11:39:01 -0700752 return '<%s.%s: %r>' % (
753 cls.__name__,
754 '|'.join([str(m._name_ or m._value_) for m in members]),
755 self._value_,
756 )
Ethan Furmanee47e5c2016-08-31 00:12:15 -0700757
758 def __str__(self):
759 cls = self.__class__
760 if self._name_ is not None:
761 return '%s.%s' % (cls.__name__, self._name_)
Ethan Furman3515dcc2016-09-18 13:15:41 -0700762 members, uncovered = _decompose(cls, self._value_)
Ethan Furmanee47e5c2016-08-31 00:12:15 -0700763 if len(members) == 1 and members[0]._name_ is None:
764 return '%s.%r' % (cls.__name__, members[0]._value_)
765 else:
766 return '%s.%s' % (
767 cls.__name__,
768 '|'.join([str(m._name_ or m._value_) for m in members]),
769 )
770
Ethan Furman25d94bb2016-09-02 16:32:32 -0700771 def __bool__(self):
772 return bool(self._value_)
773
Ethan Furmanee47e5c2016-08-31 00:12:15 -0700774 def __or__(self, other):
775 if not isinstance(other, self.__class__):
776 return NotImplemented
777 return self.__class__(self._value_ | other._value_)
778
779 def __and__(self, other):
780 if not isinstance(other, self.__class__):
781 return NotImplemented
782 return self.__class__(self._value_ & other._value_)
783
784 def __xor__(self, other):
785 if not isinstance(other, self.__class__):
786 return NotImplemented
787 return self.__class__(self._value_ ^ other._value_)
788
789 def __invert__(self):
Ethan Furman3515dcc2016-09-18 13:15:41 -0700790 members, uncovered = _decompose(self.__class__, self._value_)
Serhiy Storchaka81108372017-09-26 00:55:55 +0300791 inverted = self.__class__(0)
792 for m in self.__class__:
793 if m not in members and not (m._value_ & self._value_):
794 inverted = inverted | m
Ethan Furmanee47e5c2016-08-31 00:12:15 -0700795 return self.__class__(inverted)
796
797
Ethan Furman65a5a472016-09-01 23:55:19 -0700798class IntFlag(int, Flag):
Ethan Furmanee47e5c2016-08-31 00:12:15 -0700799 """Support for integer-based Flags"""
800
801 @classmethod
Ethan Furman3515dcc2016-09-18 13:15:41 -0700802 def _missing_(cls, value):
803 if not isinstance(value, int):
Walter Dörwald323842c2019-07-18 20:37:13 +0200804 raise ValueError("%r is not a valid %s" % (value, cls.__qualname__))
Ethan Furman3515dcc2016-09-18 13:15:41 -0700805 new_member = cls._create_pseudo_member_(value)
806 return new_member
807
808 @classmethod
Ethan Furmanee47e5c2016-08-31 00:12:15 -0700809 def _create_pseudo_member_(cls, value):
810 pseudo_member = cls._value2member_map_.get(value, None)
811 if pseudo_member is None:
Ethan Furman3515dcc2016-09-18 13:15:41 -0700812 need_to_create = [value]
813 # get unaccounted for bits
814 _, extra_flags = _decompose(cls, value)
815 # timer = 10
816 while extra_flags:
817 # timer -= 1
818 bit = _high_bit(extra_flags)
819 flag_value = 2 ** bit
820 if (flag_value not in cls._value2member_map_ and
821 flag_value not in need_to_create
822 ):
823 need_to_create.append(flag_value)
824 if extra_flags == -flag_value:
825 extra_flags = 0
826 else:
827 extra_flags ^= flag_value
828 for value in reversed(need_to_create):
829 # construct singleton pseudo-members
830 pseudo_member = int.__new__(cls, value)
831 pseudo_member._name_ = None
832 pseudo_member._value_ = value
Ethan Furman28cf6632017-01-24 12:12:06 -0800833 # use setdefault in case another thread already created a composite
834 # with this value
835 pseudo_member = cls._value2member_map_.setdefault(value, pseudo_member)
Ethan Furmanee47e5c2016-08-31 00:12:15 -0700836 return pseudo_member
837
Ethan Furmanee47e5c2016-08-31 00:12:15 -0700838 def __or__(self, other):
839 if not isinstance(other, (self.__class__, int)):
840 return NotImplemented
Ethan Furman3515dcc2016-09-18 13:15:41 -0700841 result = self.__class__(self._value_ | self.__class__(other)._value_)
842 return result
Ethan Furmanee47e5c2016-08-31 00:12:15 -0700843
844 def __and__(self, other):
845 if not isinstance(other, (self.__class__, int)):
846 return NotImplemented
847 return self.__class__(self._value_ & self.__class__(other)._value_)
848
849 def __xor__(self, other):
850 if not isinstance(other, (self.__class__, int)):
851 return NotImplemented
852 return self.__class__(self._value_ ^ self.__class__(other)._value_)
853
854 __ror__ = __or__
855 __rand__ = __and__
856 __rxor__ = __xor__
857
858 def __invert__(self):
Ethan Furman3515dcc2016-09-18 13:15:41 -0700859 result = self.__class__(~self._value_)
860 return result
Ethan Furmanee47e5c2016-08-31 00:12:15 -0700861
862
863def _high_bit(value):
Ethan Furman04439532016-09-02 15:50:21 -0700864 """returns index of highest bit, or -1 if value is zero or negative"""
Ethan Furman3515dcc2016-09-18 13:15:41 -0700865 return value.bit_length() - 1
Ethan Furmanee47e5c2016-08-31 00:12:15 -0700866
Ethan Furmanf24bb352013-07-18 17:05:39 -0700867def unique(enumeration):
868 """Class decorator for enumerations ensuring unique member values."""
869 duplicates = []
870 for name, member in enumeration.__members__.items():
871 if name != member.name:
872 duplicates.append((name, member.name))
873 if duplicates:
874 alias_details = ', '.join(
875 ["%s -> %s" % (alias, name) for (alias, name) in duplicates])
876 raise ValueError('duplicate values found in %r: %s' %
877 (enumeration, alias_details))
878 return enumeration
Ethan Furman3515dcc2016-09-18 13:15:41 -0700879
880def _decompose(flag, value):
881 """Extract all members from the value."""
882 # _decompose is only called if the value is not named
883 not_covered = value
884 negative = value < 0
Ethan Furman3515dcc2016-09-18 13:15:41 -0700885 members = []
HongWeipeng0b41a922019-11-27 06:36:02 +0800886 for member in flag:
887 member_value = member.value
Ethan Furman3515dcc2016-09-18 13:15:41 -0700888 if member_value and member_value & value == member_value:
889 members.append(member)
890 not_covered &= ~member_value
HongWeipeng0b41a922019-11-27 06:36:02 +0800891 if not negative:
892 tmp = not_covered
893 while tmp:
894 flag_value = 2 ** _high_bit(tmp)
895 if flag_value in flag._value2member_map_:
896 members.append(flag._value2member_map_[flag_value])
897 not_covered &= ~flag_value
898 tmp &= ~flag_value
Ethan Furman3515dcc2016-09-18 13:15:41 -0700899 if not members and value in flag._value2member_map_:
900 members.append(flag._value2member_map_[value])
901 members.sort(key=lambda m: m._value_, reverse=True)
902 if len(members) > 1 and members[0].value == value:
903 # we have the breakdown, don't need the value member itself
904 members.pop(0)
905 return members, not_covered