blob: a958ed8748afea531f336f0b33c978bd88affd02 [file] [log] [blame]
Ethan Furman6b3d64a2013-06-14 16:55:46 -07001import sys
Ethan Furmane03ea372013-09-25 07:14:41 -07002from types import MappingProxyType, DynamicClassAttribute
Ethan Furman6b3d64a2013-06-14 16:55:46 -07003
Ethan Furmane5754ab2015-09-17 22:03:52 -07004
Ethan Furmanc16595e2016-09-10 23:36:59 -07005__all__ = [
6 'EnumMeta',
7 'Enum', 'IntEnum', 'Flag', 'IntFlag',
8 'auto', 'unique',
9 ]
Ethan Furman6b3d64a2013-06-14 16:55:46 -070010
11
Ethan Furman101e0742013-09-15 12:34:36 -070012def _is_descriptor(obj):
13 """Returns True if obj is a descriptor, False otherwise."""
14 return (
15 hasattr(obj, '__get__') or
16 hasattr(obj, '__set__') or
17 hasattr(obj, '__delete__'))
18
19
Ethan Furman6b3d64a2013-06-14 16:55:46 -070020def _is_dunder(name):
21 """Returns True if a __dunder__ name, False otherwise."""
22 return (name[:2] == name[-2:] == '__' and
23 name[2:3] != '_' and
Ethan Furman648f8602013-10-06 17:19:54 -070024 name[-3:-2] != '_' and
25 len(name) > 4)
Ethan Furman6b3d64a2013-06-14 16:55:46 -070026
27
28def _is_sunder(name):
29 """Returns True if a _sunder_ name, False otherwise."""
30 return (name[0] == name[-1] == '_' and
31 name[1:2] != '_' and
Ethan Furman648f8602013-10-06 17:19:54 -070032 name[-2:-1] != '_' and
33 len(name) > 2)
Ethan Furman6b3d64a2013-06-14 16:55:46 -070034
Ethan Furman6b3d64a2013-06-14 16:55:46 -070035def _make_class_unpicklable(cls):
36 """Make the given class un-picklable."""
Ethan Furmanca1b7942014-02-08 11:36:27 -080037 def _break_on_call_reduce(self, proto):
Ethan Furman6b3d64a2013-06-14 16:55:46 -070038 raise TypeError('%r cannot be pickled' % self)
Ethan Furmanca1b7942014-02-08 11:36:27 -080039 cls.__reduce_ex__ = _break_on_call_reduce
Ethan Furman6b3d64a2013-06-14 16:55:46 -070040 cls.__module__ = '<unknown>'
41
Ethan Furman3515dcc2016-09-18 13:15:41 -070042_auto_null = object()
Ethan Furmanc16595e2016-09-10 23:36:59 -070043class auto:
44 """
45 Instances are replaced with an appropriate value in Enum class suites.
46 """
Ethan Furman3515dcc2016-09-18 13:15:41 -070047 value = _auto_null
Ethan Furmanc16595e2016-09-10 23:36:59 -070048
Ethan Furman101e0742013-09-15 12:34:36 -070049
Ethan Furman6b3d64a2013-06-14 16:55:46 -070050class _EnumDict(dict):
Ethan Furman101e0742013-09-15 12:34:36 -070051 """Track enum member order and ensure member names are not reused.
Ethan Furman6b3d64a2013-06-14 16:55:46 -070052
53 EnumMeta will use the names found in self._member_names as the
54 enumeration member names.
55
56 """
57 def __init__(self):
58 super().__init__()
59 self._member_names = []
Ethan Furmanc16595e2016-09-10 23:36:59 -070060 self._last_values = []
Ethan Furmana4b1bb42018-01-22 07:56:37 -080061 self._ignore = []
Ethan Furman6b3d64a2013-06-14 16:55:46 -070062
63 def __setitem__(self, key, value):
Ethan Furman101e0742013-09-15 12:34:36 -070064 """Changes anything not dundered or not a descriptor.
Ethan Furman6b3d64a2013-06-14 16:55:46 -070065
66 If an enum member name is used twice, an error is raised; duplicate
67 values are not checked for.
68
69 Single underscore (sunder) names are reserved.
70
71 """
72 if _is_sunder(key):
Ethan Furmanee47e5c2016-08-31 00:12:15 -070073 if key not in (
Ethan Furman3515dcc2016-09-18 13:15:41 -070074 '_order_', '_create_pseudo_member_',
Ethan Furmana4b1bb42018-01-22 07:56:37 -080075 '_generate_next_value_', '_missing_', '_ignore_',
Ethan Furmanee47e5c2016-08-31 00:12:15 -070076 ):
Ethan Furmane8e61272016-08-20 07:19:31 -070077 raise ValueError('_names_ are reserved for future Enum use')
Ethan Furmanc16595e2016-09-10 23:36:59 -070078 if key == '_generate_next_value_':
79 setattr(self, '_generate_next_value', value)
Ethan Furmana4b1bb42018-01-22 07:56:37 -080080 elif key == '_ignore_':
81 if isinstance(value, str):
82 value = value.replace(',',' ').split()
83 else:
84 value = list(value)
85 self._ignore = value
86 already = set(value) & set(self._member_names)
87 if already:
88 raise ValueError('_ignore_ cannot specify already set names: %r' % (already, ))
Ethan Furman101e0742013-09-15 12:34:36 -070089 elif _is_dunder(key):
Ethan Furmane8e61272016-08-20 07:19:31 -070090 if key == '__order__':
91 key = '_order_'
Ethan Furman101e0742013-09-15 12:34:36 -070092 elif key in self._member_names:
93 # descriptor overwriting an enum?
94 raise TypeError('Attempted to reuse key: %r' % key)
Ethan Furmana4b1bb42018-01-22 07:56:37 -080095 elif key in self._ignore:
96 pass
Ethan Furman101e0742013-09-15 12:34:36 -070097 elif not _is_descriptor(value):
98 if key in self:
99 # enum overwriting a descriptor?
Ethan Furmanee47e5c2016-08-31 00:12:15 -0700100 raise TypeError('%r already defined as: %r' % (key, self[key]))
Ethan Furmanc16595e2016-09-10 23:36:59 -0700101 if isinstance(value, auto):
Ethan Furman3515dcc2016-09-18 13:15:41 -0700102 if value.value == _auto_null:
103 value.value = self._generate_next_value(key, 1, len(self._member_names), self._last_values[:])
104 value = value.value
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700105 self._member_names.append(key)
Ethan Furmanc16595e2016-09-10 23:36:59 -0700106 self._last_values.append(value)
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700107 super().__setitem__(key, value)
108
109
Ezio Melotti9a3777e2013-08-17 15:53:55 +0300110# Dummy value for Enum as EnumMeta explicitly checks for it, but of course
111# until EnumMeta finishes running the first time the Enum class doesn't exist.
112# This is also why there are checks in EnumMeta like `if Enum is not None`
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700113Enum = None
114
Ethan Furman332dbc72016-08-20 00:00:52 -0700115
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700116class EnumMeta(type):
117 """Metaclass for Enum"""
118 @classmethod
Ethan Furman332dbc72016-08-20 00:00:52 -0700119 def __prepare__(metacls, cls, bases):
Ethan Furmanee47e5c2016-08-31 00:12:15 -0700120 # create the namespace dict
121 enum_dict = _EnumDict()
122 # inherit previous flags and _generate_next_value_ function
123 member_type, first_enum = metacls._get_mixins_(bases)
124 if first_enum is not None:
125 enum_dict['_generate_next_value_'] = getattr(first_enum, '_generate_next_value_', None)
126 return enum_dict
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700127
Ethan Furman65a5a472016-09-01 23:55:19 -0700128 def __new__(metacls, cls, bases, classdict):
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700129 # an Enum class is final once enumeration items have been defined; it
130 # cannot be mixed with other types (int, float, etc.) if it has an
131 # inherited __new__ unless a new __new__ is defined (or the resulting
132 # class will fail).
Ethan Furmana4b1bb42018-01-22 07:56:37 -0800133 #
134 # remove any keys listed in _ignore_
135 classdict.setdefault('_ignore_', []).append('_ignore_')
136 ignore = classdict['_ignore_']
137 for key in ignore:
138 classdict.pop(key, None)
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700139 member_type, first_enum = metacls._get_mixins_(bases)
140 __new__, save_new, use_args = metacls._find_new_(classdict, member_type,
141 first_enum)
142
143 # save enum items into separate mapping so they don't get baked into
144 # the new class
Ethan Furmanee47e5c2016-08-31 00:12:15 -0700145 enum_members = {k: classdict[k] for k in classdict._member_names}
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700146 for name in classdict._member_names:
147 del classdict[name]
148
Ethan Furmane8e61272016-08-20 07:19:31 -0700149 # adjust the sunders
150 _order_ = classdict.pop('_order_', None)
151
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700152 # check for illegal enum names (any others?)
Ethan Furmanee47e5c2016-08-31 00:12:15 -0700153 invalid_names = set(enum_members) & {'mro', }
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700154 if invalid_names:
155 raise ValueError('Invalid enum member name: {0}'.format(
156 ','.join(invalid_names)))
157
Ethan Furman48a724f2015-04-11 23:23:06 -0700158 # create a default docstring if one has not been provided
159 if '__doc__' not in classdict:
160 classdict['__doc__'] = 'An enumeration.'
161
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700162 # create our new Enum type
163 enum_class = super().__new__(metacls, cls, bases, classdict)
Ethan Furman520ad572013-07-19 19:47:21 -0700164 enum_class._member_names_ = [] # names in definition order
INADA Naokie57f91a2018-06-19 01:14:26 +0900165 enum_class._member_map_ = {} # name->value map
Ethan Furman5e5a8232013-08-04 08:42:23 -0700166 enum_class._member_type_ = member_type
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700167
orlnub1230fb9fad2018-09-12 20:28:53 +0300168 # save DynamicClassAttribute attributes from super classes so we know
169 # if we can take the shortcut of storing members in the class dict
170 dynamic_attributes = {k for c in enum_class.mro()
171 for k, v in c.__dict__.items()
172 if isinstance(v, DynamicClassAttribute)}
Ethan Furman354ecf12015-03-11 08:43:12 -0700173
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700174 # Reverse value->name map for hashable values.
Ethan Furman520ad572013-07-19 19:47:21 -0700175 enum_class._value2member_map_ = {}
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700176
Ethan Furman2da95042014-03-03 12:42:52 -0800177 # If a custom type is mixed into the Enum, and it does not know how
178 # to pickle itself, pickle.dumps will succeed but pickle.loads will
179 # fail. Rather than have the error show up later and possibly far
180 # from the source, sabotage the pickle protocol for this class so
181 # that pickle.dumps also fails.
182 #
183 # However, if the new class implements its own __reduce_ex__, do not
184 # sabotage -- it's on them to make sure it works correctly. We use
185 # __reduce_ex__ instead of any of the others as it is preferred by
186 # pickle over __reduce__, and it handles all pickle protocols.
187 if '__reduce_ex__' not in classdict:
Ethan Furmandc870522014-02-18 12:37:12 -0800188 if member_type is not object:
189 methods = ('__getnewargs_ex__', '__getnewargs__',
190 '__reduce_ex__', '__reduce__')
Ethan Furman2da95042014-03-03 12:42:52 -0800191 if not any(m in member_type.__dict__ for m in methods):
Ethan Furmandc870522014-02-18 12:37:12 -0800192 _make_class_unpicklable(enum_class)
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700193
194 # instantiate them, checking for duplicates as we go
195 # we instantiate first instead of checking for duplicates first in case
196 # a custom __new__ is doing something funky with the values -- such as
197 # auto-numbering ;)
198 for member_name in classdict._member_names:
Ethan Furmanee47e5c2016-08-31 00:12:15 -0700199 value = enum_members[member_name]
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700200 if not isinstance(value, tuple):
201 args = (value, )
202 else:
203 args = value
204 if member_type is tuple: # special case for tuple enums
205 args = (args, ) # wrap it one more time
206 if not use_args:
207 enum_member = __new__(enum_class)
Ethan Furmanb41803e2013-07-25 13:50:45 -0700208 if not hasattr(enum_member, '_value_'):
209 enum_member._value_ = value
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700210 else:
211 enum_member = __new__(enum_class, *args)
Ethan Furmanb41803e2013-07-25 13:50:45 -0700212 if not hasattr(enum_member, '_value_'):
Ethan Furmanee47e5c2016-08-31 00:12:15 -0700213 if member_type is object:
214 enum_member._value_ = value
215 else:
216 enum_member._value_ = member_type(*args)
Ethan Furman520ad572013-07-19 19:47:21 -0700217 value = enum_member._value_
Ethan Furman520ad572013-07-19 19:47:21 -0700218 enum_member._name_ = member_name
Ethan Furmanc850f342013-09-15 16:59:35 -0700219 enum_member.__objclass__ = enum_class
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700220 enum_member.__init__(*args)
221 # If another member with the same value was already defined, the
222 # new member becomes an alias to the existing one.
Ethan Furman520ad572013-07-19 19:47:21 -0700223 for name, canonical_member in enum_class._member_map_.items():
Ethan Furman0081f232014-09-16 17:31:23 -0700224 if canonical_member._value_ == enum_member._value_:
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700225 enum_member = canonical_member
226 break
227 else:
228 # Aliases don't appear in member names (only in __members__).
Ethan Furman520ad572013-07-19 19:47:21 -0700229 enum_class._member_names_.append(member_name)
Ethan Furman354ecf12015-03-11 08:43:12 -0700230 # performance boost for any member that would not shadow
231 # a DynamicClassAttribute
orlnub1230fb9fad2018-09-12 20:28:53 +0300232 if member_name not in dynamic_attributes:
Ethan Furman354ecf12015-03-11 08:43:12 -0700233 setattr(enum_class, member_name, enum_member)
234 # now add to _member_map_
Ethan Furman520ad572013-07-19 19:47:21 -0700235 enum_class._member_map_[member_name] = enum_member
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700236 try:
237 # This may fail if value is not hashable. We can't add the value
238 # to the map, and by-value lookups for this value will be
239 # linear.
Ethan Furman520ad572013-07-19 19:47:21 -0700240 enum_class._value2member_map_[value] = enum_member
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700241 except TypeError:
242 pass
243
244 # double check that repr and friends are not the mixin's or various
245 # things break (such as pickle)
Ethan Furmandc870522014-02-18 12:37:12 -0800246 for name in ('__repr__', '__str__', '__format__', '__reduce_ex__'):
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700247 class_method = getattr(enum_class, name)
248 obj_method = getattr(member_type, name, None)
249 enum_method = getattr(first_enum, name, None)
250 if obj_method is not None and obj_method is class_method:
251 setattr(enum_class, name, enum_method)
252
253 # replace any other __new__ with our own (as long as Enum is not None,
254 # anyway) -- again, this is to support pickle
255 if Enum is not None:
256 # if the user defined their own __new__, save it before it gets
257 # clobbered in case they subclass later
258 if save_new:
259 enum_class.__new_member__ = __new__
260 enum_class.__new__ = Enum.__new__
Ethan Furmane8e61272016-08-20 07:19:31 -0700261
262 # py3 support for definition order (helps keep py2/py3 code in sync)
263 if _order_ is not None:
264 if isinstance(_order_, str):
265 _order_ = _order_.replace(',', ' ').split()
266 if _order_ != enum_class._member_names_:
267 raise TypeError('member order does not match _order_')
268
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700269 return enum_class
270
Ethan Furman5de67b12016-04-13 23:52:09 -0700271 def __bool__(self):
272 """
273 classes/types should always be True.
274 """
275 return True
276
Ethan Furmand9925a12014-09-16 20:35:55 -0700277 def __call__(cls, value, names=None, *, module=None, qualname=None, type=None, start=1):
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700278 """Either returns an existing member, or creates a new enum class.
279
280 This method is used both when an enum class is given a value to match
281 to an enumeration member (i.e. Color(3)) and for the functional API
Ethan Furman23bb6f42016-11-21 09:22:05 -0800282 (i.e. Color = Enum('Color', names='RED GREEN BLUE')).
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700283
Ethan Furman2da95042014-03-03 12:42:52 -0800284 When used for the functional API:
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700285
Ethan Furman2da95042014-03-03 12:42:52 -0800286 `value` will be the name of the new class.
287
288 `names` should be either a string of white-space/comma delimited names
Ethan Furmand9925a12014-09-16 20:35:55 -0700289 (values will start at `start`), or an iterator/mapping of name, value pairs.
Ethan Furman2da95042014-03-03 12:42:52 -0800290
291 `module` should be set to the module this class is being created in;
292 if it is not set, an attempt to find that module will be made, but if
293 it fails the class will not be picklable.
294
295 `qualname` should be set to the actual location this class can be found
296 at in its module; by default it is set to the global scope. If this is
297 not correct, unpickling will fail in some circumstances.
298
299 `type`, if set, will be mixed in as the first base class.
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700300
301 """
302 if names is None: # simple value lookup
303 return cls.__new__(cls, value)
304 # otherwise, functional API: we're creating a new Enum type
Ethan Furmand9925a12014-09-16 20:35:55 -0700305 return cls._create_(value, names, module=module, qualname=qualname, type=type, start=start)
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700306
307 def __contains__(cls, member):
Rahul Jha94306522018-09-10 23:51:04 +0530308 if not isinstance(member, Enum):
309 raise TypeError(
310 "unsupported operand type(s) for 'in': '%s' and '%s'" % (
311 type(member).__qualname__, cls.__class__.__qualname__))
Ethan Furman0081f232014-09-16 17:31:23 -0700312 return isinstance(member, cls) and member._name_ in cls._member_map_
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700313
Ethan Furman64a99722013-09-22 16:18:19 -0700314 def __delattr__(cls, attr):
315 # nicer error message when someone tries to delete an attribute
316 # (see issue19025).
317 if attr in cls._member_map_:
318 raise AttributeError(
319 "%s: cannot delete Enum member." % cls.__name__)
320 super().__delattr__(attr)
321
Ethan Furman388a3922013-08-12 06:51:41 -0700322 def __dir__(self):
Ethan Furman64a99722013-09-22 16:18:19 -0700323 return (['__class__', '__doc__', '__members__', '__module__'] +
324 self._member_names_)
Ethan Furman388a3922013-08-12 06:51:41 -0700325
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700326 def __getattr__(cls, name):
327 """Return the enum member matching `name`
328
329 We use __getattr__ instead of descriptors or inserting into the enum
330 class' __dict__ in order to support `name` and `value` being both
331 properties for enum members (which live in the class' __dict__) and
332 enum members themselves.
333
334 """
335 if _is_dunder(name):
336 raise AttributeError(name)
337 try:
Ethan Furman520ad572013-07-19 19:47:21 -0700338 return cls._member_map_[name]
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700339 except KeyError:
340 raise AttributeError(name) from None
341
342 def __getitem__(cls, name):
Ethan Furman520ad572013-07-19 19:47:21 -0700343 return cls._member_map_[name]
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700344
345 def __iter__(cls):
Ethan Furman520ad572013-07-19 19:47:21 -0700346 return (cls._member_map_[name] for name in cls._member_names_)
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700347
348 def __len__(cls):
Ethan Furman520ad572013-07-19 19:47:21 -0700349 return len(cls._member_names_)
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700350
Ethan Furman2131a4a2013-09-14 18:11:24 -0700351 @property
352 def __members__(cls):
353 """Returns a mapping of member name->value.
354
355 This mapping lists all enum members, including aliases. Note that this
356 is a read-only view of the internal mapping.
357
358 """
359 return MappingProxyType(cls._member_map_)
360
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700361 def __repr__(cls):
362 return "<enum %r>" % cls.__name__
363
Ethan Furman2131a4a2013-09-14 18:11:24 -0700364 def __reversed__(cls):
365 return (cls._member_map_[name] for name in reversed(cls._member_names_))
366
Ethan Furmanf203f2d2013-09-06 07:16:48 -0700367 def __setattr__(cls, name, value):
368 """Block attempts to reassign Enum members.
369
370 A simple assignment to the class namespace only changes one of the
371 several possible ways to get an Enum member from the Enum class,
372 resulting in an inconsistent Enumeration.
373
374 """
375 member_map = cls.__dict__.get('_member_map_', {})
376 if name in member_map:
377 raise AttributeError('Cannot reassign members.')
378 super().__setattr__(name, value)
379
anentropicb8e21f12018-04-16 04:40:35 +0100380 def _create_(cls, class_name, names, *, module=None, qualname=None, type=None, start=1):
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700381 """Convenience method to create a new Enum class.
382
383 `names` can be:
384
385 * A string containing member names, separated either with spaces or
Ethan Furmand9925a12014-09-16 20:35:55 -0700386 commas. Values are incremented by 1 from `start`.
387 * An iterable of member names. Values are incremented by 1 from `start`.
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700388 * An iterable of (member name, value) pairs.
Ethan Furmand9925a12014-09-16 20:35:55 -0700389 * A mapping of member name -> value pairs.
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700390
391 """
392 metacls = cls.__class__
393 bases = (cls, ) if type is None else (type, cls)
Ethan Furmanee47e5c2016-08-31 00:12:15 -0700394 _, first_enum = cls._get_mixins_(bases)
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700395 classdict = metacls.__prepare__(class_name, bases)
396
397 # special processing needed for names?
398 if isinstance(names, str):
399 names = names.replace(',', ' ').split()
Dong-hee Nadcc8ce42017-06-22 01:52:32 +0900400 if isinstance(names, (tuple, list)) and names and isinstance(names[0], str):
Ethan Furmanee47e5c2016-08-31 00:12:15 -0700401 original_names, names = names, []
Ethan Furmanc16595e2016-09-10 23:36:59 -0700402 last_values = []
Ethan Furmanee47e5c2016-08-31 00:12:15 -0700403 for count, name in enumerate(original_names):
Ethan Furmanc16595e2016-09-10 23:36:59 -0700404 value = first_enum._generate_next_value_(name, start, count, last_values[:])
405 last_values.append(value)
406 names.append((name, value))
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700407
408 # Here, names is either an iterable of (name, value) or a mapping.
409 for item in names:
410 if isinstance(item, str):
411 member_name, member_value = item, names[item]
412 else:
413 member_name, member_value = item
414 classdict[member_name] = member_value
415 enum_class = metacls.__new__(metacls, class_name, bases, classdict)
416
417 # TODO: replace the frame hack if a blessed way to know the calling
418 # module is ever developed
419 if module is None:
420 try:
421 module = sys._getframe(2).f_globals['__name__']
Rémi Lapeyre1fd06f12019-01-24 20:43:13 +0100422 except (AttributeError, ValueError, KeyError) as exc:
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700423 pass
424 if module is None:
425 _make_class_unpicklable(enum_class)
426 else:
427 enum_class.__module__ = module
Ethan Furmanca1b7942014-02-08 11:36:27 -0800428 if qualname is not None:
429 enum_class.__qualname__ = qualname
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700430
431 return enum_class
432
orlnub1230fb9fad2018-09-12 20:28:53 +0300433 def _convert_(cls, name, module, filter, source=None):
434 """
435 Create a new Enum subclass that replaces a collection of global constants
436 """
437 # convert all constants from source (or module) that pass filter() to
438 # a new Enum called name, and export the enum and its members back to
439 # module;
440 # also, replace the __reduce_ex__ method so unpickling works in
441 # previous Python versions
442 module_globals = vars(sys.modules[module])
443 if source:
444 source = vars(source)
445 else:
446 source = module_globals
447 # _value2member_map_ is populated in the same order every time
448 # for a consistent reverse mapping of number to name when there
449 # are multiple names for the same number.
450 members = [
451 (name, value)
452 for name, value in source.items()
453 if filter(name)]
454 try:
455 # sort by value
456 members.sort(key=lambda t: (t[1], t[0]))
457 except TypeError:
458 # unless some values aren't comparable, in which case sort by name
459 members.sort(key=lambda t: t[0])
460 cls = cls(name, members, module=module)
461 cls.__reduce_ex__ = _reduce_ex_by_name
462 module_globals.update(cls.__members__)
463 module_globals[name] = cls
464 return cls
465
466 def _convert(cls, *args, **kwargs):
467 import warnings
468 warnings.warn("_convert is deprecated and will be removed in 3.9, use "
469 "_convert_ instead.", DeprecationWarning, stacklevel=2)
470 return cls._convert_(*args, **kwargs)
471
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700472 @staticmethod
473 def _get_mixins_(bases):
474 """Returns the type for creating enum members, and the first inherited
475 enum class.
476
477 bases: the tuple of bases that was given to __new__
478
479 """
480 if not bases:
481 return object, Enum
482
Ethan Furman5bdab642018-09-21 19:03:09 -0700483 def _find_data_type(bases):
484 for chain in bases:
485 for base in chain.__mro__:
486 if base is object:
487 continue
488 elif '__new__' in base.__dict__:
Ethan Furmancd453852018-10-05 23:29:36 -0700489 if issubclass(base, Enum):
Ethan Furman5bdab642018-09-21 19:03:09 -0700490 continue
491 return base
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700492
Ethan Furman5bdab642018-09-21 19:03:09 -0700493 # ensure final parent class is an Enum derivative, find any concrete
494 # data type, and check that Enum has no members
495 first_enum = bases[-1]
496 if not issubclass(first_enum, Enum):
497 raise TypeError("new enumerations should be created as "
498 "`EnumName([mixin_type, ...] [data_type,] enum_type)`")
499 member_type = _find_data_type(bases) or object
500 if first_enum._member_names_:
501 raise TypeError("Cannot extend enumerations")
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700502 return member_type, first_enum
503
504 @staticmethod
505 def _find_new_(classdict, member_type, first_enum):
506 """Returns the __new__ to be used for creating the enum members.
507
508 classdict: the class dictionary given to __new__
509 member_type: the data type whose __new__ will be used by default
510 first_enum: enumeration to check for an overriding __new__
511
512 """
513 # now find the correct __new__, checking to see of one was defined
514 # by the user; also check earlier enum classes in case a __new__ was
515 # saved as __new_member__
516 __new__ = classdict.get('__new__', None)
517
518 # should __new__ be saved as __new_member__ later?
519 save_new = __new__ is not None
520
521 if __new__ is None:
522 # check all possibles for __new_member__ before falling back to
523 # __new__
524 for method in ('__new_member__', '__new__'):
525 for possible in (member_type, first_enum):
526 target = getattr(possible, method, None)
527 if target not in {
528 None,
529 None.__new__,
530 object.__new__,
531 Enum.__new__,
532 }:
533 __new__ = target
534 break
535 if __new__ is not None:
536 break
537 else:
538 __new__ = object.__new__
539
540 # if a non-object.__new__ is used then whatever value/tuple was
541 # assigned to the enum member name will be passed to __new__ and to the
542 # new enum member's __init__
543 if __new__ is object.__new__:
544 use_args = False
545 else:
546 use_args = True
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700547 return __new__, save_new, use_args
548
549
550class Enum(metaclass=EnumMeta):
551 """Generic enumeration.
552
553 Derive from this class to define new enumerations.
554
555 """
556 def __new__(cls, value):
557 # all enum instances are actually created during class construction
558 # without calling this method; this method is called by the metaclass'
559 # __call__ (i.e. Color(3) ), and by pickle
560 if type(value) is cls:
Ethan Furman23bb6f42016-11-21 09:22:05 -0800561 # For lookups like Color(Color.RED)
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700562 return value
563 # by-value search for a matching enum member
564 # see if it's in the reverse mapping (for hashable values)
Ethan Furman2aa27322013-07-19 19:35:56 -0700565 try:
Andrew Svetlov34ae04f2018-12-26 20:45:33 +0200566 return cls._value2member_map_[value]
567 except KeyError:
568 # Not found, no need to do long O(n) search
569 pass
Ethan Furman2aa27322013-07-19 19:35:56 -0700570 except TypeError:
571 # not there, now do long search -- O(n) behavior
Ethan Furman520ad572013-07-19 19:47:21 -0700572 for member in cls._member_map_.values():
Ethan Furman0081f232014-09-16 17:31:23 -0700573 if member._value_ == value:
Ethan Furman2aa27322013-07-19 19:35:56 -0700574 return member
Ethan Furmanee47e5c2016-08-31 00:12:15 -0700575 # still not found -- try _missing_ hook
Ethan Furman019f0a02018-09-12 11:43:34 -0700576 try:
577 exc = None
578 result = cls._missing_(value)
579 except Exception as e:
580 exc = e
581 result = None
582 if isinstance(result, cls):
583 return result
584 else:
585 ve_exc = ValueError("%r is not a valid %s" % (value, cls.__name__))
586 if result is None and exc is None:
587 raise ve_exc
588 elif exc is None:
589 exc = TypeError(
590 'error in %s._missing_: returned %r instead of None or a valid member'
591 % (cls.__name__, result)
592 )
593 exc.__context__ = ve_exc
594 raise exc
Ethan Furmanee47e5c2016-08-31 00:12:15 -0700595
Ethan Furmanc16595e2016-09-10 23:36:59 -0700596 def _generate_next_value_(name, start, count, last_values):
597 for last_value in reversed(last_values):
598 try:
599 return last_value + 1
600 except TypeError:
601 pass
602 else:
Ethan Furmanee47e5c2016-08-31 00:12:15 -0700603 return start
Ethan Furmanc16595e2016-09-10 23:36:59 -0700604
Ethan Furmanee47e5c2016-08-31 00:12:15 -0700605 @classmethod
606 def _missing_(cls, value):
Ethan Furman0081f232014-09-16 17:31:23 -0700607 raise ValueError("%r is not a valid %s" % (value, cls.__name__))
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700608
609 def __repr__(self):
610 return "<%s.%s: %r>" % (
Ethan Furman520ad572013-07-19 19:47:21 -0700611 self.__class__.__name__, self._name_, self._value_)
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700612
613 def __str__(self):
Ethan Furman520ad572013-07-19 19:47:21 -0700614 return "%s.%s" % (self.__class__.__name__, self._name_)
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700615
Ethan Furman388a3922013-08-12 06:51:41 -0700616 def __dir__(self):
Ethan Furman0ae550b2014-10-14 08:58:32 -0700617 added_behavior = [
618 m
619 for cls in self.__class__.mro()
620 for m in cls.__dict__
Ethan Furman354ecf12015-03-11 08:43:12 -0700621 if m[0] != '_' and m not in self._member_map_
Ethan Furman0ae550b2014-10-14 08:58:32 -0700622 ]
Ethan Furmanec5f8eb2014-10-21 13:40:35 -0700623 return (['__class__', '__doc__', '__module__'] + added_behavior)
Ethan Furman388a3922013-08-12 06:51:41 -0700624
Ethan Furmanec15a822013-08-31 19:17:41 -0700625 def __format__(self, format_spec):
626 # mixed-in Enums should use the mixed-in type's __format__, otherwise
627 # we can get strange results with the Enum name showing up instead of
628 # the value
629
630 # pure Enum branch
631 if self._member_type_ is object:
632 cls = str
633 val = str(self)
634 # mix-in branch
635 else:
636 cls = self._member_type_
Ethan Furman0081f232014-09-16 17:31:23 -0700637 val = self._value_
Ethan Furmanec15a822013-08-31 19:17:41 -0700638 return cls.__format__(val, format_spec)
639
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700640 def __hash__(self):
Ethan Furman520ad572013-07-19 19:47:21 -0700641 return hash(self._name_)
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700642
Ethan Furmanca1b7942014-02-08 11:36:27 -0800643 def __reduce_ex__(self, proto):
Ethan Furmandc870522014-02-18 12:37:12 -0800644 return self.__class__, (self._value_, )
Ethan Furmanca1b7942014-02-08 11:36:27 -0800645
Ethan Furman33918c12013-09-27 23:02:02 -0700646 # DynamicClassAttribute is used to provide access to the `name` and
647 # `value` properties of enum members while keeping some measure of
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700648 # protection from modification, while still allowing for an enumeration
649 # to have members named `name` and `value`. This works because enumeration
650 # members are not set directly on the enum class -- __getattr__ is
651 # used to look them up.
652
Ethan Furmane03ea372013-09-25 07:14:41 -0700653 @DynamicClassAttribute
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700654 def name(self):
Ethan Furmanc850f342013-09-15 16:59:35 -0700655 """The name of the Enum member."""
Ethan Furman520ad572013-07-19 19:47:21 -0700656 return self._name_
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700657
Ethan Furmane03ea372013-09-25 07:14:41 -0700658 @DynamicClassAttribute
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700659 def value(self):
Ethan Furmanc850f342013-09-15 16:59:35 -0700660 """The value of the Enum member."""
Ethan Furman520ad572013-07-19 19:47:21 -0700661 return self._value_
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700662
663
664class IntEnum(int, Enum):
665 """Enum where members are also (and must be) ints"""
Ethan Furmanf24bb352013-07-18 17:05:39 -0700666
667
Ethan Furman24e837f2015-03-18 17:27:57 -0700668def _reduce_ex_by_name(self, proto):
669 return self.name
670
Ethan Furman65a5a472016-09-01 23:55:19 -0700671class Flag(Enum):
Ethan Furmanee47e5c2016-08-31 00:12:15 -0700672 """Support for flags"""
Ethan Furmanc16595e2016-09-10 23:36:59 -0700673
674 def _generate_next_value_(name, start, count, last_values):
Ethan Furmanee47e5c2016-08-31 00:12:15 -0700675 """
676 Generate the next value when not given.
677
678 name: the name of the member
679 start: the initital start value or None
680 count: the number of existing members
681 last_value: the last value assigned or None
682 """
683 if not count:
684 return start if start is not None else 1
Ethan Furmanc16595e2016-09-10 23:36:59 -0700685 for last_value in reversed(last_values):
686 try:
687 high_bit = _high_bit(last_value)
688 break
Ethan Furman3515dcc2016-09-18 13:15:41 -0700689 except Exception:
Ethan Furmanc16595e2016-09-10 23:36:59 -0700690 raise TypeError('Invalid Flag value: %r' % last_value) from None
Ethan Furmanee47e5c2016-08-31 00:12:15 -0700691 return 2 ** (high_bit+1)
692
693 @classmethod
694 def _missing_(cls, value):
695 original_value = value
696 if value < 0:
697 value = ~value
698 possible_member = cls._create_pseudo_member_(value)
Ethan Furmanee47e5c2016-08-31 00:12:15 -0700699 if original_value < 0:
700 possible_member = ~possible_member
701 return possible_member
702
703 @classmethod
704 def _create_pseudo_member_(cls, value):
Ethan Furman3515dcc2016-09-18 13:15:41 -0700705 """
706 Create a composite member iff value contains only members.
707 """
Ethan Furmanee47e5c2016-08-31 00:12:15 -0700708 pseudo_member = cls._value2member_map_.get(value, None)
709 if pseudo_member is None:
Ethan Furman3515dcc2016-09-18 13:15:41 -0700710 # verify all bits are accounted for
711 _, extra_flags = _decompose(cls, value)
712 if extra_flags:
713 raise ValueError("%r is not a valid %s" % (value, cls.__name__))
714 # construct a singleton enum pseudo-member
Ethan Furmanee47e5c2016-08-31 00:12:15 -0700715 pseudo_member = object.__new__(cls)
716 pseudo_member._name_ = None
717 pseudo_member._value_ = value
Ethan Furman28cf6632017-01-24 12:12:06 -0800718 # use setdefault in case another thread already created a composite
719 # with this value
720 pseudo_member = cls._value2member_map_.setdefault(value, pseudo_member)
Ethan Furmanee47e5c2016-08-31 00:12:15 -0700721 return pseudo_member
722
Ethan Furmanee47e5c2016-08-31 00:12:15 -0700723 def __contains__(self, other):
724 if not isinstance(other, self.__class__):
Rahul Jha94306522018-09-10 23:51:04 +0530725 raise TypeError(
726 "unsupported operand type(s) for 'in': '%s' and '%s'" % (
727 type(other).__qualname__, self.__class__.__qualname__))
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)