blob: 9dfba3c37cd2cd5d269ba2ef724f50f1ef1c52a5 [file] [log] [blame]
Ethan Furman6b3d64a2013-06-14 16:55:46 -07001import sys
Ethan Furmane03ea372013-09-25 07:14:41 -07002from types import MappingProxyType, DynamicClassAttribute
Ethan Furman6b3d64a2013-06-14 16:55:46 -07003
Ethan Furmane5754ab2015-09-17 22:03:52 -07004
Ethan Furmanc16595e2016-09-10 23:36:59 -07005__all__ = [
6 'EnumMeta',
7 'Enum', 'IntEnum', 'Flag', 'IntFlag',
8 'auto', 'unique',
9 ]
Ethan Furman6b3d64a2013-06-14 16:55:46 -070010
11
Ethan Furman101e0742013-09-15 12:34:36 -070012def _is_descriptor(obj):
13 """Returns True if obj is a descriptor, False otherwise."""
14 return (
15 hasattr(obj, '__get__') or
16 hasattr(obj, '__set__') or
17 hasattr(obj, '__delete__'))
18
19
Ethan Furman6b3d64a2013-06-14 16:55:46 -070020def _is_dunder(name):
21 """Returns True if a __dunder__ name, False otherwise."""
Brennan D Baraban8b914d22019-03-03 14:09:11 -080022 return (len(name) > 4 and
23 name[:2] == name[-2:] == '__' and
24 name[2] != '_' and
25 name[-3] != '_')
Ethan Furman6b3d64a2013-06-14 16:55:46 -070026
27
28def _is_sunder(name):
29 """Returns True if a _sunder_ name, False otherwise."""
Brennan D Baraban8b914d22019-03-03 14:09:11 -080030 return (len(name) > 2 and
31 name[0] == name[-1] == '_' and
Ethan Furman6b3d64a2013-06-14 16:55:46 -070032 name[1:2] != '_' and
Brennan D Baraban8b914d22019-03-03 14:09:11 -080033 name[-2:-1] != '_')
34
Ethan Furman6b3d64a2013-06-14 16:55:46 -070035
Ethan Furman6b3d64a2013-06-14 16:55:46 -070036def _make_class_unpicklable(cls):
37 """Make the given class un-picklable."""
Ethan Furmanca1b7942014-02-08 11:36:27 -080038 def _break_on_call_reduce(self, proto):
Ethan Furman6b3d64a2013-06-14 16:55:46 -070039 raise TypeError('%r cannot be pickled' % self)
Ethan Furmanca1b7942014-02-08 11:36:27 -080040 cls.__reduce_ex__ = _break_on_call_reduce
Ethan Furman6b3d64a2013-06-14 16:55:46 -070041 cls.__module__ = '<unknown>'
42
Ethan Furman3515dcc2016-09-18 13:15:41 -070043_auto_null = object()
Ethan Furmanc16595e2016-09-10 23:36:59 -070044class auto:
45 """
46 Instances are replaced with an appropriate value in Enum class suites.
47 """
Ethan Furman3515dcc2016-09-18 13:15:41 -070048 value = _auto_null
Ethan Furmanc16595e2016-09-10 23:36:59 -070049
Ethan Furman101e0742013-09-15 12:34:36 -070050
Ethan Furman6b3d64a2013-06-14 16:55:46 -070051class _EnumDict(dict):
Ethan Furman101e0742013-09-15 12:34:36 -070052 """Track enum member order and ensure member names are not reused.
Ethan Furman6b3d64a2013-06-14 16:55:46 -070053
54 EnumMeta will use the names found in self._member_names as the
55 enumeration member names.
56
57 """
58 def __init__(self):
59 super().__init__()
60 self._member_names = []
Ethan Furmanc16595e2016-09-10 23:36:59 -070061 self._last_values = []
Ethan Furmana4b1bb42018-01-22 07:56:37 -080062 self._ignore = []
Ethan Onstottd9a43e22020-04-28 13:20:55 -040063 self._auto_called = False
Ethan Furman6b3d64a2013-06-14 16:55:46 -070064
65 def __setitem__(self, key, value):
Ethan Furman101e0742013-09-15 12:34:36 -070066 """Changes anything not dundered or not a descriptor.
Ethan Furman6b3d64a2013-06-14 16:55:46 -070067
68 If an enum member name is used twice, an error is raised; duplicate
69 values are not checked for.
70
71 Single underscore (sunder) names are reserved.
72
73 """
74 if _is_sunder(key):
Ethan Furmanee47e5c2016-08-31 00:12:15 -070075 if key not in (
Ethan Furman3515dcc2016-09-18 13:15:41 -070076 '_order_', '_create_pseudo_member_',
Ethan Furmana4b1bb42018-01-22 07:56:37 -080077 '_generate_next_value_', '_missing_', '_ignore_',
Ethan Furmanee47e5c2016-08-31 00:12:15 -070078 ):
Ethan Furmane8e61272016-08-20 07:19:31 -070079 raise ValueError('_names_ are reserved for future Enum use')
Ethan Furmanc16595e2016-09-10 23:36:59 -070080 if key == '_generate_next_value_':
Ethan Onstottd9a43e22020-04-28 13:20:55 -040081 # check if members already defined as auto()
82 if self._auto_called:
83 raise TypeError("_generate_next_value_ must be defined before members")
Ethan Furmanc16595e2016-09-10 23:36:59 -070084 setattr(self, '_generate_next_value', value)
Ethan Furmana4b1bb42018-01-22 07:56:37 -080085 elif key == '_ignore_':
86 if isinstance(value, str):
87 value = value.replace(',',' ').split()
88 else:
89 value = list(value)
90 self._ignore = value
91 already = set(value) & set(self._member_names)
92 if already:
93 raise ValueError('_ignore_ cannot specify already set names: %r' % (already, ))
Ethan Furman101e0742013-09-15 12:34:36 -070094 elif _is_dunder(key):
Ethan Furmane8e61272016-08-20 07:19:31 -070095 if key == '__order__':
96 key = '_order_'
Ethan Furman101e0742013-09-15 12:34:36 -070097 elif key in self._member_names:
98 # descriptor overwriting an enum?
99 raise TypeError('Attempted to reuse key: %r' % key)
Ethan Furmana4b1bb42018-01-22 07:56:37 -0800100 elif key in self._ignore:
101 pass
Ethan Furman101e0742013-09-15 12:34:36 -0700102 elif not _is_descriptor(value):
103 if key in self:
104 # enum overwriting a descriptor?
Ethan Furmanee47e5c2016-08-31 00:12:15 -0700105 raise TypeError('%r already defined as: %r' % (key, self[key]))
Ethan Furmanc16595e2016-09-10 23:36:59 -0700106 if isinstance(value, auto):
Ethan Onstottd9a43e22020-04-28 13:20:55 -0400107 self._auto_called = True
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?)
Brennan D Baraban8b914d22019-03-03 14:09:11 -0800159 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
INADA Naokie57f91a2018-06-19 01:14:26 +0900171 enum_class._member_map_ = {} # 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
orlnub1230fb9fad2018-09-12 20:28:53 +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
orlnub1230fb9fad2018-09-12 20:28:53 +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):
Rahul Jha94306522018-09-10 23:51:04 +0530314 if not isinstance(member, Enum):
315 raise TypeError(
316 "unsupported operand type(s) for 'in': '%s' and '%s'" % (
317 type(member).__qualname__, cls.__class__.__qualname__))
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
anentropicb8e21f12018-04-16 04:40:35 +0100386 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__']
Pablo Galindo293dd232019-11-19 21:34:03 +0000428 except (AttributeError, ValueError, KeyError):
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700429 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
orlnub1230fb9fad2018-09-12 20:28:53 +0300439 def _convert_(cls, name, module, filter, source=None):
440 """
441 Create a new Enum subclass that replaces a collection of global constants
442 """
443 # convert all constants from source (or module) that pass filter() to
444 # a new Enum called name, and export the enum and its members back to
445 # module;
446 # also, replace the __reduce_ex__ method so unpickling works in
447 # previous Python versions
448 module_globals = vars(sys.modules[module])
449 if source:
450 source = vars(source)
451 else:
452 source = module_globals
453 # _value2member_map_ is populated in the same order every time
454 # for a consistent reverse mapping of number to name when there
455 # are multiple names for the same number.
456 members = [
457 (name, value)
458 for name, value in source.items()
459 if filter(name)]
460 try:
461 # sort by value
462 members.sort(key=lambda t: (t[1], t[0]))
463 except TypeError:
464 # unless some values aren't comparable, in which case sort by name
465 members.sort(key=lambda t: t[0])
466 cls = cls(name, members, module=module)
467 cls.__reduce_ex__ = _reduce_ex_by_name
468 module_globals.update(cls.__members__)
469 module_globals[name] = cls
470 return cls
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):
Miss Islington (bot)95b81e22020-09-15 16:59:48 -0700484 data_types = []
Ethan Furman5bdab642018-09-21 19:03:09 -0700485 for chain in bases:
Miss Islington (bot)95b81e22020-09-15 16:59:48 -0700486 candidate = None
Ethan Furman5bdab642018-09-21 19:03:09 -0700487 for base in chain.__mro__:
488 if base is object:
489 continue
490 elif '__new__' in base.__dict__:
Ethan Furmancd453852018-10-05 23:29:36 -0700491 if issubclass(base, Enum):
Ethan Furman5bdab642018-09-21 19:03:09 -0700492 continue
Miss Islington (bot)95b81e22020-09-15 16:59:48 -0700493 data_types.append(candidate or base)
494 break
495 elif not issubclass(base, Enum):
496 candidate = base
497 if len(data_types) > 1:
498 raise TypeError('too many data types: %r' % data_types)
499 elif data_types:
500 return data_types[0]
501 else:
502 return None
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700503
Ethan Furman5bdab642018-09-21 19:03:09 -0700504 # ensure final parent class is an Enum derivative, find any concrete
505 # data type, and check that Enum has no members
506 first_enum = bases[-1]
507 if not issubclass(first_enum, Enum):
508 raise TypeError("new enumerations should be created as "
509 "`EnumName([mixin_type, ...] [data_type,] enum_type)`")
510 member_type = _find_data_type(bases) or object
511 if first_enum._member_names_:
512 raise TypeError("Cannot extend enumerations")
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700513 return member_type, first_enum
514
515 @staticmethod
516 def _find_new_(classdict, member_type, first_enum):
517 """Returns the __new__ to be used for creating the enum members.
518
519 classdict: the class dictionary given to __new__
520 member_type: the data type whose __new__ will be used by default
521 first_enum: enumeration to check for an overriding __new__
522
523 """
524 # now find the correct __new__, checking to see of one was defined
525 # by the user; also check earlier enum classes in case a __new__ was
526 # saved as __new_member__
527 __new__ = classdict.get('__new__', None)
528
529 # should __new__ be saved as __new_member__ later?
530 save_new = __new__ is not None
531
532 if __new__ is None:
533 # check all possibles for __new_member__ before falling back to
534 # __new__
535 for method in ('__new_member__', '__new__'):
536 for possible in (member_type, first_enum):
537 target = getattr(possible, method, None)
538 if target not in {
539 None,
540 None.__new__,
541 object.__new__,
542 Enum.__new__,
543 }:
544 __new__ = target
545 break
546 if __new__ is not None:
547 break
548 else:
549 __new__ = object.__new__
550
551 # if a non-object.__new__ is used then whatever value/tuple was
552 # assigned to the enum member name will be passed to __new__ and to the
553 # new enum member's __init__
554 if __new__ is object.__new__:
555 use_args = False
556 else:
557 use_args = True
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700558 return __new__, save_new, use_args
559
560
561class Enum(metaclass=EnumMeta):
562 """Generic enumeration.
563
564 Derive from this class to define new enumerations.
565
566 """
567 def __new__(cls, value):
568 # all enum instances are actually created during class construction
569 # without calling this method; this method is called by the metaclass'
570 # __call__ (i.e. Color(3) ), and by pickle
571 if type(value) is cls:
Ethan Furman23bb6f42016-11-21 09:22:05 -0800572 # For lookups like Color(Color.RED)
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700573 return value
574 # by-value search for a matching enum member
575 # see if it's in the reverse mapping (for hashable values)
Ethan Furman2aa27322013-07-19 19:35:56 -0700576 try:
Andrew Svetlov34ae04f2018-12-26 20:45:33 +0200577 return cls._value2member_map_[value]
578 except KeyError:
579 # Not found, no need to do long O(n) search
580 pass
Ethan Furman2aa27322013-07-19 19:35:56 -0700581 except TypeError:
582 # not there, now do long search -- O(n) behavior
Ethan Furman520ad572013-07-19 19:47:21 -0700583 for member in cls._member_map_.values():
Ethan Furman0081f232014-09-16 17:31:23 -0700584 if member._value_ == value:
Ethan Furman2aa27322013-07-19 19:35:56 -0700585 return member
Ethan Furmanee47e5c2016-08-31 00:12:15 -0700586 # still not found -- try _missing_ hook
Ethan Furman019f0a02018-09-12 11:43:34 -0700587 try:
588 exc = None
589 result = cls._missing_(value)
590 except Exception as e:
591 exc = e
592 result = None
593 if isinstance(result, cls):
594 return result
595 else:
Walter Dörwald323842c2019-07-18 20:37:13 +0200596 ve_exc = ValueError("%r is not a valid %s" % (value, cls.__qualname__))
Ethan Furman019f0a02018-09-12 11:43:34 -0700597 if result is None and exc is None:
598 raise ve_exc
599 elif exc is None:
600 exc = TypeError(
601 'error in %s._missing_: returned %r instead of None or a valid member'
602 % (cls.__name__, result)
603 )
604 exc.__context__ = ve_exc
605 raise exc
Ethan Furmanee47e5c2016-08-31 00:12:15 -0700606
Ethan Furmanc16595e2016-09-10 23:36:59 -0700607 def _generate_next_value_(name, start, count, last_values):
608 for last_value in reversed(last_values):
609 try:
610 return last_value + 1
611 except TypeError:
612 pass
613 else:
Ethan Furmanee47e5c2016-08-31 00:12:15 -0700614 return start
Ethan Furmanc16595e2016-09-10 23:36:59 -0700615
Ethan Furmanee47e5c2016-08-31 00:12:15 -0700616 @classmethod
617 def _missing_(cls, value):
Walter Dörwald323842c2019-07-18 20:37:13 +0200618 raise ValueError("%r is not a valid %s" % (value, cls.__qualname__))
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700619
620 def __repr__(self):
621 return "<%s.%s: %r>" % (
Ethan Furman520ad572013-07-19 19:47:21 -0700622 self.__class__.__name__, self._name_, self._value_)
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700623
624 def __str__(self):
Ethan Furman520ad572013-07-19 19:47:21 -0700625 return "%s.%s" % (self.__class__.__name__, self._name_)
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700626
Ethan Furman388a3922013-08-12 06:51:41 -0700627 def __dir__(self):
Ethan Furman0ae550b2014-10-14 08:58:32 -0700628 added_behavior = [
629 m
630 for cls in self.__class__.mro()
631 for m in cls.__dict__
Ethan Furman354ecf12015-03-11 08:43:12 -0700632 if m[0] != '_' and m not in self._member_map_
Ethan Furman0ae550b2014-10-14 08:58:32 -0700633 ]
Ethan Furmanec5f8eb2014-10-21 13:40:35 -0700634 return (['__class__', '__doc__', '__module__'] + added_behavior)
Ethan Furman388a3922013-08-12 06:51:41 -0700635
Ethan Furmanec15a822013-08-31 19:17:41 -0700636 def __format__(self, format_spec):
637 # mixed-in Enums should use the mixed-in type's __format__, otherwise
638 # we can get strange results with the Enum name showing up instead of
639 # the value
640
thatneat2f19e822019-07-04 11:28:37 -0700641 # pure Enum branch, or branch with __str__ explicitly overridden
642 str_overridden = type(self).__str__ != Enum.__str__
643 if self._member_type_ is object or str_overridden:
Ethan Furmanec15a822013-08-31 19:17:41 -0700644 cls = str
645 val = str(self)
646 # mix-in branch
647 else:
648 cls = self._member_type_
Ethan Furman0081f232014-09-16 17:31:23 -0700649 val = self._value_
Ethan Furmanec15a822013-08-31 19:17:41 -0700650 return cls.__format__(val, format_spec)
651
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700652 def __hash__(self):
Ethan Furman520ad572013-07-19 19:47:21 -0700653 return hash(self._name_)
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700654
Ethan Furmanca1b7942014-02-08 11:36:27 -0800655 def __reduce_ex__(self, proto):
Ethan Furmandc870522014-02-18 12:37:12 -0800656 return self.__class__, (self._value_, )
Ethan Furmanca1b7942014-02-08 11:36:27 -0800657
Ethan Furman33918c12013-09-27 23:02:02 -0700658 # DynamicClassAttribute is used to provide access to the `name` and
659 # `value` properties of enum members while keeping some measure of
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700660 # protection from modification, while still allowing for an enumeration
661 # to have members named `name` and `value`. This works because enumeration
662 # members are not set directly on the enum class -- __getattr__ is
663 # used to look them up.
664
Ethan Furmane03ea372013-09-25 07:14:41 -0700665 @DynamicClassAttribute
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700666 def name(self):
Ethan Furmanc850f342013-09-15 16:59:35 -0700667 """The name of the Enum member."""
Ethan Furman520ad572013-07-19 19:47:21 -0700668 return self._name_
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700669
Ethan Furmane03ea372013-09-25 07:14:41 -0700670 @DynamicClassAttribute
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700671 def value(self):
Ethan Furmanc850f342013-09-15 16:59:35 -0700672 """The value of the Enum member."""
Ethan Furman520ad572013-07-19 19:47:21 -0700673 return self._value_
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700674
675
676class IntEnum(int, Enum):
677 """Enum where members are also (and must be) ints"""
Ethan Furmanf24bb352013-07-18 17:05:39 -0700678
679
Ethan Furman24e837f2015-03-18 17:27:57 -0700680def _reduce_ex_by_name(self, proto):
681 return self.name
682
Ethan Furman65a5a472016-09-01 23:55:19 -0700683class Flag(Enum):
Ethan Furmanee47e5c2016-08-31 00:12:15 -0700684 """Support for flags"""
Ethan Furmanc16595e2016-09-10 23:36:59 -0700685
686 def _generate_next_value_(name, start, count, last_values):
Ethan Furmanee47e5c2016-08-31 00:12:15 -0700687 """
688 Generate the next value when not given.
689
690 name: the name of the member
HongWeipengbb16fb22019-09-21 13:22:54 +0800691 start: the initial start value or None
Ethan Furmanee47e5c2016-08-31 00:12:15 -0700692 count: the number of existing members
693 last_value: the last value assigned or None
694 """
695 if not count:
696 return start if start is not None else 1
Ethan Furmanc16595e2016-09-10 23:36:59 -0700697 for last_value in reversed(last_values):
698 try:
699 high_bit = _high_bit(last_value)
700 break
Ethan Furman3515dcc2016-09-18 13:15:41 -0700701 except Exception:
Ethan Furmanc16595e2016-09-10 23:36:59 -0700702 raise TypeError('Invalid Flag value: %r' % last_value) from None
Ethan Furmanee47e5c2016-08-31 00:12:15 -0700703 return 2 ** (high_bit+1)
704
705 @classmethod
706 def _missing_(cls, value):
707 original_value = value
708 if value < 0:
709 value = ~value
710 possible_member = cls._create_pseudo_member_(value)
Ethan Furmanee47e5c2016-08-31 00:12:15 -0700711 if original_value < 0:
712 possible_member = ~possible_member
713 return possible_member
714
715 @classmethod
716 def _create_pseudo_member_(cls, value):
Ethan Furman3515dcc2016-09-18 13:15:41 -0700717 """
718 Create a composite member iff value contains only members.
719 """
Ethan Furmanee47e5c2016-08-31 00:12:15 -0700720 pseudo_member = cls._value2member_map_.get(value, None)
721 if pseudo_member is None:
Ethan Furman3515dcc2016-09-18 13:15:41 -0700722 # verify all bits are accounted for
723 _, extra_flags = _decompose(cls, value)
724 if extra_flags:
Walter Dörwald323842c2019-07-18 20:37:13 +0200725 raise ValueError("%r is not a valid %s" % (value, cls.__qualname__))
Ethan Furman3515dcc2016-09-18 13:15:41 -0700726 # construct a singleton enum pseudo-member
Ethan Furmanee47e5c2016-08-31 00:12:15 -0700727 pseudo_member = object.__new__(cls)
728 pseudo_member._name_ = None
729 pseudo_member._value_ = value
Ethan Furman28cf6632017-01-24 12:12:06 -0800730 # use setdefault in case another thread already created a composite
731 # with this value
732 pseudo_member = cls._value2member_map_.setdefault(value, pseudo_member)
Ethan Furmanee47e5c2016-08-31 00:12:15 -0700733 return pseudo_member
734
Ethan Furmanee47e5c2016-08-31 00:12:15 -0700735 def __contains__(self, other):
736 if not isinstance(other, self.__class__):
Rahul Jha94306522018-09-10 23:51:04 +0530737 raise TypeError(
738 "unsupported operand type(s) for 'in': '%s' and '%s'" % (
739 type(other).__qualname__, self.__class__.__qualname__))
Ethan Furmanee47e5c2016-08-31 00:12:15 -0700740 return other._value_ & self._value_ == other._value_
741
Ethan Furmanee47e5c2016-08-31 00:12:15 -0700742 def __repr__(self):
743 cls = self.__class__
744 if self._name_ is not None:
745 return '<%s.%s: %r>' % (cls.__name__, self._name_, self._value_)
Ethan Furman3515dcc2016-09-18 13:15:41 -0700746 members, uncovered = _decompose(cls, self._value_)
Ethan Furman27682d22016-09-04 11:39:01 -0700747 return '<%s.%s: %r>' % (
748 cls.__name__,
749 '|'.join([str(m._name_ or m._value_) for m in members]),
750 self._value_,
751 )
Ethan Furmanee47e5c2016-08-31 00:12:15 -0700752
753 def __str__(self):
754 cls = self.__class__
755 if self._name_ is not None:
756 return '%s.%s' % (cls.__name__, self._name_)
Ethan Furman3515dcc2016-09-18 13:15:41 -0700757 members, uncovered = _decompose(cls, self._value_)
Ethan Furmanee47e5c2016-08-31 00:12:15 -0700758 if len(members) == 1 and members[0]._name_ is None:
759 return '%s.%r' % (cls.__name__, members[0]._value_)
760 else:
761 return '%s.%s' % (
762 cls.__name__,
763 '|'.join([str(m._name_ or m._value_) for m in members]),
764 )
765
Ethan Furman25d94bb2016-09-02 16:32:32 -0700766 def __bool__(self):
767 return bool(self._value_)
768
Ethan Furmanee47e5c2016-08-31 00:12:15 -0700769 def __or__(self, other):
770 if not isinstance(other, self.__class__):
771 return NotImplemented
772 return self.__class__(self._value_ | other._value_)
773
774 def __and__(self, other):
775 if not isinstance(other, self.__class__):
776 return NotImplemented
777 return self.__class__(self._value_ & other._value_)
778
779 def __xor__(self, other):
780 if not isinstance(other, self.__class__):
781 return NotImplemented
782 return self.__class__(self._value_ ^ other._value_)
783
784 def __invert__(self):
Ethan Furman3515dcc2016-09-18 13:15:41 -0700785 members, uncovered = _decompose(self.__class__, self._value_)
Serhiy Storchaka81108372017-09-26 00:55:55 +0300786 inverted = self.__class__(0)
787 for m in self.__class__:
788 if m not in members and not (m._value_ & self._value_):
789 inverted = inverted | m
Ethan Furmanee47e5c2016-08-31 00:12:15 -0700790 return self.__class__(inverted)
791
792
Ethan Furman65a5a472016-09-01 23:55:19 -0700793class IntFlag(int, Flag):
Ethan Furmanee47e5c2016-08-31 00:12:15 -0700794 """Support for integer-based Flags"""
795
796 @classmethod
Ethan Furman3515dcc2016-09-18 13:15:41 -0700797 def _missing_(cls, value):
798 if not isinstance(value, int):
Walter Dörwald323842c2019-07-18 20:37:13 +0200799 raise ValueError("%r is not a valid %s" % (value, cls.__qualname__))
Ethan Furman3515dcc2016-09-18 13:15:41 -0700800 new_member = cls._create_pseudo_member_(value)
801 return new_member
802
803 @classmethod
Ethan Furmanee47e5c2016-08-31 00:12:15 -0700804 def _create_pseudo_member_(cls, value):
805 pseudo_member = cls._value2member_map_.get(value, None)
806 if pseudo_member is None:
Ethan Furman3515dcc2016-09-18 13:15:41 -0700807 need_to_create = [value]
808 # get unaccounted for bits
809 _, extra_flags = _decompose(cls, value)
810 # timer = 10
811 while extra_flags:
812 # timer -= 1
813 bit = _high_bit(extra_flags)
814 flag_value = 2 ** bit
815 if (flag_value not in cls._value2member_map_ and
816 flag_value not in need_to_create
817 ):
818 need_to_create.append(flag_value)
819 if extra_flags == -flag_value:
820 extra_flags = 0
821 else:
822 extra_flags ^= flag_value
823 for value in reversed(need_to_create):
824 # construct singleton pseudo-members
825 pseudo_member = int.__new__(cls, value)
826 pseudo_member._name_ = None
827 pseudo_member._value_ = value
Ethan Furman28cf6632017-01-24 12:12:06 -0800828 # use setdefault in case another thread already created a composite
829 # with this value
830 pseudo_member = cls._value2member_map_.setdefault(value, pseudo_member)
Ethan Furmanee47e5c2016-08-31 00:12:15 -0700831 return pseudo_member
832
Ethan Furmanee47e5c2016-08-31 00:12:15 -0700833 def __or__(self, other):
834 if not isinstance(other, (self.__class__, int)):
835 return NotImplemented
Ethan Furman3515dcc2016-09-18 13:15:41 -0700836 result = self.__class__(self._value_ | self.__class__(other)._value_)
837 return result
Ethan Furmanee47e5c2016-08-31 00:12:15 -0700838
839 def __and__(self, other):
840 if not isinstance(other, (self.__class__, int)):
841 return NotImplemented
842 return self.__class__(self._value_ & self.__class__(other)._value_)
843
844 def __xor__(self, other):
845 if not isinstance(other, (self.__class__, int)):
846 return NotImplemented
847 return self.__class__(self._value_ ^ self.__class__(other)._value_)
848
849 __ror__ = __or__
850 __rand__ = __and__
851 __rxor__ = __xor__
852
853 def __invert__(self):
Ethan Furman3515dcc2016-09-18 13:15:41 -0700854 result = self.__class__(~self._value_)
855 return result
Ethan Furmanee47e5c2016-08-31 00:12:15 -0700856
857
858def _high_bit(value):
Ethan Furman04439532016-09-02 15:50:21 -0700859 """returns index of highest bit, or -1 if value is zero or negative"""
Ethan Furman3515dcc2016-09-18 13:15:41 -0700860 return value.bit_length() - 1
Ethan Furmanee47e5c2016-08-31 00:12:15 -0700861
Ethan Furmanf24bb352013-07-18 17:05:39 -0700862def unique(enumeration):
863 """Class decorator for enumerations ensuring unique member values."""
864 duplicates = []
865 for name, member in enumeration.__members__.items():
866 if name != member.name:
867 duplicates.append((name, member.name))
868 if duplicates:
869 alias_details = ', '.join(
870 ["%s -> %s" % (alias, name) for (alias, name) in duplicates])
871 raise ValueError('duplicate values found in %r: %s' %
872 (enumeration, alias_details))
873 return enumeration
Ethan Furman3515dcc2016-09-18 13:15:41 -0700874
875def _decompose(flag, value):
876 """Extract all members from the value."""
877 # _decompose is only called if the value is not named
878 not_covered = value
879 negative = value < 0
Ethan Furman3515dcc2016-09-18 13:15:41 -0700880 members = []
HongWeipeng0b41a922019-11-27 06:36:02 +0800881 for member in flag:
882 member_value = member.value
Ethan Furman3515dcc2016-09-18 13:15:41 -0700883 if member_value and member_value & value == member_value:
884 members.append(member)
885 not_covered &= ~member_value
HongWeipeng0b41a922019-11-27 06:36:02 +0800886 if not negative:
887 tmp = not_covered
888 while tmp:
889 flag_value = 2 ** _high_bit(tmp)
890 if flag_value in flag._value2member_map_:
891 members.append(flag._value2member_map_[flag_value])
892 not_covered &= ~flag_value
893 tmp &= ~flag_value
Ethan Furman3515dcc2016-09-18 13:15:41 -0700894 if not members and value in flag._value2member_map_:
895 members.append(flag._value2member_map_[value])
896 members.sort(key=lambda m: m._value_, reverse=True)
897 if len(members) > 1 and members[0].value == value:
898 # we have the breakdown, don't need the value member itself
899 members.pop(0)
900 return members, not_covered