blob: 69b41fe7cbc08cafe1cacbc91f66c44e77b07892 [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."""
28 return (name[:2] == name[-2:] == '__' and
29 name[2:3] != '_' and
Ethan Furman648f8602013-10-06 17:19:54 -070030 name[-3:-2] != '_' and
31 len(name) > 4)
Ethan Furman6b3d64a2013-06-14 16:55:46 -070032
33
34def _is_sunder(name):
35 """Returns True if a _sunder_ name, False otherwise."""
36 return (name[0] == name[-1] == '_' and
37 name[1:2] != '_' and
Ethan Furman648f8602013-10-06 17:19:54 -070038 name[-2:-1] != '_' and
39 len(name) > 2)
Ethan Furman6b3d64a2013-06-14 16:55:46 -070040
Ethan Furman6b3d64a2013-06-14 16:55:46 -070041def _make_class_unpicklable(cls):
42 """Make the given class un-picklable."""
Ethan Furmanca1b7942014-02-08 11:36:27 -080043 def _break_on_call_reduce(self, proto):
Ethan Furman6b3d64a2013-06-14 16:55:46 -070044 raise TypeError('%r cannot be pickled' % self)
Ethan Furmanca1b7942014-02-08 11:36:27 -080045 cls.__reduce_ex__ = _break_on_call_reduce
Ethan Furman6b3d64a2013-06-14 16:55:46 -070046 cls.__module__ = '<unknown>'
47
Ethan Furman3515dcc2016-09-18 13:15:41 -070048_auto_null = object()
Ethan Furmanc16595e2016-09-10 23:36:59 -070049class auto:
50 """
51 Instances are replaced with an appropriate value in Enum class suites.
52 """
Ethan Furman3515dcc2016-09-18 13:15:41 -070053 value = _auto_null
Ethan Furmanc16595e2016-09-10 23:36:59 -070054
Ethan Furman101e0742013-09-15 12:34:36 -070055
Ethan Furman6b3d64a2013-06-14 16:55:46 -070056class _EnumDict(dict):
Ethan Furman101e0742013-09-15 12:34:36 -070057 """Track enum member order and ensure member names are not reused.
Ethan Furman6b3d64a2013-06-14 16:55:46 -070058
59 EnumMeta will use the names found in self._member_names as the
60 enumeration member names.
61
62 """
63 def __init__(self):
64 super().__init__()
65 self._member_names = []
Ethan Furmanc16595e2016-09-10 23:36:59 -070066 self._last_values = []
Ethan Furmana4b1bb42018-01-22 07:56:37 -080067 self._ignore = []
Ethan Furman6b3d64a2013-06-14 16:55:46 -070068
69 def __setitem__(self, key, value):
Ethan Furman101e0742013-09-15 12:34:36 -070070 """Changes anything not dundered or not a descriptor.
Ethan Furman6b3d64a2013-06-14 16:55:46 -070071
72 If an enum member name is used twice, an error is raised; duplicate
73 values are not checked for.
74
75 Single underscore (sunder) names are reserved.
76
77 """
78 if _is_sunder(key):
Ethan Furmanee47e5c2016-08-31 00:12:15 -070079 if key not in (
Ethan Furman3515dcc2016-09-18 13:15:41 -070080 '_order_', '_create_pseudo_member_',
Ethan Furmana4b1bb42018-01-22 07:56:37 -080081 '_generate_next_value_', '_missing_', '_ignore_',
Ethan Furmanee47e5c2016-08-31 00:12:15 -070082 ):
Ethan Furmane8e61272016-08-20 07:19:31 -070083 raise ValueError('_names_ are reserved for future Enum use')
Ethan Furmanc16595e2016-09-10 23:36:59 -070084 if key == '_generate_next_value_':
85 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 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?)
Ethan Furmanee47e5c2016-08-31 00:12:15 -0700159 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
171 enum_class._member_map_ = OrderedDict() # 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
orlnub123c0d63bf2018-09-10 19:39:48 +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
orlnub123c0d63bf2018-09-10 19:39:48 +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 Furmandc870522014-02-18 12:37:12 -0800252 for name in ('__repr__', '__str__', '__format__', '__reduce_ex__'):
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700253 class_method = getattr(enum_class, name)
254 obj_method = getattr(member_type, name, None)
255 enum_method = getattr(first_enum, name, None)
256 if obj_method is not None and obj_method is class_method:
257 setattr(enum_class, name, enum_method)
258
259 # replace any other __new__ with our own (as long as Enum is not None,
260 # anyway) -- again, this is to support pickle
261 if Enum is not None:
262 # if the user defined their own __new__, save it before it gets
263 # clobbered in case they subclass later
264 if save_new:
265 enum_class.__new_member__ = __new__
266 enum_class.__new__ = Enum.__new__
Ethan Furmane8e61272016-08-20 07:19:31 -0700267
268 # py3 support for definition order (helps keep py2/py3 code in sync)
269 if _order_ is not None:
270 if isinstance(_order_, str):
271 _order_ = _order_.replace(',', ' ').split()
272 if _order_ != enum_class._member_names_:
273 raise TypeError('member order does not match _order_')
274
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700275 return enum_class
276
Ethan Furman5de67b12016-04-13 23:52:09 -0700277 def __bool__(self):
278 """
279 classes/types should always be True.
280 """
281 return True
282
Ethan Furmand9925a12014-09-16 20:35:55 -0700283 def __call__(cls, value, names=None, *, module=None, qualname=None, type=None, start=1):
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700284 """Either returns an existing member, or creates a new enum class.
285
286 This method is used both when an enum class is given a value to match
287 to an enumeration member (i.e. Color(3)) and for the functional API
Ethan Furman23bb6f42016-11-21 09:22:05 -0800288 (i.e. Color = Enum('Color', names='RED GREEN BLUE')).
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700289
Ethan Furman2da95042014-03-03 12:42:52 -0800290 When used for the functional API:
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700291
Ethan Furman2da95042014-03-03 12:42:52 -0800292 `value` will be the name of the new class.
293
294 `names` should be either a string of white-space/comma delimited names
Ethan Furmand9925a12014-09-16 20:35:55 -0700295 (values will start at `start`), or an iterator/mapping of name, value pairs.
Ethan Furman2da95042014-03-03 12:42:52 -0800296
297 `module` should be set to the module this class is being created in;
298 if it is not set, an attempt to find that module will be made, but if
299 it fails the class will not be picklable.
300
301 `qualname` should be set to the actual location this class can be found
302 at in its module; by default it is set to the global scope. If this is
303 not correct, unpickling will fail in some circumstances.
304
305 `type`, if set, will be mixed in as the first base class.
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700306
307 """
308 if names is None: # simple value lookup
309 return cls.__new__(cls, value)
310 # otherwise, functional API: we're creating a new Enum type
Ethan Furmand9925a12014-09-16 20:35:55 -0700311 return cls._create_(value, names, module=module, qualname=qualname, type=type, start=start)
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700312
313 def __contains__(cls, member):
Ethan Furman37151762018-04-11 18:56:25 -0700314 if not isinstance(member, Enum):
315 import warnings
316 warnings.warn(
317 "using non-Enums in containment checks will raise "
318 "TypeError in Python 3.8",
319 DeprecationWarning, 2)
Ethan Furman0081f232014-09-16 17:31:23 -0700320 return isinstance(member, cls) and member._name_ in cls._member_map_
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700321
Ethan Furman64a99722013-09-22 16:18:19 -0700322 def __delattr__(cls, attr):
323 # nicer error message when someone tries to delete an attribute
324 # (see issue19025).
325 if attr in cls._member_map_:
326 raise AttributeError(
327 "%s: cannot delete Enum member." % cls.__name__)
328 super().__delattr__(attr)
329
Ethan Furman388a3922013-08-12 06:51:41 -0700330 def __dir__(self):
Ethan Furman64a99722013-09-22 16:18:19 -0700331 return (['__class__', '__doc__', '__members__', '__module__'] +
332 self._member_names_)
Ethan Furman388a3922013-08-12 06:51:41 -0700333
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700334 def __getattr__(cls, name):
335 """Return the enum member matching `name`
336
337 We use __getattr__ instead of descriptors or inserting into the enum
338 class' __dict__ in order to support `name` and `value` being both
339 properties for enum members (which live in the class' __dict__) and
340 enum members themselves.
341
342 """
343 if _is_dunder(name):
344 raise AttributeError(name)
345 try:
Ethan Furman520ad572013-07-19 19:47:21 -0700346 return cls._member_map_[name]
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700347 except KeyError:
348 raise AttributeError(name) from None
349
350 def __getitem__(cls, name):
Ethan Furman520ad572013-07-19 19:47:21 -0700351 return cls._member_map_[name]
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700352
353 def __iter__(cls):
Ethan Furman520ad572013-07-19 19:47:21 -0700354 return (cls._member_map_[name] for name in cls._member_names_)
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700355
356 def __len__(cls):
Ethan Furman520ad572013-07-19 19:47:21 -0700357 return len(cls._member_names_)
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700358
Ethan Furman2131a4a2013-09-14 18:11:24 -0700359 @property
360 def __members__(cls):
361 """Returns a mapping of member name->value.
362
363 This mapping lists all enum members, including aliases. Note that this
364 is a read-only view of the internal mapping.
365
366 """
367 return MappingProxyType(cls._member_map_)
368
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700369 def __repr__(cls):
370 return "<enum %r>" % cls.__name__
371
Ethan Furman2131a4a2013-09-14 18:11:24 -0700372 def __reversed__(cls):
373 return (cls._member_map_[name] for name in reversed(cls._member_names_))
374
Ethan Furmanf203f2d2013-09-06 07:16:48 -0700375 def __setattr__(cls, name, value):
376 """Block attempts to reassign Enum members.
377
378 A simple assignment to the class namespace only changes one of the
379 several possible ways to get an Enum member from the Enum class,
380 resulting in an inconsistent Enumeration.
381
382 """
383 member_map = cls.__dict__.get('_member_map_', {})
384 if name in member_map:
385 raise AttributeError('Cannot reassign members.')
386 super().__setattr__(name, value)
387
Miss Islington (bot)3bcca482018-04-16 13:49:34 -0700388 def _create_(cls, class_name, names, *, module=None, qualname=None, type=None, start=1):
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700389 """Convenience method to create a new Enum class.
390
391 `names` can be:
392
393 * A string containing member names, separated either with spaces or
Ethan Furmand9925a12014-09-16 20:35:55 -0700394 commas. Values are incremented by 1 from `start`.
395 * An iterable of member names. Values are incremented by 1 from `start`.
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700396 * An iterable of (member name, value) pairs.
Ethan Furmand9925a12014-09-16 20:35:55 -0700397 * A mapping of member name -> value pairs.
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700398
399 """
400 metacls = cls.__class__
401 bases = (cls, ) if type is None else (type, cls)
Ethan Furmanee47e5c2016-08-31 00:12:15 -0700402 _, first_enum = cls._get_mixins_(bases)
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700403 classdict = metacls.__prepare__(class_name, bases)
404
405 # special processing needed for names?
406 if isinstance(names, str):
407 names = names.replace(',', ' ').split()
Dong-hee Nadcc8ce42017-06-22 01:52:32 +0900408 if isinstance(names, (tuple, list)) and names and isinstance(names[0], str):
Ethan Furmanee47e5c2016-08-31 00:12:15 -0700409 original_names, names = names, []
Ethan Furmanc16595e2016-09-10 23:36:59 -0700410 last_values = []
Ethan Furmanee47e5c2016-08-31 00:12:15 -0700411 for count, name in enumerate(original_names):
Ethan Furmanc16595e2016-09-10 23:36:59 -0700412 value = first_enum._generate_next_value_(name, start, count, last_values[:])
413 last_values.append(value)
414 names.append((name, value))
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700415
416 # Here, names is either an iterable of (name, value) or a mapping.
417 for item in names:
418 if isinstance(item, str):
419 member_name, member_value = item, names[item]
420 else:
421 member_name, member_value = item
422 classdict[member_name] = member_value
423 enum_class = metacls.__new__(metacls, class_name, bases, classdict)
424
425 # TODO: replace the frame hack if a blessed way to know the calling
426 # module is ever developed
427 if module is None:
428 try:
429 module = sys._getframe(2).f_globals['__name__']
430 except (AttributeError, ValueError) as exc:
431 pass
432 if module is None:
433 _make_class_unpicklable(enum_class)
434 else:
435 enum_class.__module__ = module
Ethan Furmanca1b7942014-02-08 11:36:27 -0800436 if qualname is not None:
437 enum_class.__qualname__ = qualname
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700438
439 return enum_class
440
441 @staticmethod
442 def _get_mixins_(bases):
443 """Returns the type for creating enum members, and the first inherited
444 enum class.
445
446 bases: the tuple of bases that was given to __new__
447
448 """
449 if not bases:
450 return object, Enum
451
452 # double check that we are not subclassing a class with existing
453 # enumeration members; while we're at it, see if any other data
454 # type has been mixed in so we can use the correct __new__
455 member_type = first_enum = None
456 for base in bases:
457 if (base is not Enum and
458 issubclass(base, Enum) and
Ethan Furman520ad572013-07-19 19:47:21 -0700459 base._member_names_):
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700460 raise TypeError("Cannot extend enumerations")
461 # base is now the last base in bases
462 if not issubclass(base, Enum):
463 raise TypeError("new enumerations must be created as "
464 "`ClassName([mixin_type,] enum_type)`")
465
466 # get correct mix-in type (either mix-in type of Enum subclass, or
467 # first base if last base is Enum)
468 if not issubclass(bases[0], Enum):
469 member_type = bases[0] # first data type
470 first_enum = bases[-1] # enum type
471 else:
472 for base in bases[0].__mro__:
473 # most common: (IntEnum, int, Enum, object)
474 # possible: (<Enum 'AutoIntEnum'>, <Enum 'IntEnum'>,
475 # <class 'int'>, <Enum 'Enum'>,
476 # <class 'object'>)
477 if issubclass(base, Enum):
478 if first_enum is None:
479 first_enum = base
480 else:
481 if member_type is None:
482 member_type = base
483
484 return member_type, first_enum
485
486 @staticmethod
487 def _find_new_(classdict, member_type, first_enum):
488 """Returns the __new__ to be used for creating the enum members.
489
490 classdict: the class dictionary given to __new__
491 member_type: the data type whose __new__ will be used by default
492 first_enum: enumeration to check for an overriding __new__
493
494 """
495 # now find the correct __new__, checking to see of one was defined
496 # by the user; also check earlier enum classes in case a __new__ was
497 # saved as __new_member__
498 __new__ = classdict.get('__new__', None)
499
500 # should __new__ be saved as __new_member__ later?
501 save_new = __new__ is not None
502
503 if __new__ is None:
504 # check all possibles for __new_member__ before falling back to
505 # __new__
506 for method in ('__new_member__', '__new__'):
507 for possible in (member_type, first_enum):
508 target = getattr(possible, method, None)
509 if target not in {
510 None,
511 None.__new__,
512 object.__new__,
513 Enum.__new__,
514 }:
515 __new__ = target
516 break
517 if __new__ is not None:
518 break
519 else:
520 __new__ = object.__new__
521
522 # if a non-object.__new__ is used then whatever value/tuple was
523 # assigned to the enum member name will be passed to __new__ and to the
524 # new enum member's __init__
525 if __new__ is object.__new__:
526 use_args = False
527 else:
528 use_args = True
529
530 return __new__, save_new, use_args
531
532
533class Enum(metaclass=EnumMeta):
534 """Generic enumeration.
535
536 Derive from this class to define new enumerations.
537
538 """
539 def __new__(cls, value):
540 # all enum instances are actually created during class construction
541 # without calling this method; this method is called by the metaclass'
542 # __call__ (i.e. Color(3) ), and by pickle
543 if type(value) is cls:
Ethan Furman23bb6f42016-11-21 09:22:05 -0800544 # For lookups like Color(Color.RED)
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700545 return value
546 # by-value search for a matching enum member
547 # see if it's in the reverse mapping (for hashable values)
Ethan Furman2aa27322013-07-19 19:35:56 -0700548 try:
Ethan Furman520ad572013-07-19 19:47:21 -0700549 if value in cls._value2member_map_:
550 return cls._value2member_map_[value]
Ethan Furman2aa27322013-07-19 19:35:56 -0700551 except TypeError:
552 # not there, now do long search -- O(n) behavior
Ethan Furman520ad572013-07-19 19:47:21 -0700553 for member in cls._member_map_.values():
Ethan Furman0081f232014-09-16 17:31:23 -0700554 if member._value_ == value:
Ethan Furman2aa27322013-07-19 19:35:56 -0700555 return member
Ethan Furmanee47e5c2016-08-31 00:12:15 -0700556 # still not found -- try _missing_ hook
557 return cls._missing_(value)
558
Ethan Furmanc16595e2016-09-10 23:36:59 -0700559 def _generate_next_value_(name, start, count, last_values):
560 for last_value in reversed(last_values):
561 try:
562 return last_value + 1
563 except TypeError:
564 pass
565 else:
Ethan Furmanee47e5c2016-08-31 00:12:15 -0700566 return start
Ethan Furmanc16595e2016-09-10 23:36:59 -0700567
Ethan Furmanee47e5c2016-08-31 00:12:15 -0700568 @classmethod
569 def _missing_(cls, value):
Ethan Furman0081f232014-09-16 17:31:23 -0700570 raise ValueError("%r is not a valid %s" % (value, cls.__name__))
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700571
572 def __repr__(self):
573 return "<%s.%s: %r>" % (
Ethan Furman520ad572013-07-19 19:47:21 -0700574 self.__class__.__name__, self._name_, self._value_)
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700575
576 def __str__(self):
Ethan Furman520ad572013-07-19 19:47:21 -0700577 return "%s.%s" % (self.__class__.__name__, self._name_)
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700578
Ethan Furman388a3922013-08-12 06:51:41 -0700579 def __dir__(self):
Ethan Furman0ae550b2014-10-14 08:58:32 -0700580 added_behavior = [
581 m
582 for cls in self.__class__.mro()
583 for m in cls.__dict__
Ethan Furman354ecf12015-03-11 08:43:12 -0700584 if m[0] != '_' and m not in self._member_map_
Ethan Furman0ae550b2014-10-14 08:58:32 -0700585 ]
Ethan Furmanec5f8eb2014-10-21 13:40:35 -0700586 return (['__class__', '__doc__', '__module__'] + added_behavior)
Ethan Furman388a3922013-08-12 06:51:41 -0700587
Ethan Furmanec15a822013-08-31 19:17:41 -0700588 def __format__(self, format_spec):
589 # mixed-in Enums should use the mixed-in type's __format__, otherwise
590 # we can get strange results with the Enum name showing up instead of
591 # the value
592
593 # pure Enum branch
594 if self._member_type_ is object:
595 cls = str
596 val = str(self)
597 # mix-in branch
598 else:
599 cls = self._member_type_
Ethan Furman0081f232014-09-16 17:31:23 -0700600 val = self._value_
Ethan Furmanec15a822013-08-31 19:17:41 -0700601 return cls.__format__(val, format_spec)
602
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700603 def __hash__(self):
Ethan Furman520ad572013-07-19 19:47:21 -0700604 return hash(self._name_)
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700605
Ethan Furmanca1b7942014-02-08 11:36:27 -0800606 def __reduce_ex__(self, proto):
Ethan Furmandc870522014-02-18 12:37:12 -0800607 return self.__class__, (self._value_, )
Ethan Furmanca1b7942014-02-08 11:36:27 -0800608
Ethan Furman33918c12013-09-27 23:02:02 -0700609 # DynamicClassAttribute is used to provide access to the `name` and
610 # `value` properties of enum members while keeping some measure of
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700611 # protection from modification, while still allowing for an enumeration
612 # to have members named `name` and `value`. This works because enumeration
613 # members are not set directly on the enum class -- __getattr__ is
614 # used to look them up.
615
Ethan Furmane03ea372013-09-25 07:14:41 -0700616 @DynamicClassAttribute
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700617 def name(self):
Ethan Furmanc850f342013-09-15 16:59:35 -0700618 """The name of the Enum member."""
Ethan Furman520ad572013-07-19 19:47:21 -0700619 return self._name_
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700620
Ethan Furmane03ea372013-09-25 07:14:41 -0700621 @DynamicClassAttribute
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700622 def value(self):
Ethan Furmanc850f342013-09-15 16:59:35 -0700623 """The value of the Enum member."""
Ethan Furman520ad572013-07-19 19:47:21 -0700624 return self._value_
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700625
Ethan Furman24e837f2015-03-18 17:27:57 -0700626 @classmethod
627 def _convert(cls, name, module, filter, source=None):
628 """
629 Create a new Enum subclass that replaces a collection of global constants
630 """
631 # convert all constants from source (or module) that pass filter() to
632 # a new Enum called name, and export the enum and its members back to
633 # module;
634 # also, replace the __reduce_ex__ method so unpickling works in
635 # previous Python versions
636 module_globals = vars(sys.modules[module])
637 if source:
638 source = vars(source)
639 else:
640 source = module_globals
Gregory P. Smith ext:(%20%5BGoogle%20Inc.%5D)6f20bd62016-06-03 19:14:52 +0000641 # We use an OrderedDict of sorted source keys so that the
642 # _value2member_map is populated in the same order every time
643 # for a consistent reverse mapping of number to name when there
644 # are multiple names for the same number rather than varying
645 # between runs due to hash randomization of the module dictionary.
Ethan Furman06339e72016-09-11 13:25:26 -0700646 members = [
647 (name, source[name])
648 for name in source.keys()
649 if filter(name)]
650 try:
651 # sort by value
652 members.sort(key=lambda t: (t[1], t[0]))
653 except TypeError:
654 # unless some values aren't comparable, in which case sort by name
655 members.sort(key=lambda t: t[0])
Ethan Furman24e837f2015-03-18 17:27:57 -0700656 cls = cls(name, members, module=module)
657 cls.__reduce_ex__ = _reduce_ex_by_name
658 module_globals.update(cls.__members__)
659 module_globals[name] = cls
660 return cls
661
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700662
663class IntEnum(int, Enum):
664 """Enum where members are also (and must be) ints"""
Ethan Furmanf24bb352013-07-18 17:05:39 -0700665
666
Ethan Furman24e837f2015-03-18 17:27:57 -0700667def _reduce_ex_by_name(self, proto):
668 return self.name
669
Ethan Furman65a5a472016-09-01 23:55:19 -0700670class Flag(Enum):
Ethan Furmanee47e5c2016-08-31 00:12:15 -0700671 """Support for flags"""
Ethan Furmanc16595e2016-09-10 23:36:59 -0700672
673 def _generate_next_value_(name, start, count, last_values):
Ethan Furmanee47e5c2016-08-31 00:12:15 -0700674 """
675 Generate the next value when not given.
676
677 name: the name of the member
678 start: the initital start value or None
679 count: the number of existing members
680 last_value: the last value assigned or None
681 """
682 if not count:
683 return start if start is not None else 1
Ethan Furmanc16595e2016-09-10 23:36:59 -0700684 for last_value in reversed(last_values):
685 try:
686 high_bit = _high_bit(last_value)
687 break
Ethan Furman3515dcc2016-09-18 13:15:41 -0700688 except Exception:
Ethan Furmanc16595e2016-09-10 23:36:59 -0700689 raise TypeError('Invalid Flag value: %r' % last_value) from None
Ethan Furmanee47e5c2016-08-31 00:12:15 -0700690 return 2 ** (high_bit+1)
691
692 @classmethod
693 def _missing_(cls, value):
694 original_value = value
695 if value < 0:
696 value = ~value
697 possible_member = cls._create_pseudo_member_(value)
Ethan Furmanee47e5c2016-08-31 00:12:15 -0700698 if original_value < 0:
699 possible_member = ~possible_member
700 return possible_member
701
702 @classmethod
703 def _create_pseudo_member_(cls, value):
Ethan Furman3515dcc2016-09-18 13:15:41 -0700704 """
705 Create a composite member iff value contains only members.
706 """
Ethan Furmanee47e5c2016-08-31 00:12:15 -0700707 pseudo_member = cls._value2member_map_.get(value, None)
708 if pseudo_member is None:
Ethan Furman3515dcc2016-09-18 13:15:41 -0700709 # verify all bits are accounted for
710 _, extra_flags = _decompose(cls, value)
711 if extra_flags:
712 raise ValueError("%r is not a valid %s" % (value, cls.__name__))
713 # construct a singleton enum pseudo-member
Ethan Furmanee47e5c2016-08-31 00:12:15 -0700714 pseudo_member = object.__new__(cls)
715 pseudo_member._name_ = None
716 pseudo_member._value_ = value
Ethan Furman28cf6632017-01-24 12:12:06 -0800717 # use setdefault in case another thread already created a composite
718 # with this value
719 pseudo_member = cls._value2member_map_.setdefault(value, pseudo_member)
Ethan Furmanee47e5c2016-08-31 00:12:15 -0700720 return pseudo_member
721
Ethan Furmanee47e5c2016-08-31 00:12:15 -0700722 def __contains__(self, other):
723 if not isinstance(other, self.__class__):
Ethan Furman37151762018-04-11 18:56:25 -0700724 import warnings
725 warnings.warn(
726 "using non-Flags in containment checks will raise "
727 "TypeError in Python 3.8",
728 DeprecationWarning, 2)
729 return False
Ethan Furmanee47e5c2016-08-31 00:12:15 -0700730 return other._value_ & self._value_ == other._value_
731
Ethan Furmanee47e5c2016-08-31 00:12:15 -0700732 def __repr__(self):
733 cls = self.__class__
734 if self._name_ is not None:
735 return '<%s.%s: %r>' % (cls.__name__, self._name_, self._value_)
Ethan Furman3515dcc2016-09-18 13:15:41 -0700736 members, uncovered = _decompose(cls, self._value_)
Ethan Furman27682d22016-09-04 11:39:01 -0700737 return '<%s.%s: %r>' % (
738 cls.__name__,
739 '|'.join([str(m._name_ or m._value_) for m in members]),
740 self._value_,
741 )
Ethan Furmanee47e5c2016-08-31 00:12:15 -0700742
743 def __str__(self):
744 cls = self.__class__
745 if self._name_ is not None:
746 return '%s.%s' % (cls.__name__, self._name_)
Ethan Furman3515dcc2016-09-18 13:15:41 -0700747 members, uncovered = _decompose(cls, self._value_)
Ethan Furmanee47e5c2016-08-31 00:12:15 -0700748 if len(members) == 1 and members[0]._name_ is None:
749 return '%s.%r' % (cls.__name__, members[0]._value_)
750 else:
751 return '%s.%s' % (
752 cls.__name__,
753 '|'.join([str(m._name_ or m._value_) for m in members]),
754 )
755
Ethan Furman25d94bb2016-09-02 16:32:32 -0700756 def __bool__(self):
757 return bool(self._value_)
758
Ethan Furmanee47e5c2016-08-31 00:12:15 -0700759 def __or__(self, other):
760 if not isinstance(other, self.__class__):
761 return NotImplemented
762 return self.__class__(self._value_ | other._value_)
763
764 def __and__(self, other):
765 if not isinstance(other, self.__class__):
766 return NotImplemented
767 return self.__class__(self._value_ & other._value_)
768
769 def __xor__(self, other):
770 if not isinstance(other, self.__class__):
771 return NotImplemented
772 return self.__class__(self._value_ ^ other._value_)
773
774 def __invert__(self):
Ethan Furman3515dcc2016-09-18 13:15:41 -0700775 members, uncovered = _decompose(self.__class__, self._value_)
Serhiy Storchaka81108372017-09-26 00:55:55 +0300776 inverted = self.__class__(0)
777 for m in self.__class__:
778 if m not in members and not (m._value_ & self._value_):
779 inverted = inverted | m
Ethan Furmanee47e5c2016-08-31 00:12:15 -0700780 return self.__class__(inverted)
781
782
Ethan Furman65a5a472016-09-01 23:55:19 -0700783class IntFlag(int, Flag):
Ethan Furmanee47e5c2016-08-31 00:12:15 -0700784 """Support for integer-based Flags"""
785
786 @classmethod
Ethan Furman3515dcc2016-09-18 13:15:41 -0700787 def _missing_(cls, value):
788 if not isinstance(value, int):
789 raise ValueError("%r is not a valid %s" % (value, cls.__name__))
790 new_member = cls._create_pseudo_member_(value)
791 return new_member
792
793 @classmethod
Ethan Furmanee47e5c2016-08-31 00:12:15 -0700794 def _create_pseudo_member_(cls, value):
795 pseudo_member = cls._value2member_map_.get(value, None)
796 if pseudo_member is None:
Ethan Furman3515dcc2016-09-18 13:15:41 -0700797 need_to_create = [value]
798 # get unaccounted for bits
799 _, extra_flags = _decompose(cls, value)
800 # timer = 10
801 while extra_flags:
802 # timer -= 1
803 bit = _high_bit(extra_flags)
804 flag_value = 2 ** bit
805 if (flag_value not in cls._value2member_map_ and
806 flag_value not in need_to_create
807 ):
808 need_to_create.append(flag_value)
809 if extra_flags == -flag_value:
810 extra_flags = 0
811 else:
812 extra_flags ^= flag_value
813 for value in reversed(need_to_create):
814 # construct singleton pseudo-members
815 pseudo_member = int.__new__(cls, value)
816 pseudo_member._name_ = None
817 pseudo_member._value_ = value
Ethan Furman28cf6632017-01-24 12:12:06 -0800818 # use setdefault in case another thread already created a composite
819 # with this value
820 pseudo_member = cls._value2member_map_.setdefault(value, pseudo_member)
Ethan Furmanee47e5c2016-08-31 00:12:15 -0700821 return pseudo_member
822
Ethan Furmanee47e5c2016-08-31 00:12:15 -0700823 def __or__(self, other):
824 if not isinstance(other, (self.__class__, int)):
825 return NotImplemented
Ethan Furman3515dcc2016-09-18 13:15:41 -0700826 result = self.__class__(self._value_ | self.__class__(other)._value_)
827 return result
Ethan Furmanee47e5c2016-08-31 00:12:15 -0700828
829 def __and__(self, other):
830 if not isinstance(other, (self.__class__, int)):
831 return NotImplemented
832 return self.__class__(self._value_ & self.__class__(other)._value_)
833
834 def __xor__(self, other):
835 if not isinstance(other, (self.__class__, int)):
836 return NotImplemented
837 return self.__class__(self._value_ ^ self.__class__(other)._value_)
838
839 __ror__ = __or__
840 __rand__ = __and__
841 __rxor__ = __xor__
842
843 def __invert__(self):
Ethan Furman3515dcc2016-09-18 13:15:41 -0700844 result = self.__class__(~self._value_)
845 return result
Ethan Furmanee47e5c2016-08-31 00:12:15 -0700846
847
848def _high_bit(value):
Ethan Furman04439532016-09-02 15:50:21 -0700849 """returns index of highest bit, or -1 if value is zero or negative"""
Ethan Furman3515dcc2016-09-18 13:15:41 -0700850 return value.bit_length() - 1
Ethan Furmanee47e5c2016-08-31 00:12:15 -0700851
Ethan Furmanf24bb352013-07-18 17:05:39 -0700852def unique(enumeration):
853 """Class decorator for enumerations ensuring unique member values."""
854 duplicates = []
855 for name, member in enumeration.__members__.items():
856 if name != member.name:
857 duplicates.append((name, member.name))
858 if duplicates:
859 alias_details = ', '.join(
860 ["%s -> %s" % (alias, name) for (alias, name) in duplicates])
861 raise ValueError('duplicate values found in %r: %s' %
862 (enumeration, alias_details))
863 return enumeration
Ethan Furman3515dcc2016-09-18 13:15:41 -0700864
865def _decompose(flag, value):
866 """Extract all members from the value."""
867 # _decompose is only called if the value is not named
868 not_covered = value
869 negative = value < 0
Ethan Furman28cf6632017-01-24 12:12:06 -0800870 # issue29167: wrap accesses to _value2member_map_ in a list to avoid race
Ville Skyttä49b27342017-08-03 09:00:59 +0300871 # conditions between iterating over it and having more pseudo-
Ethan Furman28cf6632017-01-24 12:12:06 -0800872 # members added to it
Ethan Furman3515dcc2016-09-18 13:15:41 -0700873 if negative:
874 # only check for named flags
875 flags_to_check = [
876 (m, v)
Ethan Furman28cf6632017-01-24 12:12:06 -0800877 for v, m in list(flag._value2member_map_.items())
Ethan Furman3515dcc2016-09-18 13:15:41 -0700878 if m.name is not None
879 ]
880 else:
881 # check for named flags and powers-of-two 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 or _power_of_two(v)
886 ]
887 members = []
888 for member, member_value in flags_to_check:
889 if member_value and member_value & value == member_value:
890 members.append(member)
891 not_covered &= ~member_value
892 if not members and value in flag._value2member_map_:
893 members.append(flag._value2member_map_[value])
894 members.sort(key=lambda m: m._value_, reverse=True)
895 if len(members) > 1 and members[0].value == value:
896 # we have the breakdown, don't need the value member itself
897 members.pop(0)
898 return members, not_covered
899
900def _power_of_two(value):
901 if value < 1:
902 return False
903 return value == 2 ** _high_bit(value)