blob: 576de031805dbc8022da4fc13ff6901cc469dbc2 [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
Ethan Furman354ecf12015-03-11 08:43:12 -0700174 # save attributes from super classes so we know if we can take
175 # the shortcut of storing members in the class dict
Ethan Furman3803ad42016-05-01 10:03:53 -0700176 base_attributes = {a for b in enum_class.mro() for a in b.__dict__}
Ethan Furman354ecf12015-03-11 08:43:12 -0700177
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700178 # Reverse value->name map for hashable values.
Ethan Furman520ad572013-07-19 19:47:21 -0700179 enum_class._value2member_map_ = {}
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700180
Ethan Furman2da95042014-03-03 12:42:52 -0800181 # If a custom type is mixed into the Enum, and it does not know how
182 # to pickle itself, pickle.dumps will succeed but pickle.loads will
183 # fail. Rather than have the error show up later and possibly far
184 # from the source, sabotage the pickle protocol for this class so
185 # that pickle.dumps also fails.
186 #
187 # However, if the new class implements its own __reduce_ex__, do not
188 # sabotage -- it's on them to make sure it works correctly. We use
189 # __reduce_ex__ instead of any of the others as it is preferred by
190 # pickle over __reduce__, and it handles all pickle protocols.
191 if '__reduce_ex__' not in classdict:
Ethan Furmandc870522014-02-18 12:37:12 -0800192 if member_type is not object:
193 methods = ('__getnewargs_ex__', '__getnewargs__',
194 '__reduce_ex__', '__reduce__')
Ethan Furman2da95042014-03-03 12:42:52 -0800195 if not any(m in member_type.__dict__ for m in methods):
Ethan Furmandc870522014-02-18 12:37:12 -0800196 _make_class_unpicklable(enum_class)
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700197
198 # instantiate them, checking for duplicates as we go
199 # we instantiate first instead of checking for duplicates first in case
200 # a custom __new__ is doing something funky with the values -- such as
201 # auto-numbering ;)
202 for member_name in classdict._member_names:
Ethan Furmanee47e5c2016-08-31 00:12:15 -0700203 value = enum_members[member_name]
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700204 if not isinstance(value, tuple):
205 args = (value, )
206 else:
207 args = value
208 if member_type is tuple: # special case for tuple enums
209 args = (args, ) # wrap it one more time
210 if not use_args:
211 enum_member = __new__(enum_class)
Ethan Furmanb41803e2013-07-25 13:50:45 -0700212 if not hasattr(enum_member, '_value_'):
213 enum_member._value_ = value
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700214 else:
215 enum_member = __new__(enum_class, *args)
Ethan Furmanb41803e2013-07-25 13:50:45 -0700216 if not hasattr(enum_member, '_value_'):
Ethan Furmanee47e5c2016-08-31 00:12:15 -0700217 if member_type is object:
218 enum_member._value_ = value
219 else:
220 enum_member._value_ = member_type(*args)
Ethan Furman520ad572013-07-19 19:47:21 -0700221 value = enum_member._value_
Ethan Furman520ad572013-07-19 19:47:21 -0700222 enum_member._name_ = member_name
Ethan Furmanc850f342013-09-15 16:59:35 -0700223 enum_member.__objclass__ = enum_class
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700224 enum_member.__init__(*args)
225 # If another member with the same value was already defined, the
226 # new member becomes an alias to the existing one.
Ethan Furman520ad572013-07-19 19:47:21 -0700227 for name, canonical_member in enum_class._member_map_.items():
Ethan Furman0081f232014-09-16 17:31:23 -0700228 if canonical_member._value_ == enum_member._value_:
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700229 enum_member = canonical_member
230 break
231 else:
232 # Aliases don't appear in member names (only in __members__).
Ethan Furman520ad572013-07-19 19:47:21 -0700233 enum_class._member_names_.append(member_name)
Ethan Furman354ecf12015-03-11 08:43:12 -0700234 # performance boost for any member that would not shadow
235 # a DynamicClassAttribute
236 if member_name not in base_attributes:
237 setattr(enum_class, member_name, enum_member)
238 # now add to _member_map_
Ethan Furman520ad572013-07-19 19:47:21 -0700239 enum_class._member_map_[member_name] = enum_member
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700240 try:
241 # This may fail if value is not hashable. We can't add the value
242 # to the map, and by-value lookups for this value will be
243 # linear.
Ethan Furman520ad572013-07-19 19:47:21 -0700244 enum_class._value2member_map_[value] = enum_member
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700245 except TypeError:
246 pass
247
248 # double check that repr and friends are not the mixin's or various
249 # things break (such as pickle)
Ethan Furmandc870522014-02-18 12:37:12 -0800250 for name in ('__repr__', '__str__', '__format__', '__reduce_ex__'):
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700251 class_method = getattr(enum_class, name)
252 obj_method = getattr(member_type, name, None)
253 enum_method = getattr(first_enum, name, None)
254 if obj_method is not None and obj_method is class_method:
255 setattr(enum_class, name, enum_method)
256
257 # replace any other __new__ with our own (as long as Enum is not None,
258 # anyway) -- again, this is to support pickle
259 if Enum is not None:
260 # if the user defined their own __new__, save it before it gets
261 # clobbered in case they subclass later
262 if save_new:
263 enum_class.__new_member__ = __new__
264 enum_class.__new__ = Enum.__new__
Ethan Furmane8e61272016-08-20 07:19:31 -0700265
266 # py3 support for definition order (helps keep py2/py3 code in sync)
267 if _order_ is not None:
268 if isinstance(_order_, str):
269 _order_ = _order_.replace(',', ' ').split()
270 if _order_ != enum_class._member_names_:
271 raise TypeError('member order does not match _order_')
272
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700273 return enum_class
274
Ethan Furman5de67b12016-04-13 23:52:09 -0700275 def __bool__(self):
276 """
277 classes/types should always be True.
278 """
279 return True
280
Ethan Furmand9925a12014-09-16 20:35:55 -0700281 def __call__(cls, value, names=None, *, module=None, qualname=None, type=None, start=1):
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700282 """Either returns an existing member, or creates a new enum class.
283
284 This method is used both when an enum class is given a value to match
285 to an enumeration member (i.e. Color(3)) and for the functional API
Ethan Furman23bb6f42016-11-21 09:22:05 -0800286 (i.e. Color = Enum('Color', names='RED GREEN BLUE')).
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700287
Ethan Furman2da95042014-03-03 12:42:52 -0800288 When used for the functional API:
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700289
Ethan Furman2da95042014-03-03 12:42:52 -0800290 `value` will be the name of the new class.
291
292 `names` should be either a string of white-space/comma delimited names
Ethan Furmand9925a12014-09-16 20:35:55 -0700293 (values will start at `start`), or an iterator/mapping of name, value pairs.
Ethan Furman2da95042014-03-03 12:42:52 -0800294
295 `module` should be set to the module this class is being created in;
296 if it is not set, an attempt to find that module will be made, but if
297 it fails the class will not be picklable.
298
299 `qualname` should be set to the actual location this class can be found
300 at in its module; by default it is set to the global scope. If this is
301 not correct, unpickling will fail in some circumstances.
302
303 `type`, if set, will be mixed in as the first base class.
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700304
305 """
306 if names is None: # simple value lookup
307 return cls.__new__(cls, value)
308 # otherwise, functional API: we're creating a new Enum type
Ethan Furmand9925a12014-09-16 20:35:55 -0700309 return cls._create_(value, names, module=module, qualname=qualname, type=type, start=start)
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700310
311 def __contains__(cls, member):
Ethan Furman37151762018-04-11 18:56:25 -0700312 if not isinstance(member, Enum):
313 import warnings
314 warnings.warn(
315 "using non-Enums in containment checks will raise "
316 "TypeError in Python 3.8",
317 DeprecationWarning, 2)
Ethan Furman0081f232014-09-16 17:31:23 -0700318 return isinstance(member, cls) and member._name_ in cls._member_map_
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700319
Ethan Furman64a99722013-09-22 16:18:19 -0700320 def __delattr__(cls, attr):
321 # nicer error message when someone tries to delete an attribute
322 # (see issue19025).
323 if attr in cls._member_map_:
324 raise AttributeError(
325 "%s: cannot delete Enum member." % cls.__name__)
326 super().__delattr__(attr)
327
Ethan Furman388a3922013-08-12 06:51:41 -0700328 def __dir__(self):
Ethan Furman64a99722013-09-22 16:18:19 -0700329 return (['__class__', '__doc__', '__members__', '__module__'] +
330 self._member_names_)
Ethan Furman388a3922013-08-12 06:51:41 -0700331
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700332 def __getattr__(cls, name):
333 """Return the enum member matching `name`
334
335 We use __getattr__ instead of descriptors or inserting into the enum
336 class' __dict__ in order to support `name` and `value` being both
337 properties for enum members (which live in the class' __dict__) and
338 enum members themselves.
339
340 """
341 if _is_dunder(name):
342 raise AttributeError(name)
343 try:
Ethan Furman520ad572013-07-19 19:47:21 -0700344 return cls._member_map_[name]
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700345 except KeyError:
346 raise AttributeError(name) from None
347
348 def __getitem__(cls, name):
Ethan Furman520ad572013-07-19 19:47:21 -0700349 return cls._member_map_[name]
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700350
351 def __iter__(cls):
Ethan Furman520ad572013-07-19 19:47:21 -0700352 return (cls._member_map_[name] for name in cls._member_names_)
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700353
354 def __len__(cls):
Ethan Furman520ad572013-07-19 19:47:21 -0700355 return len(cls._member_names_)
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700356
Ethan Furman2131a4a2013-09-14 18:11:24 -0700357 @property
358 def __members__(cls):
359 """Returns a mapping of member name->value.
360
361 This mapping lists all enum members, including aliases. Note that this
362 is a read-only view of the internal mapping.
363
364 """
365 return MappingProxyType(cls._member_map_)
366
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700367 def __repr__(cls):
368 return "<enum %r>" % cls.__name__
369
Ethan Furman2131a4a2013-09-14 18:11:24 -0700370 def __reversed__(cls):
371 return (cls._member_map_[name] for name in reversed(cls._member_names_))
372
Ethan Furmanf203f2d2013-09-06 07:16:48 -0700373 def __setattr__(cls, name, value):
374 """Block attempts to reassign Enum members.
375
376 A simple assignment to the class namespace only changes one of the
377 several possible ways to get an Enum member from the Enum class,
378 resulting in an inconsistent Enumeration.
379
380 """
381 member_map = cls.__dict__.get('_member_map_', {})
382 if name in member_map:
383 raise AttributeError('Cannot reassign members.')
384 super().__setattr__(name, value)
385
Miss Islington (bot)3bcca482018-04-16 13:49:34 -0700386 def _create_(cls, class_name, names, *, module=None, qualname=None, type=None, start=1):
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700387 """Convenience method to create a new Enum class.
388
389 `names` can be:
390
391 * A string containing member names, separated either with spaces or
Ethan Furmand9925a12014-09-16 20:35:55 -0700392 commas. Values are incremented by 1 from `start`.
393 * An iterable of member names. Values are incremented by 1 from `start`.
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700394 * An iterable of (member name, value) pairs.
Ethan Furmand9925a12014-09-16 20:35:55 -0700395 * A mapping of member name -> value pairs.
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700396
397 """
398 metacls = cls.__class__
399 bases = (cls, ) if type is None else (type, cls)
Ethan Furmanee47e5c2016-08-31 00:12:15 -0700400 _, first_enum = cls._get_mixins_(bases)
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700401 classdict = metacls.__prepare__(class_name, bases)
402
403 # special processing needed for names?
404 if isinstance(names, str):
405 names = names.replace(',', ' ').split()
Dong-hee Nadcc8ce42017-06-22 01:52:32 +0900406 if isinstance(names, (tuple, list)) and names and isinstance(names[0], str):
Ethan Furmanee47e5c2016-08-31 00:12:15 -0700407 original_names, names = names, []
Ethan Furmanc16595e2016-09-10 23:36:59 -0700408 last_values = []
Ethan Furmanee47e5c2016-08-31 00:12:15 -0700409 for count, name in enumerate(original_names):
Ethan Furmanc16595e2016-09-10 23:36:59 -0700410 value = first_enum._generate_next_value_(name, start, count, last_values[:])
411 last_values.append(value)
412 names.append((name, value))
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700413
414 # Here, names is either an iterable of (name, value) or a mapping.
415 for item in names:
416 if isinstance(item, str):
417 member_name, member_value = item, names[item]
418 else:
419 member_name, member_value = item
420 classdict[member_name] = member_value
421 enum_class = metacls.__new__(metacls, class_name, bases, classdict)
422
423 # TODO: replace the frame hack if a blessed way to know the calling
424 # module is ever developed
425 if module is None:
426 try:
427 module = sys._getframe(2).f_globals['__name__']
428 except (AttributeError, ValueError) as exc:
429 pass
430 if module is None:
431 _make_class_unpicklable(enum_class)
432 else:
433 enum_class.__module__ = module
Ethan Furmanca1b7942014-02-08 11:36:27 -0800434 if qualname is not None:
435 enum_class.__qualname__ = qualname
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700436
437 return enum_class
438
439 @staticmethod
440 def _get_mixins_(bases):
441 """Returns the type for creating enum members, and the first inherited
442 enum class.
443
444 bases: the tuple of bases that was given to __new__
445
446 """
447 if not bases:
448 return object, Enum
449
450 # double check that we are not subclassing a class with existing
451 # enumeration members; while we're at it, see if any other data
452 # type has been mixed in so we can use the correct __new__
453 member_type = first_enum = None
454 for base in bases:
455 if (base is not Enum and
456 issubclass(base, Enum) and
Ethan Furman520ad572013-07-19 19:47:21 -0700457 base._member_names_):
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700458 raise TypeError("Cannot extend enumerations")
459 # base is now the last base in bases
460 if not issubclass(base, Enum):
461 raise TypeError("new enumerations must be created as "
462 "`ClassName([mixin_type,] enum_type)`")
463
464 # get correct mix-in type (either mix-in type of Enum subclass, or
465 # first base if last base is Enum)
466 if not issubclass(bases[0], Enum):
467 member_type = bases[0] # first data type
468 first_enum = bases[-1] # enum type
469 else:
470 for base in bases[0].__mro__:
471 # most common: (IntEnum, int, Enum, object)
472 # possible: (<Enum 'AutoIntEnum'>, <Enum 'IntEnum'>,
473 # <class 'int'>, <Enum 'Enum'>,
474 # <class 'object'>)
475 if issubclass(base, Enum):
476 if first_enum is None:
477 first_enum = base
478 else:
479 if member_type is None:
480 member_type = base
481
482 return member_type, first_enum
483
484 @staticmethod
485 def _find_new_(classdict, member_type, first_enum):
486 """Returns the __new__ to be used for creating the enum members.
487
488 classdict: the class dictionary given to __new__
489 member_type: the data type whose __new__ will be used by default
490 first_enum: enumeration to check for an overriding __new__
491
492 """
493 # now find the correct __new__, checking to see of one was defined
494 # by the user; also check earlier enum classes in case a __new__ was
495 # saved as __new_member__
496 __new__ = classdict.get('__new__', None)
497
498 # should __new__ be saved as __new_member__ later?
499 save_new = __new__ is not None
500
501 if __new__ is None:
502 # check all possibles for __new_member__ before falling back to
503 # __new__
504 for method in ('__new_member__', '__new__'):
505 for possible in (member_type, first_enum):
506 target = getattr(possible, method, None)
507 if target not in {
508 None,
509 None.__new__,
510 object.__new__,
511 Enum.__new__,
512 }:
513 __new__ = target
514 break
515 if __new__ is not None:
516 break
517 else:
518 __new__ = object.__new__
519
520 # if a non-object.__new__ is used then whatever value/tuple was
521 # assigned to the enum member name will be passed to __new__ and to the
522 # new enum member's __init__
523 if __new__ is object.__new__:
524 use_args = False
525 else:
526 use_args = True
527
528 return __new__, save_new, use_args
529
530
531class Enum(metaclass=EnumMeta):
532 """Generic enumeration.
533
534 Derive from this class to define new enumerations.
535
536 """
537 def __new__(cls, value):
538 # all enum instances are actually created during class construction
539 # without calling this method; this method is called by the metaclass'
540 # __call__ (i.e. Color(3) ), and by pickle
541 if type(value) is cls:
Ethan Furman23bb6f42016-11-21 09:22:05 -0800542 # For lookups like Color(Color.RED)
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700543 return value
544 # by-value search for a matching enum member
545 # see if it's in the reverse mapping (for hashable values)
Ethan Furman2aa27322013-07-19 19:35:56 -0700546 try:
Ethan Furman520ad572013-07-19 19:47:21 -0700547 if value in cls._value2member_map_:
548 return cls._value2member_map_[value]
Ethan Furman2aa27322013-07-19 19:35:56 -0700549 except TypeError:
550 # not there, now do long search -- O(n) behavior
Ethan Furman520ad572013-07-19 19:47:21 -0700551 for member in cls._member_map_.values():
Ethan Furman0081f232014-09-16 17:31:23 -0700552 if member._value_ == value:
Ethan Furman2aa27322013-07-19 19:35:56 -0700553 return member
Ethan Furmanee47e5c2016-08-31 00:12:15 -0700554 # still not found -- try _missing_ hook
555 return cls._missing_(value)
556
Ethan Furmanc16595e2016-09-10 23:36:59 -0700557 def _generate_next_value_(name, start, count, last_values):
558 for last_value in reversed(last_values):
559 try:
560 return last_value + 1
561 except TypeError:
562 pass
563 else:
Ethan Furmanee47e5c2016-08-31 00:12:15 -0700564 return start
Ethan Furmanc16595e2016-09-10 23:36:59 -0700565
Ethan Furmanee47e5c2016-08-31 00:12:15 -0700566 @classmethod
567 def _missing_(cls, value):
Ethan Furman0081f232014-09-16 17:31:23 -0700568 raise ValueError("%r is not a valid %s" % (value, cls.__name__))
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700569
570 def __repr__(self):
571 return "<%s.%s: %r>" % (
Ethan Furman520ad572013-07-19 19:47:21 -0700572 self.__class__.__name__, self._name_, self._value_)
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700573
574 def __str__(self):
Ethan Furman520ad572013-07-19 19:47:21 -0700575 return "%s.%s" % (self.__class__.__name__, self._name_)
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700576
Ethan Furman388a3922013-08-12 06:51:41 -0700577 def __dir__(self):
Ethan Furman0ae550b2014-10-14 08:58:32 -0700578 added_behavior = [
579 m
580 for cls in self.__class__.mro()
581 for m in cls.__dict__
Ethan Furman354ecf12015-03-11 08:43:12 -0700582 if m[0] != '_' and m not in self._member_map_
Ethan Furman0ae550b2014-10-14 08:58:32 -0700583 ]
Ethan Furmanec5f8eb2014-10-21 13:40:35 -0700584 return (['__class__', '__doc__', '__module__'] + added_behavior)
Ethan Furman388a3922013-08-12 06:51:41 -0700585
Ethan Furmanec15a822013-08-31 19:17:41 -0700586 def __format__(self, format_spec):
587 # mixed-in Enums should use the mixed-in type's __format__, otherwise
588 # we can get strange results with the Enum name showing up instead of
589 # the value
590
591 # pure Enum branch
592 if self._member_type_ is object:
593 cls = str
594 val = str(self)
595 # mix-in branch
596 else:
597 cls = self._member_type_
Ethan Furman0081f232014-09-16 17:31:23 -0700598 val = self._value_
Ethan Furmanec15a822013-08-31 19:17:41 -0700599 return cls.__format__(val, format_spec)
600
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700601 def __hash__(self):
Ethan Furman520ad572013-07-19 19:47:21 -0700602 return hash(self._name_)
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700603
Ethan Furmanca1b7942014-02-08 11:36:27 -0800604 def __reduce_ex__(self, proto):
Ethan Furmandc870522014-02-18 12:37:12 -0800605 return self.__class__, (self._value_, )
Ethan Furmanca1b7942014-02-08 11:36:27 -0800606
Ethan Furman33918c12013-09-27 23:02:02 -0700607 # DynamicClassAttribute is used to provide access to the `name` and
608 # `value` properties of enum members while keeping some measure of
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700609 # protection from modification, while still allowing for an enumeration
610 # to have members named `name` and `value`. This works because enumeration
611 # members are not set directly on the enum class -- __getattr__ is
612 # used to look them up.
613
Ethan Furmane03ea372013-09-25 07:14:41 -0700614 @DynamicClassAttribute
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700615 def name(self):
Ethan Furmanc850f342013-09-15 16:59:35 -0700616 """The name of the Enum member."""
Ethan Furman520ad572013-07-19 19:47:21 -0700617 return self._name_
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700618
Ethan Furmane03ea372013-09-25 07:14:41 -0700619 @DynamicClassAttribute
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700620 def value(self):
Ethan Furmanc850f342013-09-15 16:59:35 -0700621 """The value of the Enum member."""
Ethan Furman520ad572013-07-19 19:47:21 -0700622 return self._value_
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700623
Ethan Furman24e837f2015-03-18 17:27:57 -0700624 @classmethod
625 def _convert(cls, name, module, filter, source=None):
626 """
627 Create a new Enum subclass that replaces a collection of global constants
628 """
629 # convert all constants from source (or module) that pass filter() to
630 # a new Enum called name, and export the enum and its members back to
631 # module;
632 # also, replace the __reduce_ex__ method so unpickling works in
633 # previous Python versions
634 module_globals = vars(sys.modules[module])
635 if source:
636 source = vars(source)
637 else:
638 source = module_globals
Gregory P. Smith ext:(%20%5BGoogle%20Inc.%5D)6f20bd62016-06-03 19:14:52 +0000639 # We use an OrderedDict of sorted source keys so that the
640 # _value2member_map is populated in the same order every time
641 # for a consistent reverse mapping of number to name when there
642 # are multiple names for the same number rather than varying
643 # between runs due to hash randomization of the module dictionary.
Ethan Furman06339e72016-09-11 13:25:26 -0700644 members = [
645 (name, source[name])
646 for name in source.keys()
647 if filter(name)]
648 try:
649 # sort by value
650 members.sort(key=lambda t: (t[1], t[0]))
651 except TypeError:
652 # unless some values aren't comparable, in which case sort by name
653 members.sort(key=lambda t: t[0])
Ethan Furman24e837f2015-03-18 17:27:57 -0700654 cls = cls(name, members, module=module)
655 cls.__reduce_ex__ = _reduce_ex_by_name
656 module_globals.update(cls.__members__)
657 module_globals[name] = cls
658 return cls
659
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700660
661class IntEnum(int, Enum):
662 """Enum where members are also (and must be) ints"""
Ethan Furmanf24bb352013-07-18 17:05:39 -0700663
664
Ethan Furman24e837f2015-03-18 17:27:57 -0700665def _reduce_ex_by_name(self, proto):
666 return self.name
667
Ethan Furman65a5a472016-09-01 23:55:19 -0700668class Flag(Enum):
Ethan Furmanee47e5c2016-08-31 00:12:15 -0700669 """Support for flags"""
Ethan Furmanc16595e2016-09-10 23:36:59 -0700670
671 def _generate_next_value_(name, start, count, last_values):
Ethan Furmanee47e5c2016-08-31 00:12:15 -0700672 """
673 Generate the next value when not given.
674
675 name: the name of the member
676 start: the initital start value or None
677 count: the number of existing members
678 last_value: the last value assigned or None
679 """
680 if not count:
681 return start if start is not None else 1
Ethan Furmanc16595e2016-09-10 23:36:59 -0700682 for last_value in reversed(last_values):
683 try:
684 high_bit = _high_bit(last_value)
685 break
Ethan Furman3515dcc2016-09-18 13:15:41 -0700686 except Exception:
Ethan Furmanc16595e2016-09-10 23:36:59 -0700687 raise TypeError('Invalid Flag value: %r' % last_value) from None
Ethan Furmanee47e5c2016-08-31 00:12:15 -0700688 return 2 ** (high_bit+1)
689
690 @classmethod
691 def _missing_(cls, value):
692 original_value = value
693 if value < 0:
694 value = ~value
695 possible_member = cls._create_pseudo_member_(value)
Ethan Furmanee47e5c2016-08-31 00:12:15 -0700696 if original_value < 0:
697 possible_member = ~possible_member
698 return possible_member
699
700 @classmethod
701 def _create_pseudo_member_(cls, value):
Ethan Furman3515dcc2016-09-18 13:15:41 -0700702 """
703 Create a composite member iff value contains only members.
704 """
Ethan Furmanee47e5c2016-08-31 00:12:15 -0700705 pseudo_member = cls._value2member_map_.get(value, None)
706 if pseudo_member is None:
Ethan Furman3515dcc2016-09-18 13:15:41 -0700707 # verify all bits are accounted for
708 _, extra_flags = _decompose(cls, value)
709 if extra_flags:
710 raise ValueError("%r is not a valid %s" % (value, cls.__name__))
711 # construct a singleton enum pseudo-member
Ethan Furmanee47e5c2016-08-31 00:12:15 -0700712 pseudo_member = object.__new__(cls)
713 pseudo_member._name_ = None
714 pseudo_member._value_ = value
Ethan Furman28cf6632017-01-24 12:12:06 -0800715 # use setdefault in case another thread already created a composite
716 # with this value
717 pseudo_member = cls._value2member_map_.setdefault(value, pseudo_member)
Ethan Furmanee47e5c2016-08-31 00:12:15 -0700718 return pseudo_member
719
Ethan Furmanee47e5c2016-08-31 00:12:15 -0700720 def __contains__(self, other):
721 if not isinstance(other, self.__class__):
Ethan Furman37151762018-04-11 18:56:25 -0700722 import warnings
723 warnings.warn(
724 "using non-Flags in containment checks will raise "
725 "TypeError in Python 3.8",
726 DeprecationWarning, 2)
727 return False
Ethan Furmanee47e5c2016-08-31 00:12:15 -0700728 return other._value_ & self._value_ == other._value_
729
Ethan Furmanee47e5c2016-08-31 00:12:15 -0700730 def __repr__(self):
731 cls = self.__class__
732 if self._name_ is not None:
733 return '<%s.%s: %r>' % (cls.__name__, self._name_, self._value_)
Ethan Furman3515dcc2016-09-18 13:15:41 -0700734 members, uncovered = _decompose(cls, self._value_)
Ethan Furman27682d22016-09-04 11:39:01 -0700735 return '<%s.%s: %r>' % (
736 cls.__name__,
737 '|'.join([str(m._name_ or m._value_) for m in members]),
738 self._value_,
739 )
Ethan Furmanee47e5c2016-08-31 00:12:15 -0700740
741 def __str__(self):
742 cls = self.__class__
743 if self._name_ is not None:
744 return '%s.%s' % (cls.__name__, self._name_)
Ethan Furman3515dcc2016-09-18 13:15:41 -0700745 members, uncovered = _decompose(cls, self._value_)
Ethan Furmanee47e5c2016-08-31 00:12:15 -0700746 if len(members) == 1 and members[0]._name_ is None:
747 return '%s.%r' % (cls.__name__, members[0]._value_)
748 else:
749 return '%s.%s' % (
750 cls.__name__,
751 '|'.join([str(m._name_ or m._value_) for m in members]),
752 )
753
Ethan Furman25d94bb2016-09-02 16:32:32 -0700754 def __bool__(self):
755 return bool(self._value_)
756
Ethan Furmanee47e5c2016-08-31 00:12:15 -0700757 def __or__(self, other):
758 if not isinstance(other, self.__class__):
759 return NotImplemented
760 return self.__class__(self._value_ | other._value_)
761
762 def __and__(self, other):
763 if not isinstance(other, self.__class__):
764 return NotImplemented
765 return self.__class__(self._value_ & other._value_)
766
767 def __xor__(self, other):
768 if not isinstance(other, self.__class__):
769 return NotImplemented
770 return self.__class__(self._value_ ^ other._value_)
771
772 def __invert__(self):
Ethan Furman3515dcc2016-09-18 13:15:41 -0700773 members, uncovered = _decompose(self.__class__, self._value_)
Serhiy Storchaka81108372017-09-26 00:55:55 +0300774 inverted = self.__class__(0)
775 for m in self.__class__:
776 if m not in members and not (m._value_ & self._value_):
777 inverted = inverted | m
Ethan Furmanee47e5c2016-08-31 00:12:15 -0700778 return self.__class__(inverted)
779
780
Ethan Furman65a5a472016-09-01 23:55:19 -0700781class IntFlag(int, Flag):
Ethan Furmanee47e5c2016-08-31 00:12:15 -0700782 """Support for integer-based Flags"""
783
784 @classmethod
Ethan Furman3515dcc2016-09-18 13:15:41 -0700785 def _missing_(cls, value):
786 if not isinstance(value, int):
787 raise ValueError("%r is not a valid %s" % (value, cls.__name__))
788 new_member = cls._create_pseudo_member_(value)
789 return new_member
790
791 @classmethod
Ethan Furmanee47e5c2016-08-31 00:12:15 -0700792 def _create_pseudo_member_(cls, value):
793 pseudo_member = cls._value2member_map_.get(value, None)
794 if pseudo_member is None:
Ethan Furman3515dcc2016-09-18 13:15:41 -0700795 need_to_create = [value]
796 # get unaccounted for bits
797 _, extra_flags = _decompose(cls, value)
798 # timer = 10
799 while extra_flags:
800 # timer -= 1
801 bit = _high_bit(extra_flags)
802 flag_value = 2 ** bit
803 if (flag_value not in cls._value2member_map_ and
804 flag_value not in need_to_create
805 ):
806 need_to_create.append(flag_value)
807 if extra_flags == -flag_value:
808 extra_flags = 0
809 else:
810 extra_flags ^= flag_value
811 for value in reversed(need_to_create):
812 # construct singleton pseudo-members
813 pseudo_member = int.__new__(cls, value)
814 pseudo_member._name_ = None
815 pseudo_member._value_ = value
Ethan Furman28cf6632017-01-24 12:12:06 -0800816 # use setdefault in case another thread already created a composite
817 # with this value
818 pseudo_member = cls._value2member_map_.setdefault(value, pseudo_member)
Ethan Furmanee47e5c2016-08-31 00:12:15 -0700819 return pseudo_member
820
Ethan Furmanee47e5c2016-08-31 00:12:15 -0700821 def __or__(self, other):
822 if not isinstance(other, (self.__class__, int)):
823 return NotImplemented
Ethan Furman3515dcc2016-09-18 13:15:41 -0700824 result = self.__class__(self._value_ | self.__class__(other)._value_)
825 return result
Ethan Furmanee47e5c2016-08-31 00:12:15 -0700826
827 def __and__(self, other):
828 if not isinstance(other, (self.__class__, int)):
829 return NotImplemented
830 return self.__class__(self._value_ & self.__class__(other)._value_)
831
832 def __xor__(self, other):
833 if not isinstance(other, (self.__class__, int)):
834 return NotImplemented
835 return self.__class__(self._value_ ^ self.__class__(other)._value_)
836
837 __ror__ = __or__
838 __rand__ = __and__
839 __rxor__ = __xor__
840
841 def __invert__(self):
Ethan Furman3515dcc2016-09-18 13:15:41 -0700842 result = self.__class__(~self._value_)
843 return result
Ethan Furmanee47e5c2016-08-31 00:12:15 -0700844
845
846def _high_bit(value):
Ethan Furman04439532016-09-02 15:50:21 -0700847 """returns index of highest bit, or -1 if value is zero or negative"""
Ethan Furman3515dcc2016-09-18 13:15:41 -0700848 return value.bit_length() - 1
Ethan Furmanee47e5c2016-08-31 00:12:15 -0700849
Ethan Furmanf24bb352013-07-18 17:05:39 -0700850def unique(enumeration):
851 """Class decorator for enumerations ensuring unique member values."""
852 duplicates = []
853 for name, member in enumeration.__members__.items():
854 if name != member.name:
855 duplicates.append((name, member.name))
856 if duplicates:
857 alias_details = ', '.join(
858 ["%s -> %s" % (alias, name) for (alias, name) in duplicates])
859 raise ValueError('duplicate values found in %r: %s' %
860 (enumeration, alias_details))
861 return enumeration
Ethan Furman3515dcc2016-09-18 13:15:41 -0700862
863def _decompose(flag, value):
864 """Extract all members from the value."""
865 # _decompose is only called if the value is not named
866 not_covered = value
867 negative = value < 0
Ethan Furman28cf6632017-01-24 12:12:06 -0800868 # issue29167: wrap accesses to _value2member_map_ in a list to avoid race
Ville Skyttä49b27342017-08-03 09:00:59 +0300869 # conditions between iterating over it and having more pseudo-
Ethan Furman28cf6632017-01-24 12:12:06 -0800870 # members added to it
Ethan Furman3515dcc2016-09-18 13:15:41 -0700871 if negative:
872 # only check for named flags
873 flags_to_check = [
874 (m, v)
Ethan Furman28cf6632017-01-24 12:12:06 -0800875 for v, m in list(flag._value2member_map_.items())
Ethan Furman3515dcc2016-09-18 13:15:41 -0700876 if m.name is not None
877 ]
878 else:
879 # check for named flags and powers-of-two flags
880 flags_to_check = [
881 (m, v)
Ethan Furman28cf6632017-01-24 12:12:06 -0800882 for v, m in list(flag._value2member_map_.items())
Ethan Furman3515dcc2016-09-18 13:15:41 -0700883 if m.name is not None or _power_of_two(v)
884 ]
885 members = []
886 for member, member_value in flags_to_check:
887 if member_value and member_value & value == member_value:
888 members.append(member)
889 not_covered &= ~member_value
890 if not members and value in flag._value2member_map_:
891 members.append(flag._value2member_map_[value])
892 members.sort(key=lambda m: m._value_, reverse=True)
893 if len(members) > 1 and members[0].value == value:
894 # we have the breakdown, don't need the value member itself
895 members.pop(0)
896 return members, not_covered
897
898def _power_of_two(value):
899 if value < 1:
900 return False
901 return value == 2 ** _high_bit(value)