blob: 5e97a9e8d810c1b77df7ece149c1f6348951ea11 [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 Furmanc7915072015-09-17 22:55:40 -07004# try _collections first to reduce startup cost
Ethan Furmane5754ab2015-09-17 22:03:52 -07005try:
6 from _collections import OrderedDict
7except ImportError:
8 from collections import OrderedDict
9
10
Ethan Furmanc16595e2016-09-10 23:36:59 -070011__all__ = [
12 'EnumMeta',
13 'Enum', 'IntEnum', 'Flag', 'IntFlag',
14 'auto', 'unique',
15 ]
Ethan Furman6b3d64a2013-06-14 16:55:46 -070016
17
Ethan Furman101e0742013-09-15 12:34:36 -070018def _is_descriptor(obj):
19 """Returns True if obj is a descriptor, False otherwise."""
20 return (
21 hasattr(obj, '__get__') or
22 hasattr(obj, '__set__') or
23 hasattr(obj, '__delete__'))
24
25
Ethan Furman6b3d64a2013-06-14 16:55:46 -070026def _is_dunder(name):
27 """Returns True if a __dunder__ name, False otherwise."""
Miss Islington (bot)8755f0a2019-03-08 13:44:21 -080028 return (len(name) > 4 and
29 name[:2] == name[-2:] == '__' and
30 name[2] != '_' and
31 name[-3] != '_')
Ethan Furman6b3d64a2013-06-14 16:55:46 -070032
33
34def _is_sunder(name):
35 """Returns True if a _sunder_ name, False otherwise."""
Miss Islington (bot)8755f0a2019-03-08 13:44:21 -080036 return (len(name) > 2 and
37 name[0] == name[-1] == '_' and
Ethan Furman6b3d64a2013-06-14 16:55:46 -070038 name[1:2] != '_' and
Miss Islington (bot)8755f0a2019-03-08 13:44:21 -080039 name[-2:-1] != '_')
40
Ethan Furman6b3d64a2013-06-14 16:55:46 -070041
Ethan Furman6b3d64a2013-06-14 16:55:46 -070042def _make_class_unpicklable(cls):
43 """Make the given class un-picklable."""
Ethan Furmanca1b7942014-02-08 11:36:27 -080044 def _break_on_call_reduce(self, proto):
Ethan Furman6b3d64a2013-06-14 16:55:46 -070045 raise TypeError('%r cannot be pickled' % self)
Ethan Furmanca1b7942014-02-08 11:36:27 -080046 cls.__reduce_ex__ = _break_on_call_reduce
Ethan Furman6b3d64a2013-06-14 16:55:46 -070047 cls.__module__ = '<unknown>'
48
Ethan Furman3515dcc2016-09-18 13:15:41 -070049_auto_null = object()
Ethan Furmanc16595e2016-09-10 23:36:59 -070050class auto:
51 """
52 Instances are replaced with an appropriate value in Enum class suites.
53 """
Ethan Furman3515dcc2016-09-18 13:15:41 -070054 value = _auto_null
Ethan Furmanc16595e2016-09-10 23:36:59 -070055
Ethan Furman101e0742013-09-15 12:34:36 -070056
Ethan Furman6b3d64a2013-06-14 16:55:46 -070057class _EnumDict(dict):
Ethan Furman101e0742013-09-15 12:34:36 -070058 """Track enum member order and ensure member names are not reused.
Ethan Furman6b3d64a2013-06-14 16:55:46 -070059
60 EnumMeta will use the names found in self._member_names as the
61 enumeration member names.
62
63 """
64 def __init__(self):
65 super().__init__()
66 self._member_names = []
Ethan Furmanc16595e2016-09-10 23:36:59 -070067 self._last_values = []
Ethan Furmana4b1bb42018-01-22 07:56:37 -080068 self._ignore = []
Ethan Furman6b3d64a2013-06-14 16:55:46 -070069
70 def __setitem__(self, key, value):
Ethan Furman101e0742013-09-15 12:34:36 -070071 """Changes anything not dundered or not a descriptor.
Ethan Furman6b3d64a2013-06-14 16:55:46 -070072
73 If an enum member name is used twice, an error is raised; duplicate
74 values are not checked for.
75
76 Single underscore (sunder) names are reserved.
77
78 """
79 if _is_sunder(key):
Ethan Furmanee47e5c2016-08-31 00:12:15 -070080 if key not in (
Ethan Furman3515dcc2016-09-18 13:15:41 -070081 '_order_', '_create_pseudo_member_',
Ethan Furmana4b1bb42018-01-22 07:56:37 -080082 '_generate_next_value_', '_missing_', '_ignore_',
Ethan Furmanee47e5c2016-08-31 00:12:15 -070083 ):
Ethan Furmane8e61272016-08-20 07:19:31 -070084 raise ValueError('_names_ are reserved for future Enum use')
Ethan Furmanc16595e2016-09-10 23:36:59 -070085 if key == '_generate_next_value_':
86 setattr(self, '_generate_next_value', value)
Ethan Furmana4b1bb42018-01-22 07:56:37 -080087 elif key == '_ignore_':
88 if isinstance(value, str):
89 value = value.replace(',',' ').split()
90 else:
91 value = list(value)
92 self._ignore = value
93 already = set(value) & set(self._member_names)
94 if already:
95 raise ValueError('_ignore_ cannot specify already set names: %r' % (already, ))
Ethan Furman101e0742013-09-15 12:34:36 -070096 elif _is_dunder(key):
Ethan Furmane8e61272016-08-20 07:19:31 -070097 if key == '__order__':
98 key = '_order_'
Ethan Furman101e0742013-09-15 12:34:36 -070099 elif key in self._member_names:
100 # descriptor overwriting an enum?
101 raise TypeError('Attempted to reuse key: %r' % key)
Ethan Furmana4b1bb42018-01-22 07:56:37 -0800102 elif key in self._ignore:
103 pass
Ethan Furman101e0742013-09-15 12:34:36 -0700104 elif not _is_descriptor(value):
105 if key in self:
106 # enum overwriting a descriptor?
Ethan Furmanee47e5c2016-08-31 00:12:15 -0700107 raise TypeError('%r already defined as: %r' % (key, self[key]))
Ethan Furmanc16595e2016-09-10 23:36:59 -0700108 if isinstance(value, auto):
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?)
Miss Islington (bot)8755f0a2019-03-08 13:44:21 -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
172 enum_class._member_map_ = OrderedDict() # 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
orlnub123c0d63bf2018-09-10 19:39:48 +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
orlnub123c0d63bf2018-09-10 19:39:48 +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 Furmandc870522014-02-18 12:37:12 -0800253 for name in ('__repr__', '__str__', '__format__', '__reduce_ex__'):
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700254 class_method = getattr(enum_class, name)
255 obj_method = getattr(member_type, name, None)
256 enum_method = getattr(first_enum, name, None)
257 if obj_method is not None and obj_method is class_method:
258 setattr(enum_class, name, enum_method)
259
260 # replace any other __new__ with our own (as long as Enum is not None,
261 # anyway) -- again, this is to support pickle
262 if Enum is not None:
263 # if the user defined their own __new__, save it before it gets
264 # clobbered in case they subclass later
265 if save_new:
266 enum_class.__new_member__ = __new__
267 enum_class.__new__ = Enum.__new__
Ethan Furmane8e61272016-08-20 07:19:31 -0700268
269 # py3 support for definition order (helps keep py2/py3 code in sync)
270 if _order_ is not None:
271 if isinstance(_order_, str):
272 _order_ = _order_.replace(',', ' ').split()
273 if _order_ != enum_class._member_names_:
274 raise TypeError('member order does not match _order_')
275
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700276 return enum_class
277
Ethan Furman5de67b12016-04-13 23:52:09 -0700278 def __bool__(self):
279 """
280 classes/types should always be True.
281 """
282 return True
283
Ethan Furmand9925a12014-09-16 20:35:55 -0700284 def __call__(cls, value, names=None, *, module=None, qualname=None, type=None, start=1):
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700285 """Either returns an existing member, or creates a new enum class.
286
287 This method is used both when an enum class is given a value to match
288 to an enumeration member (i.e. Color(3)) and for the functional API
Ethan Furman23bb6f42016-11-21 09:22:05 -0800289 (i.e. Color = Enum('Color', names='RED GREEN BLUE')).
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700290
Ethan Furman2da95042014-03-03 12:42:52 -0800291 When used for the functional API:
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700292
Ethan Furman2da95042014-03-03 12:42:52 -0800293 `value` will be the name of the new class.
294
295 `names` should be either a string of white-space/comma delimited names
Ethan Furmand9925a12014-09-16 20:35:55 -0700296 (values will start at `start`), or an iterator/mapping of name, value pairs.
Ethan Furman2da95042014-03-03 12:42:52 -0800297
298 `module` should be set to the module this class is being created in;
299 if it is not set, an attempt to find that module will be made, but if
300 it fails the class will not be picklable.
301
302 `qualname` should be set to the actual location this class can be found
303 at in its module; by default it is set to the global scope. If this is
304 not correct, unpickling will fail in some circumstances.
305
306 `type`, if set, will be mixed in as the first base class.
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700307
308 """
309 if names is None: # simple value lookup
310 return cls.__new__(cls, value)
311 # otherwise, functional API: we're creating a new Enum type
Ethan Furmand9925a12014-09-16 20:35:55 -0700312 return cls._create_(value, names, module=module, qualname=qualname, type=type, start=start)
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700313
314 def __contains__(cls, member):
Ethan Furman37151762018-04-11 18:56:25 -0700315 if not isinstance(member, Enum):
316 import warnings
317 warnings.warn(
318 "using non-Enums in containment checks will raise "
319 "TypeError in Python 3.8",
320 DeprecationWarning, 2)
Ethan Furman0081f232014-09-16 17:31:23 -0700321 return isinstance(member, cls) and member._name_ in cls._member_map_
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700322
Ethan Furman64a99722013-09-22 16:18:19 -0700323 def __delattr__(cls, attr):
324 # nicer error message when someone tries to delete an attribute
325 # (see issue19025).
326 if attr in cls._member_map_:
327 raise AttributeError(
328 "%s: cannot delete Enum member." % cls.__name__)
329 super().__delattr__(attr)
330
Ethan Furman388a3922013-08-12 06:51:41 -0700331 def __dir__(self):
Ethan Furman64a99722013-09-22 16:18:19 -0700332 return (['__class__', '__doc__', '__members__', '__module__'] +
333 self._member_names_)
Ethan Furman388a3922013-08-12 06:51:41 -0700334
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700335 def __getattr__(cls, name):
336 """Return the enum member matching `name`
337
338 We use __getattr__ instead of descriptors or inserting into the enum
339 class' __dict__ in order to support `name` and `value` being both
340 properties for enum members (which live in the class' __dict__) and
341 enum members themselves.
342
343 """
344 if _is_dunder(name):
345 raise AttributeError(name)
346 try:
Ethan Furman520ad572013-07-19 19:47:21 -0700347 return cls._member_map_[name]
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700348 except KeyError:
349 raise AttributeError(name) from None
350
351 def __getitem__(cls, name):
Ethan Furman520ad572013-07-19 19:47:21 -0700352 return cls._member_map_[name]
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700353
354 def __iter__(cls):
Ethan Furman520ad572013-07-19 19:47:21 -0700355 return (cls._member_map_[name] for name in cls._member_names_)
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700356
357 def __len__(cls):
Ethan Furman520ad572013-07-19 19:47:21 -0700358 return len(cls._member_names_)
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700359
Ethan Furman2131a4a2013-09-14 18:11:24 -0700360 @property
361 def __members__(cls):
362 """Returns a mapping of member name->value.
363
364 This mapping lists all enum members, including aliases. Note that this
365 is a read-only view of the internal mapping.
366
367 """
368 return MappingProxyType(cls._member_map_)
369
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700370 def __repr__(cls):
371 return "<enum %r>" % cls.__name__
372
Ethan Furman2131a4a2013-09-14 18:11:24 -0700373 def __reversed__(cls):
374 return (cls._member_map_[name] for name in reversed(cls._member_names_))
375
Ethan Furmanf203f2d2013-09-06 07:16:48 -0700376 def __setattr__(cls, name, value):
377 """Block attempts to reassign Enum members.
378
379 A simple assignment to the class namespace only changes one of the
380 several possible ways to get an Enum member from the Enum class,
381 resulting in an inconsistent Enumeration.
382
383 """
384 member_map = cls.__dict__.get('_member_map_', {})
385 if name in member_map:
386 raise AttributeError('Cannot reassign members.')
387 super().__setattr__(name, value)
388
Miss Islington (bot)3bcca482018-04-16 13:49:34 -0700389 def _create_(cls, class_name, names, *, module=None, qualname=None, type=None, start=1):
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700390 """Convenience method to create a new Enum class.
391
392 `names` can be:
393
394 * A string containing member names, separated either with spaces or
Ethan Furmand9925a12014-09-16 20:35:55 -0700395 commas. Values are incremented by 1 from `start`.
396 * An iterable of member names. Values are incremented by 1 from `start`.
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700397 * An iterable of (member name, value) pairs.
Ethan Furmand9925a12014-09-16 20:35:55 -0700398 * A mapping of member name -> value pairs.
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700399
400 """
401 metacls = cls.__class__
402 bases = (cls, ) if type is None else (type, cls)
Ethan Furmanee47e5c2016-08-31 00:12:15 -0700403 _, first_enum = cls._get_mixins_(bases)
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700404 classdict = metacls.__prepare__(class_name, bases)
405
406 # special processing needed for names?
407 if isinstance(names, str):
408 names = names.replace(',', ' ').split()
Dong-hee Nadcc8ce42017-06-22 01:52:32 +0900409 if isinstance(names, (tuple, list)) and names and isinstance(names[0], str):
Ethan Furmanee47e5c2016-08-31 00:12:15 -0700410 original_names, names = names, []
Ethan Furmanc16595e2016-09-10 23:36:59 -0700411 last_values = []
Ethan Furmanee47e5c2016-08-31 00:12:15 -0700412 for count, name in enumerate(original_names):
Ethan Furmanc16595e2016-09-10 23:36:59 -0700413 value = first_enum._generate_next_value_(name, start, count, last_values[:])
414 last_values.append(value)
415 names.append((name, value))
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700416
417 # Here, names is either an iterable of (name, value) or a mapping.
418 for item in names:
419 if isinstance(item, str):
420 member_name, member_value = item, names[item]
421 else:
422 member_name, member_value = item
423 classdict[member_name] = member_value
424 enum_class = metacls.__new__(metacls, class_name, bases, classdict)
425
426 # TODO: replace the frame hack if a blessed way to know the calling
427 # module is ever developed
428 if module is None:
429 try:
430 module = sys._getframe(2).f_globals['__name__']
Miss Islington (bot)1c798912019-01-30 09:36:51 -0800431 except (AttributeError, ValueError, KeyError) as exc:
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700432 pass
433 if module is None:
434 _make_class_unpicklable(enum_class)
435 else:
436 enum_class.__module__ = module
Ethan Furmanca1b7942014-02-08 11:36:27 -0800437 if qualname is not None:
438 enum_class.__qualname__ = qualname
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700439
440 return enum_class
441
442 @staticmethod
443 def _get_mixins_(bases):
444 """Returns the type for creating enum members, and the first inherited
445 enum class.
446
447 bases: the tuple of bases that was given to __new__
448
449 """
450 if not bases:
451 return object, Enum
452
Ethan Furman0c076ca2018-09-21 22:26:32 -0700453 def _find_data_type(bases):
454 for chain in bases:
455 for base in chain.__mro__:
456 if base is object:
457 continue
458 elif '__new__' in base.__dict__:
Miss Islington (bot)453b3b02018-10-06 00:43:20 -0700459 if issubclass(base, Enum):
Ethan Furman0c076ca2018-09-21 22:26:32 -0700460 continue
461 return base
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700462
Ethan Furman0c076ca2018-09-21 22:26:32 -0700463 # ensure final parent class is an Enum derivative, find any concrete
464 # data type, and check that Enum has no members
465 first_enum = bases[-1]
466 if not issubclass(first_enum, Enum):
467 raise TypeError("new enumerations should be created as "
468 "`EnumName([mixin_type, ...] [data_type,] enum_type)`")
469 member_type = _find_data_type(bases) or object
470 if first_enum._member_names_:
471 raise TypeError("Cannot extend enumerations")
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700472 return member_type, first_enum
473
474 @staticmethod
475 def _find_new_(classdict, member_type, first_enum):
476 """Returns the __new__ to be used for creating the enum members.
477
478 classdict: the class dictionary given to __new__
479 member_type: the data type whose __new__ will be used by default
480 first_enum: enumeration to check for an overriding __new__
481
482 """
483 # now find the correct __new__, checking to see of one was defined
484 # by the user; also check earlier enum classes in case a __new__ was
485 # saved as __new_member__
486 __new__ = classdict.get('__new__', None)
487
488 # should __new__ be saved as __new_member__ later?
489 save_new = __new__ is not None
490
491 if __new__ is None:
492 # check all possibles for __new_member__ before falling back to
493 # __new__
494 for method in ('__new_member__', '__new__'):
495 for possible in (member_type, first_enum):
496 target = getattr(possible, method, None)
497 if target not in {
498 None,
499 None.__new__,
500 object.__new__,
501 Enum.__new__,
502 }:
503 __new__ = target
504 break
505 if __new__ is not None:
506 break
507 else:
508 __new__ = object.__new__
509
510 # if a non-object.__new__ is used then whatever value/tuple was
511 # assigned to the enum member name will be passed to __new__ and to the
512 # new enum member's __init__
513 if __new__ is object.__new__:
514 use_args = False
515 else:
516 use_args = True
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700517 return __new__, save_new, use_args
518
519
520class Enum(metaclass=EnumMeta):
521 """Generic enumeration.
522
523 Derive from this class to define new enumerations.
524
525 """
526 def __new__(cls, value):
527 # all enum instances are actually created during class construction
528 # without calling this method; this method is called by the metaclass'
529 # __call__ (i.e. Color(3) ), and by pickle
530 if type(value) is cls:
Ethan Furman23bb6f42016-11-21 09:22:05 -0800531 # For lookups like Color(Color.RED)
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700532 return value
533 # by-value search for a matching enum member
534 # see if it's in the reverse mapping (for hashable values)
Ethan Furman2aa27322013-07-19 19:35:56 -0700535 try:
Miss Islington (bot)705b5992018-12-26 12:48:55 -0800536 return cls._value2member_map_[value]
537 except KeyError:
538 # Not found, no need to do long O(n) search
539 pass
Ethan Furman2aa27322013-07-19 19:35:56 -0700540 except TypeError:
541 # not there, now do long search -- O(n) behavior
Ethan Furman520ad572013-07-19 19:47:21 -0700542 for member in cls._member_map_.values():
Ethan Furman0081f232014-09-16 17:31:23 -0700543 if member._value_ == value:
Ethan Furman2aa27322013-07-19 19:35:56 -0700544 return member
Ethan Furmanee47e5c2016-08-31 00:12:15 -0700545 # still not found -- try _missing_ hook
Victor Stinner0f2fc8b2018-10-20 01:49:30 +0200546 try:
547 exc = None
548 result = cls._missing_(value)
549 except Exception as e:
550 exc = e
551 result = None
552 if isinstance(result, cls):
553 return result
554 else:
555 ve_exc = ValueError("%r is not a valid %s" % (value, cls.__name__))
556 if result is None and exc is None:
557 raise ve_exc
558 elif exc is None:
559 exc = TypeError(
560 'error in %s._missing_: returned %r instead of None or a valid member'
561 % (cls.__name__, result)
562 )
563 exc.__context__ = ve_exc
564 raise exc
Ethan Furmanee47e5c2016-08-31 00:12:15 -0700565
Ethan Furmanc16595e2016-09-10 23:36:59 -0700566 def _generate_next_value_(name, start, count, last_values):
567 for last_value in reversed(last_values):
568 try:
569 return last_value + 1
570 except TypeError:
571 pass
572 else:
Ethan Furmanee47e5c2016-08-31 00:12:15 -0700573 return start
Ethan Furmanc16595e2016-09-10 23:36:59 -0700574
Ethan Furmanee47e5c2016-08-31 00:12:15 -0700575 @classmethod
576 def _missing_(cls, value):
Ethan Furman0081f232014-09-16 17:31:23 -0700577 raise ValueError("%r is not a valid %s" % (value, cls.__name__))
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700578
579 def __repr__(self):
580 return "<%s.%s: %r>" % (
Ethan Furman520ad572013-07-19 19:47:21 -0700581 self.__class__.__name__, self._name_, self._value_)
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700582
583 def __str__(self):
Ethan Furman520ad572013-07-19 19:47:21 -0700584 return "%s.%s" % (self.__class__.__name__, self._name_)
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700585
Ethan Furman388a3922013-08-12 06:51:41 -0700586 def __dir__(self):
Ethan Furman0ae550b2014-10-14 08:58:32 -0700587 added_behavior = [
588 m
589 for cls in self.__class__.mro()
590 for m in cls.__dict__
Ethan Furman354ecf12015-03-11 08:43:12 -0700591 if m[0] != '_' and m not in self._member_map_
Ethan Furman0ae550b2014-10-14 08:58:32 -0700592 ]
Ethan Furmanec5f8eb2014-10-21 13:40:35 -0700593 return (['__class__', '__doc__', '__module__'] + added_behavior)
Ethan Furman388a3922013-08-12 06:51:41 -0700594
Ethan Furmanec15a822013-08-31 19:17:41 -0700595 def __format__(self, format_spec):
596 # mixed-in Enums should use the mixed-in type's __format__, otherwise
597 # we can get strange results with the Enum name showing up instead of
598 # the value
599
600 # pure Enum branch
601 if self._member_type_ is object:
602 cls = str
603 val = str(self)
604 # mix-in branch
605 else:
606 cls = self._member_type_
Ethan Furman0081f232014-09-16 17:31:23 -0700607 val = self._value_
Ethan Furmanec15a822013-08-31 19:17:41 -0700608 return cls.__format__(val, format_spec)
609
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700610 def __hash__(self):
Ethan Furman520ad572013-07-19 19:47:21 -0700611 return hash(self._name_)
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700612
Ethan Furmanca1b7942014-02-08 11:36:27 -0800613 def __reduce_ex__(self, proto):
Ethan Furmandc870522014-02-18 12:37:12 -0800614 return self.__class__, (self._value_, )
Ethan Furmanca1b7942014-02-08 11:36:27 -0800615
Ethan Furman33918c12013-09-27 23:02:02 -0700616 # DynamicClassAttribute is used to provide access to the `name` and
617 # `value` properties of enum members while keeping some measure of
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700618 # protection from modification, while still allowing for an enumeration
619 # to have members named `name` and `value`. This works because enumeration
620 # members are not set directly on the enum class -- __getattr__ is
621 # used to look them up.
622
Ethan Furmane03ea372013-09-25 07:14:41 -0700623 @DynamicClassAttribute
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700624 def name(self):
Ethan Furmanc850f342013-09-15 16:59:35 -0700625 """The name of the Enum member."""
Ethan Furman520ad572013-07-19 19:47:21 -0700626 return self._name_
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700627
Ethan Furmane03ea372013-09-25 07:14:41 -0700628 @DynamicClassAttribute
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700629 def value(self):
Ethan Furmanc850f342013-09-15 16:59:35 -0700630 """The value of the Enum member."""
Ethan Furman520ad572013-07-19 19:47:21 -0700631 return self._value_
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700632
Ethan Furman24e837f2015-03-18 17:27:57 -0700633 @classmethod
634 def _convert(cls, name, module, filter, source=None):
635 """
636 Create a new Enum subclass that replaces a collection of global constants
637 """
638 # convert all constants from source (or module) that pass filter() to
639 # a new Enum called name, and export the enum and its members back to
640 # module;
641 # also, replace the __reduce_ex__ method so unpickling works in
642 # previous Python versions
643 module_globals = vars(sys.modules[module])
644 if source:
645 source = vars(source)
646 else:
647 source = module_globals
Gregory P. Smith ext:(%20%5BGoogle%20Inc.%5D)6f20bd62016-06-03 19:14:52 +0000648 # We use an OrderedDict of sorted source keys so that the
649 # _value2member_map is populated in the same order every time
650 # for a consistent reverse mapping of number to name when there
651 # are multiple names for the same number rather than varying
652 # between runs due to hash randomization of the module dictionary.
Ethan Furman06339e72016-09-11 13:25:26 -0700653 members = [
654 (name, source[name])
655 for name in source.keys()
656 if filter(name)]
657 try:
658 # sort by value
659 members.sort(key=lambda t: (t[1], t[0]))
660 except TypeError:
661 # unless some values aren't comparable, in which case sort by name
662 members.sort(key=lambda t: t[0])
Ethan Furman24e837f2015-03-18 17:27:57 -0700663 cls = cls(name, members, module=module)
664 cls.__reduce_ex__ = _reduce_ex_by_name
665 module_globals.update(cls.__members__)
666 module_globals[name] = cls
667 return cls
668
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700669
670class IntEnum(int, Enum):
671 """Enum where members are also (and must be) ints"""
Ethan Furmanf24bb352013-07-18 17:05:39 -0700672
673
Ethan Furman24e837f2015-03-18 17:27:57 -0700674def _reduce_ex_by_name(self, proto):
675 return self.name
676
Ethan Furman65a5a472016-09-01 23:55:19 -0700677class Flag(Enum):
Ethan Furmanee47e5c2016-08-31 00:12:15 -0700678 """Support for flags"""
Ethan Furmanc16595e2016-09-10 23:36:59 -0700679
680 def _generate_next_value_(name, start, count, last_values):
Ethan Furmanee47e5c2016-08-31 00:12:15 -0700681 """
682 Generate the next value when not given.
683
684 name: the name of the member
685 start: the initital start value or None
686 count: the number of existing members
687 last_value: the last value assigned or None
688 """
689 if not count:
690 return start if start is not None else 1
Ethan Furmanc16595e2016-09-10 23:36:59 -0700691 for last_value in reversed(last_values):
692 try:
693 high_bit = _high_bit(last_value)
694 break
Ethan Furman3515dcc2016-09-18 13:15:41 -0700695 except Exception:
Ethan Furmanc16595e2016-09-10 23:36:59 -0700696 raise TypeError('Invalid Flag value: %r' % last_value) from None
Ethan Furmanee47e5c2016-08-31 00:12:15 -0700697 return 2 ** (high_bit+1)
698
699 @classmethod
700 def _missing_(cls, value):
701 original_value = value
702 if value < 0:
703 value = ~value
704 possible_member = cls._create_pseudo_member_(value)
Ethan Furmanee47e5c2016-08-31 00:12:15 -0700705 if original_value < 0:
706 possible_member = ~possible_member
707 return possible_member
708
709 @classmethod
710 def _create_pseudo_member_(cls, value):
Ethan Furman3515dcc2016-09-18 13:15:41 -0700711 """
712 Create a composite member iff value contains only members.
713 """
Ethan Furmanee47e5c2016-08-31 00:12:15 -0700714 pseudo_member = cls._value2member_map_.get(value, None)
715 if pseudo_member is None:
Ethan Furman3515dcc2016-09-18 13:15:41 -0700716 # verify all bits are accounted for
717 _, extra_flags = _decompose(cls, value)
718 if extra_flags:
719 raise ValueError("%r is not a valid %s" % (value, cls.__name__))
720 # construct a singleton enum pseudo-member
Ethan Furmanee47e5c2016-08-31 00:12:15 -0700721 pseudo_member = object.__new__(cls)
722 pseudo_member._name_ = None
723 pseudo_member._value_ = value
Ethan Furman28cf6632017-01-24 12:12:06 -0800724 # use setdefault in case another thread already created a composite
725 # with this value
726 pseudo_member = cls._value2member_map_.setdefault(value, pseudo_member)
Ethan Furmanee47e5c2016-08-31 00:12:15 -0700727 return pseudo_member
728
Ethan Furmanee47e5c2016-08-31 00:12:15 -0700729 def __contains__(self, other):
730 if not isinstance(other, self.__class__):
Ethan Furman37151762018-04-11 18:56:25 -0700731 import warnings
732 warnings.warn(
733 "using non-Flags in containment checks will raise "
734 "TypeError in Python 3.8",
735 DeprecationWarning, 2)
736 return False
Ethan Furmanee47e5c2016-08-31 00:12:15 -0700737 return other._value_ & self._value_ == other._value_
738
Ethan Furmanee47e5c2016-08-31 00:12:15 -0700739 def __repr__(self):
740 cls = self.__class__
741 if self._name_ is not None:
742 return '<%s.%s: %r>' % (cls.__name__, self._name_, self._value_)
Ethan Furman3515dcc2016-09-18 13:15:41 -0700743 members, uncovered = _decompose(cls, self._value_)
Ethan Furman27682d22016-09-04 11:39:01 -0700744 return '<%s.%s: %r>' % (
745 cls.__name__,
746 '|'.join([str(m._name_ or m._value_) for m in members]),
747 self._value_,
748 )
Ethan Furmanee47e5c2016-08-31 00:12:15 -0700749
750 def __str__(self):
751 cls = self.__class__
752 if self._name_ is not None:
753 return '%s.%s' % (cls.__name__, self._name_)
Ethan Furman3515dcc2016-09-18 13:15:41 -0700754 members, uncovered = _decompose(cls, self._value_)
Ethan Furmanee47e5c2016-08-31 00:12:15 -0700755 if len(members) == 1 and members[0]._name_ is None:
756 return '%s.%r' % (cls.__name__, members[0]._value_)
757 else:
758 return '%s.%s' % (
759 cls.__name__,
760 '|'.join([str(m._name_ or m._value_) for m in members]),
761 )
762
Ethan Furman25d94bb2016-09-02 16:32:32 -0700763 def __bool__(self):
764 return bool(self._value_)
765
Ethan Furmanee47e5c2016-08-31 00:12:15 -0700766 def __or__(self, other):
767 if not isinstance(other, self.__class__):
768 return NotImplemented
769 return self.__class__(self._value_ | other._value_)
770
771 def __and__(self, other):
772 if not isinstance(other, self.__class__):
773 return NotImplemented
774 return self.__class__(self._value_ & other._value_)
775
776 def __xor__(self, other):
777 if not isinstance(other, self.__class__):
778 return NotImplemented
779 return self.__class__(self._value_ ^ other._value_)
780
781 def __invert__(self):
Ethan Furman3515dcc2016-09-18 13:15:41 -0700782 members, uncovered = _decompose(self.__class__, self._value_)
Serhiy Storchaka81108372017-09-26 00:55:55 +0300783 inverted = self.__class__(0)
784 for m in self.__class__:
785 if m not in members and not (m._value_ & self._value_):
786 inverted = inverted | m
Ethan Furmanee47e5c2016-08-31 00:12:15 -0700787 return self.__class__(inverted)
788
789
Ethan Furman65a5a472016-09-01 23:55:19 -0700790class IntFlag(int, Flag):
Ethan Furmanee47e5c2016-08-31 00:12:15 -0700791 """Support for integer-based Flags"""
792
793 @classmethod
Ethan Furman3515dcc2016-09-18 13:15:41 -0700794 def _missing_(cls, value):
795 if not isinstance(value, int):
796 raise ValueError("%r is not a valid %s" % (value, cls.__name__))
797 new_member = cls._create_pseudo_member_(value)
798 return new_member
799
800 @classmethod
Ethan Furmanee47e5c2016-08-31 00:12:15 -0700801 def _create_pseudo_member_(cls, value):
802 pseudo_member = cls._value2member_map_.get(value, None)
803 if pseudo_member is None:
Ethan Furman3515dcc2016-09-18 13:15:41 -0700804 need_to_create = [value]
805 # get unaccounted for bits
806 _, extra_flags = _decompose(cls, value)
807 # timer = 10
808 while extra_flags:
809 # timer -= 1
810 bit = _high_bit(extra_flags)
811 flag_value = 2 ** bit
812 if (flag_value not in cls._value2member_map_ and
813 flag_value not in need_to_create
814 ):
815 need_to_create.append(flag_value)
816 if extra_flags == -flag_value:
817 extra_flags = 0
818 else:
819 extra_flags ^= flag_value
820 for value in reversed(need_to_create):
821 # construct singleton pseudo-members
822 pseudo_member = int.__new__(cls, value)
823 pseudo_member._name_ = None
824 pseudo_member._value_ = value
Ethan Furman28cf6632017-01-24 12:12:06 -0800825 # use setdefault in case another thread already created a composite
826 # with this value
827 pseudo_member = cls._value2member_map_.setdefault(value, pseudo_member)
Ethan Furmanee47e5c2016-08-31 00:12:15 -0700828 return pseudo_member
829
Ethan Furmanee47e5c2016-08-31 00:12:15 -0700830 def __or__(self, other):
831 if not isinstance(other, (self.__class__, int)):
832 return NotImplemented
Ethan Furman3515dcc2016-09-18 13:15:41 -0700833 result = self.__class__(self._value_ | self.__class__(other)._value_)
834 return result
Ethan Furmanee47e5c2016-08-31 00:12:15 -0700835
836 def __and__(self, other):
837 if not isinstance(other, (self.__class__, int)):
838 return NotImplemented
839 return self.__class__(self._value_ & self.__class__(other)._value_)
840
841 def __xor__(self, other):
842 if not isinstance(other, (self.__class__, int)):
843 return NotImplemented
844 return self.__class__(self._value_ ^ self.__class__(other)._value_)
845
846 __ror__ = __or__
847 __rand__ = __and__
848 __rxor__ = __xor__
849
850 def __invert__(self):
Ethan Furman3515dcc2016-09-18 13:15:41 -0700851 result = self.__class__(~self._value_)
852 return result
Ethan Furmanee47e5c2016-08-31 00:12:15 -0700853
854
855def _high_bit(value):
Ethan Furman04439532016-09-02 15:50:21 -0700856 """returns index of highest bit, or -1 if value is zero or negative"""
Ethan Furman3515dcc2016-09-18 13:15:41 -0700857 return value.bit_length() - 1
Ethan Furmanee47e5c2016-08-31 00:12:15 -0700858
Ethan Furmanf24bb352013-07-18 17:05:39 -0700859def unique(enumeration):
860 """Class decorator for enumerations ensuring unique member values."""
861 duplicates = []
862 for name, member in enumeration.__members__.items():
863 if name != member.name:
864 duplicates.append((name, member.name))
865 if duplicates:
866 alias_details = ', '.join(
867 ["%s -> %s" % (alias, name) for (alias, name) in duplicates])
868 raise ValueError('duplicate values found in %r: %s' %
869 (enumeration, alias_details))
870 return enumeration
Ethan Furman3515dcc2016-09-18 13:15:41 -0700871
872def _decompose(flag, value):
873 """Extract all members from the value."""
874 # _decompose is only called if the value is not named
875 not_covered = value
876 negative = value < 0
Ethan Furman28cf6632017-01-24 12:12:06 -0800877 # issue29167: wrap accesses to _value2member_map_ in a list to avoid race
Ville Skyttä49b27342017-08-03 09:00:59 +0300878 # conditions between iterating over it and having more pseudo-
Ethan Furman28cf6632017-01-24 12:12:06 -0800879 # members added to it
Ethan Furman3515dcc2016-09-18 13:15:41 -0700880 if negative:
881 # only check for named flags
882 flags_to_check = [
883 (m, v)
Ethan Furman28cf6632017-01-24 12:12:06 -0800884 for v, m in list(flag._value2member_map_.items())
Ethan Furman3515dcc2016-09-18 13:15:41 -0700885 if m.name is not None
886 ]
887 else:
888 # check for named flags and powers-of-two flags
889 flags_to_check = [
890 (m, v)
Ethan Furman28cf6632017-01-24 12:12:06 -0800891 for v, m in list(flag._value2member_map_.items())
Ethan Furman3515dcc2016-09-18 13:15:41 -0700892 if m.name is not None or _power_of_two(v)
893 ]
894 members = []
895 for member, member_value in flags_to_check:
896 if member_value and member_value & value == member_value:
897 members.append(member)
898 not_covered &= ~member_value
899 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
906
907def _power_of_two(value):
908 if value < 1:
909 return False
910 return value == 2 ** _high_bit(value)