blob: 5e0088ee89fee3b2aa0f25e82bad3398b86f2c3b [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 ):
Zackery Spytz2ec67522020-09-13 14:27:51 -060079 raise ValueError(f'_sunder_ names, such as "{key}", are '
80 'reserved for future Enum use')
Ethan Furmanc16595e2016-09-10 23:36:59 -070081 if key == '_generate_next_value_':
Ethan Onstottd9a43e22020-04-28 13:20:55 -040082 # check if members already defined as auto()
83 if self._auto_called:
84 raise TypeError("_generate_next_value_ must be defined before members")
Ethan Furmanc16595e2016-09-10 23:36:59 -070085 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 Onstottd9a43e22020-04-28 13:20:55 -0400108 self._auto_called = True
Ethan Furman3515dcc2016-09-18 13:15:41 -0700109 if value.value == _auto_null:
110 value.value = self._generate_next_value(key, 1, len(self._member_names), self._last_values[:])
111 value = value.value
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700112 self._member_names.append(key)
Ethan Furmanc16595e2016-09-10 23:36:59 -0700113 self._last_values.append(value)
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700114 super().__setitem__(key, value)
115
116
Ezio Melotti9a3777e2013-08-17 15:53:55 +0300117# Dummy value for Enum as EnumMeta explicitly checks for it, but of course
118# until EnumMeta finishes running the first time the Enum class doesn't exist.
119# This is also why there are checks in EnumMeta like `if Enum is not None`
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700120Enum = None
121
Ethan Furman332dbc72016-08-20 00:00:52 -0700122
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700123class EnumMeta(type):
124 """Metaclass for Enum"""
125 @classmethod
Ethan Furman332dbc72016-08-20 00:00:52 -0700126 def __prepare__(metacls, cls, bases):
Ethan Furmanee47e5c2016-08-31 00:12:15 -0700127 # create the namespace dict
128 enum_dict = _EnumDict()
129 # inherit previous flags and _generate_next_value_ function
130 member_type, first_enum = metacls._get_mixins_(bases)
131 if first_enum is not None:
132 enum_dict['_generate_next_value_'] = getattr(first_enum, '_generate_next_value_', None)
133 return enum_dict
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700134
Ethan Furman65a5a472016-09-01 23:55:19 -0700135 def __new__(metacls, cls, bases, classdict):
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700136 # an Enum class is final once enumeration items have been defined; it
137 # cannot be mixed with other types (int, float, etc.) if it has an
138 # inherited __new__ unless a new __new__ is defined (or the resulting
139 # class will fail).
Ethan Furmana4b1bb42018-01-22 07:56:37 -0800140 #
141 # remove any keys listed in _ignore_
142 classdict.setdefault('_ignore_', []).append('_ignore_')
143 ignore = classdict['_ignore_']
144 for key in ignore:
145 classdict.pop(key, None)
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700146 member_type, first_enum = metacls._get_mixins_(bases)
147 __new__, save_new, use_args = metacls._find_new_(classdict, member_type,
148 first_enum)
149
150 # save enum items into separate mapping so they don't get baked into
151 # the new class
Ethan Furmanee47e5c2016-08-31 00:12:15 -0700152 enum_members = {k: classdict[k] for k in classdict._member_names}
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700153 for name in classdict._member_names:
154 del classdict[name]
155
Ethan Furmane8e61272016-08-20 07:19:31 -0700156 # adjust the sunders
157 _order_ = classdict.pop('_order_', None)
158
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700159 # check for illegal enum names (any others?)
Brennan D Baraban8b914d22019-03-03 14:09:11 -0800160 invalid_names = set(enum_members) & {'mro', ''}
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700161 if invalid_names:
162 raise ValueError('Invalid enum member name: {0}'.format(
163 ','.join(invalid_names)))
164
Ethan Furman48a724f2015-04-11 23:23:06 -0700165 # create a default docstring if one has not been provided
166 if '__doc__' not in classdict:
167 classdict['__doc__'] = 'An enumeration.'
168
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700169 # create our new Enum type
170 enum_class = super().__new__(metacls, cls, bases, classdict)
Ethan Furman520ad572013-07-19 19:47:21 -0700171 enum_class._member_names_ = [] # names in definition order
INADA Naokie57f91a2018-06-19 01:14:26 +0900172 enum_class._member_map_ = {} # name->value map
Ethan Furman5e5a8232013-08-04 08:42:23 -0700173 enum_class._member_type_ = member_type
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700174
orlnub1230fb9fad2018-09-12 20:28:53 +0300175 # save DynamicClassAttribute attributes from super classes so we know
176 # if we can take the shortcut of storing members in the class dict
177 dynamic_attributes = {k for c in enum_class.mro()
178 for k, v in c.__dict__.items()
179 if isinstance(v, DynamicClassAttribute)}
Ethan Furman354ecf12015-03-11 08:43:12 -0700180
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700181 # Reverse value->name map for hashable values.
Ethan Furman520ad572013-07-19 19:47:21 -0700182 enum_class._value2member_map_ = {}
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700183
Ethan Furman2da95042014-03-03 12:42:52 -0800184 # If a custom type is mixed into the Enum, and it does not know how
185 # to pickle itself, pickle.dumps will succeed but pickle.loads will
186 # fail. Rather than have the error show up later and possibly far
187 # from the source, sabotage the pickle protocol for this class so
188 # that pickle.dumps also fails.
189 #
190 # However, if the new class implements its own __reduce_ex__, do not
191 # sabotage -- it's on them to make sure it works correctly. We use
192 # __reduce_ex__ instead of any of the others as it is preferred by
193 # pickle over __reduce__, and it handles all pickle protocols.
194 if '__reduce_ex__' not in classdict:
Ethan Furmandc870522014-02-18 12:37:12 -0800195 if member_type is not object:
196 methods = ('__getnewargs_ex__', '__getnewargs__',
197 '__reduce_ex__', '__reduce__')
Ethan Furman2da95042014-03-03 12:42:52 -0800198 if not any(m in member_type.__dict__ for m in methods):
Ethan Furmandc870522014-02-18 12:37:12 -0800199 _make_class_unpicklable(enum_class)
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700200
201 # instantiate them, checking for duplicates as we go
202 # we instantiate first instead of checking for duplicates first in case
203 # a custom __new__ is doing something funky with the values -- such as
204 # auto-numbering ;)
205 for member_name in classdict._member_names:
Ethan Furmanee47e5c2016-08-31 00:12:15 -0700206 value = enum_members[member_name]
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700207 if not isinstance(value, tuple):
208 args = (value, )
209 else:
210 args = value
211 if member_type is tuple: # special case for tuple enums
212 args = (args, ) # wrap it one more time
213 if not use_args:
214 enum_member = __new__(enum_class)
Ethan Furmanb41803e2013-07-25 13:50:45 -0700215 if not hasattr(enum_member, '_value_'):
216 enum_member._value_ = value
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700217 else:
218 enum_member = __new__(enum_class, *args)
Ethan Furmanb41803e2013-07-25 13:50:45 -0700219 if not hasattr(enum_member, '_value_'):
Ethan Furmanee47e5c2016-08-31 00:12:15 -0700220 if member_type is object:
221 enum_member._value_ = value
222 else:
223 enum_member._value_ = member_type(*args)
Ethan Furman520ad572013-07-19 19:47:21 -0700224 value = enum_member._value_
Ethan Furman520ad572013-07-19 19:47:21 -0700225 enum_member._name_ = member_name
Ethan Furmanc850f342013-09-15 16:59:35 -0700226 enum_member.__objclass__ = enum_class
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700227 enum_member.__init__(*args)
228 # If another member with the same value was already defined, the
229 # new member becomes an alias to the existing one.
Ethan Furman520ad572013-07-19 19:47:21 -0700230 for name, canonical_member in enum_class._member_map_.items():
Ethan Furman0081f232014-09-16 17:31:23 -0700231 if canonical_member._value_ == enum_member._value_:
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700232 enum_member = canonical_member
233 break
234 else:
235 # Aliases don't appear in member names (only in __members__).
Ethan Furman520ad572013-07-19 19:47:21 -0700236 enum_class._member_names_.append(member_name)
Ethan Furman354ecf12015-03-11 08:43:12 -0700237 # performance boost for any member that would not shadow
238 # a DynamicClassAttribute
orlnub1230fb9fad2018-09-12 20:28:53 +0300239 if member_name not in dynamic_attributes:
Ethan Furman354ecf12015-03-11 08:43:12 -0700240 setattr(enum_class, member_name, enum_member)
241 # now add to _member_map_
Ethan Furman520ad572013-07-19 19:47:21 -0700242 enum_class._member_map_[member_name] = enum_member
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700243 try:
244 # This may fail if value is not hashable. We can't add the value
245 # to the map, and by-value lookups for this value will be
246 # linear.
Ethan Furman520ad572013-07-19 19:47:21 -0700247 enum_class._value2member_map_[value] = enum_member
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700248 except TypeError:
249 pass
250
251 # double check that repr and friends are not the mixin's or various
252 # things break (such as pickle)
Ethan Furmandc870522014-02-18 12:37:12 -0800253 for name in ('__repr__', '__str__', '__format__', '__reduce_ex__'):
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700254 class_method = getattr(enum_class, name)
255 obj_method = getattr(member_type, name, None)
256 enum_method = getattr(first_enum, name, None)
257 if obj_method is not None and obj_method is class_method:
258 setattr(enum_class, name, enum_method)
259
260 # replace any other __new__ with our own (as long as Enum is not None,
261 # anyway) -- again, this is to support pickle
262 if Enum is not None:
263 # if the user defined their own __new__, save it before it gets
264 # clobbered in case they subclass later
265 if save_new:
266 enum_class.__new_member__ = __new__
267 enum_class.__new__ = Enum.__new__
Ethan Furmane8e61272016-08-20 07:19:31 -0700268
269 # py3 support for definition order (helps keep py2/py3 code in sync)
270 if _order_ is not None:
271 if isinstance(_order_, str):
272 _order_ = _order_.replace(',', ' ').split()
273 if _order_ != enum_class._member_names_:
274 raise TypeError('member order does not match _order_')
275
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700276 return enum_class
277
Ethan Furman5de67b12016-04-13 23:52:09 -0700278 def __bool__(self):
279 """
280 classes/types should always be True.
281 """
282 return True
283
Ethan Furmand9925a12014-09-16 20:35:55 -0700284 def __call__(cls, value, names=None, *, module=None, qualname=None, type=None, start=1):
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700285 """Either returns an existing member, or creates a new enum class.
286
287 This method is used both when an enum class is given a value to match
288 to an enumeration member (i.e. Color(3)) and for the functional API
Ethan Furman23bb6f42016-11-21 09:22:05 -0800289 (i.e. Color = Enum('Color', names='RED GREEN BLUE')).
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700290
Ethan Furman2da95042014-03-03 12:42:52 -0800291 When used for the functional API:
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700292
Ethan Furman2da95042014-03-03 12:42:52 -0800293 `value` will be the name of the new class.
294
295 `names` should be either a string of white-space/comma delimited names
Ethan Furmand9925a12014-09-16 20:35:55 -0700296 (values will start at `start`), or an iterator/mapping of name, value pairs.
Ethan Furman2da95042014-03-03 12:42:52 -0800297
298 `module` should be set to the module this class is being created in;
299 if it is not set, an attempt to find that module will be made, but if
300 it fails the class will not be picklable.
301
302 `qualname` should be set to the actual location this class can be found
303 at in its module; by default it is set to the global scope. If this is
304 not correct, unpickling will fail in some circumstances.
305
306 `type`, if set, will be mixed in as the first base class.
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700307
308 """
309 if names is None: # simple value lookup
310 return cls.__new__(cls, value)
311 # otherwise, functional API: we're creating a new Enum type
Ethan Furmand9925a12014-09-16 20:35:55 -0700312 return cls._create_(value, names, module=module, qualname=qualname, type=type, start=start)
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700313
314 def __contains__(cls, member):
Rahul Jha94306522018-09-10 23:51:04 +0530315 if not isinstance(member, Enum):
316 raise TypeError(
317 "unsupported operand type(s) for 'in': '%s' and '%s'" % (
318 type(member).__qualname__, cls.__class__.__qualname__))
Ethan Furman0081f232014-09-16 17:31:23 -0700319 return isinstance(member, cls) and member._name_ in cls._member_map_
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700320
Ethan Furman64a99722013-09-22 16:18:19 -0700321 def __delattr__(cls, attr):
322 # nicer error message when someone tries to delete an attribute
323 # (see issue19025).
324 if attr in cls._member_map_:
325 raise AttributeError(
326 "%s: cannot delete Enum member." % cls.__name__)
327 super().__delattr__(attr)
328
Ethan Furman388a3922013-08-12 06:51:41 -0700329 def __dir__(self):
Ethan Furman64a99722013-09-22 16:18:19 -0700330 return (['__class__', '__doc__', '__members__', '__module__'] +
331 self._member_names_)
Ethan Furman388a3922013-08-12 06:51:41 -0700332
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700333 def __getattr__(cls, name):
334 """Return the enum member matching `name`
335
336 We use __getattr__ instead of descriptors or inserting into the enum
337 class' __dict__ in order to support `name` and `value` being both
338 properties for enum members (which live in the class' __dict__) and
339 enum members themselves.
340
341 """
342 if _is_dunder(name):
343 raise AttributeError(name)
344 try:
Ethan Furman520ad572013-07-19 19:47:21 -0700345 return cls._member_map_[name]
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700346 except KeyError:
347 raise AttributeError(name) from None
348
349 def __getitem__(cls, name):
Ethan Furman520ad572013-07-19 19:47:21 -0700350 return cls._member_map_[name]
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700351
352 def __iter__(cls):
Ethan Furman520ad572013-07-19 19:47:21 -0700353 return (cls._member_map_[name] for name in cls._member_names_)
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700354
355 def __len__(cls):
Ethan Furman520ad572013-07-19 19:47:21 -0700356 return len(cls._member_names_)
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700357
Ethan Furman2131a4a2013-09-14 18:11:24 -0700358 @property
359 def __members__(cls):
360 """Returns a mapping of member name->value.
361
362 This mapping lists all enum members, including aliases. Note that this
363 is a read-only view of the internal mapping.
364
365 """
366 return MappingProxyType(cls._member_map_)
367
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700368 def __repr__(cls):
369 return "<enum %r>" % cls.__name__
370
Ethan Furman2131a4a2013-09-14 18:11:24 -0700371 def __reversed__(cls):
372 return (cls._member_map_[name] for name in reversed(cls._member_names_))
373
Ethan Furmanf203f2d2013-09-06 07:16:48 -0700374 def __setattr__(cls, name, value):
375 """Block attempts to reassign Enum members.
376
377 A simple assignment to the class namespace only changes one of the
378 several possible ways to get an Enum member from the Enum class,
379 resulting in an inconsistent Enumeration.
380
381 """
382 member_map = cls.__dict__.get('_member_map_', {})
383 if name in member_map:
384 raise AttributeError('Cannot reassign members.')
385 super().__setattr__(name, value)
386
anentropicb8e21f12018-04-16 04:40:35 +0100387 def _create_(cls, class_name, names, *, module=None, qualname=None, type=None, start=1):
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700388 """Convenience method to create a new Enum class.
389
390 `names` can be:
391
392 * A string containing member names, separated either with spaces or
Ethan Furmand9925a12014-09-16 20:35:55 -0700393 commas. Values are incremented by 1 from `start`.
394 * An iterable of member names. Values are incremented by 1 from `start`.
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700395 * An iterable of (member name, value) pairs.
Ethan Furmand9925a12014-09-16 20:35:55 -0700396 * A mapping of member name -> value pairs.
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700397
398 """
399 metacls = cls.__class__
400 bases = (cls, ) if type is None else (type, cls)
Ethan Furmanee47e5c2016-08-31 00:12:15 -0700401 _, first_enum = cls._get_mixins_(bases)
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700402 classdict = metacls.__prepare__(class_name, bases)
403
404 # special processing needed for names?
405 if isinstance(names, str):
406 names = names.replace(',', ' ').split()
Dong-hee Nadcc8ce42017-06-22 01:52:32 +0900407 if isinstance(names, (tuple, list)) and names and isinstance(names[0], str):
Ethan Furmanee47e5c2016-08-31 00:12:15 -0700408 original_names, names = names, []
Ethan Furmanc16595e2016-09-10 23:36:59 -0700409 last_values = []
Ethan Furmanee47e5c2016-08-31 00:12:15 -0700410 for count, name in enumerate(original_names):
Ethan Furmanc16595e2016-09-10 23:36:59 -0700411 value = first_enum._generate_next_value_(name, start, count, last_values[:])
412 last_values.append(value)
413 names.append((name, value))
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700414
415 # Here, names is either an iterable of (name, value) or a mapping.
416 for item in names:
417 if isinstance(item, str):
418 member_name, member_value = item, names[item]
419 else:
420 member_name, member_value = item
421 classdict[member_name] = member_value
422 enum_class = metacls.__new__(metacls, class_name, bases, classdict)
423
424 # TODO: replace the frame hack if a blessed way to know the calling
425 # module is ever developed
426 if module is None:
427 try:
428 module = sys._getframe(2).f_globals['__name__']
Pablo Galindo293dd232019-11-19 21:34:03 +0000429 except (AttributeError, ValueError, KeyError):
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700430 pass
431 if module is None:
432 _make_class_unpicklable(enum_class)
433 else:
434 enum_class.__module__ = module
Ethan Furmanca1b7942014-02-08 11:36:27 -0800435 if qualname is not None:
436 enum_class.__qualname__ = qualname
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700437
438 return enum_class
439
orlnub1230fb9fad2018-09-12 20:28:53 +0300440 def _convert_(cls, name, module, filter, source=None):
441 """
442 Create a new Enum subclass that replaces a collection of global constants
443 """
444 # convert all constants from source (or module) that pass filter() to
445 # a new Enum called name, and export the enum and its members back to
446 # module;
447 # also, replace the __reduce_ex__ method so unpickling works in
448 # previous Python versions
449 module_globals = vars(sys.modules[module])
450 if source:
451 source = vars(source)
452 else:
453 source = module_globals
454 # _value2member_map_ is populated in the same order every time
455 # for a consistent reverse mapping of number to name when there
456 # are multiple names for the same number.
457 members = [
458 (name, value)
459 for name, value in source.items()
460 if filter(name)]
461 try:
462 # sort by value
463 members.sort(key=lambda t: (t[1], t[0]))
464 except TypeError:
465 # unless some values aren't comparable, in which case sort by name
466 members.sort(key=lambda t: t[0])
467 cls = cls(name, members, module=module)
468 cls.__reduce_ex__ = _reduce_ex_by_name
469 module_globals.update(cls.__members__)
470 module_globals[name] = cls
471 return cls
472
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700473 @staticmethod
474 def _get_mixins_(bases):
475 """Returns the type for creating enum members, and the first inherited
476 enum class.
477
478 bases: the tuple of bases that was given to __new__
479
480 """
481 if not bases:
482 return object, Enum
483
Ethan Furman5bdab642018-09-21 19:03:09 -0700484 def _find_data_type(bases):
Ethan Furmanbff01f32020-09-15 15:56:26 -0700485 data_types = []
Ethan Furman5bdab642018-09-21 19:03:09 -0700486 for chain in bases:
Ethan Furmanbff01f32020-09-15 15:56:26 -0700487 candidate = None
Ethan Furman5bdab642018-09-21 19:03:09 -0700488 for base in chain.__mro__:
489 if base is object:
490 continue
491 elif '__new__' in base.__dict__:
Ethan Furmancd453852018-10-05 23:29:36 -0700492 if issubclass(base, Enum):
Ethan Furman5bdab642018-09-21 19:03:09 -0700493 continue
Ethan Furmanbff01f32020-09-15 15:56:26 -0700494 data_types.append(candidate or base)
495 break
496 elif not issubclass(base, Enum):
497 candidate = base
498 if len(data_types) > 1:
499 raise TypeError('too many data types: %r' % data_types)
500 elif data_types:
501 return data_types[0]
502 else:
503 return None
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700504
Ethan Furman5bdab642018-09-21 19:03:09 -0700505 # ensure final parent class is an Enum derivative, find any concrete
506 # data type, and check that Enum has no members
507 first_enum = bases[-1]
508 if not issubclass(first_enum, Enum):
509 raise TypeError("new enumerations should be created as "
510 "`EnumName([mixin_type, ...] [data_type,] enum_type)`")
511 member_type = _find_data_type(bases) or object
512 if first_enum._member_names_:
513 raise TypeError("Cannot extend enumerations")
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700514 return member_type, first_enum
515
516 @staticmethod
517 def _find_new_(classdict, member_type, first_enum):
518 """Returns the __new__ to be used for creating the enum members.
519
520 classdict: the class dictionary given to __new__
521 member_type: the data type whose __new__ will be used by default
522 first_enum: enumeration to check for an overriding __new__
523
524 """
525 # now find the correct __new__, checking to see of one was defined
526 # by the user; also check earlier enum classes in case a __new__ was
527 # saved as __new_member__
528 __new__ = classdict.get('__new__', None)
529
530 # should __new__ be saved as __new_member__ later?
531 save_new = __new__ is not None
532
533 if __new__ is None:
534 # check all possibles for __new_member__ before falling back to
535 # __new__
536 for method in ('__new_member__', '__new__'):
537 for possible in (member_type, first_enum):
538 target = getattr(possible, method, None)
539 if target not in {
540 None,
541 None.__new__,
542 object.__new__,
543 Enum.__new__,
544 }:
545 __new__ = target
546 break
547 if __new__ is not None:
548 break
549 else:
550 __new__ = object.__new__
551
552 # if a non-object.__new__ is used then whatever value/tuple was
553 # assigned to the enum member name will be passed to __new__ and to the
554 # new enum member's __init__
555 if __new__ is object.__new__:
556 use_args = False
557 else:
558 use_args = True
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700559 return __new__, save_new, use_args
560
561
562class Enum(metaclass=EnumMeta):
563 """Generic enumeration.
564
565 Derive from this class to define new enumerations.
566
567 """
568 def __new__(cls, value):
569 # all enum instances are actually created during class construction
570 # without calling this method; this method is called by the metaclass'
571 # __call__ (i.e. Color(3) ), and by pickle
572 if type(value) is cls:
Ethan Furman23bb6f42016-11-21 09:22:05 -0800573 # For lookups like Color(Color.RED)
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700574 return value
575 # by-value search for a matching enum member
576 # see if it's in the reverse mapping (for hashable values)
Ethan Furman2aa27322013-07-19 19:35:56 -0700577 try:
Andrew Svetlov34ae04f2018-12-26 20:45:33 +0200578 return cls._value2member_map_[value]
579 except KeyError:
580 # Not found, no need to do long O(n) search
581 pass
Ethan Furman2aa27322013-07-19 19:35:56 -0700582 except TypeError:
583 # not there, now do long search -- O(n) behavior
Ethan Furman520ad572013-07-19 19:47:21 -0700584 for member in cls._member_map_.values():
Ethan Furman0081f232014-09-16 17:31:23 -0700585 if member._value_ == value:
Ethan Furman2aa27322013-07-19 19:35:56 -0700586 return member
Ethan Furmanee47e5c2016-08-31 00:12:15 -0700587 # still not found -- try _missing_ hook
Ethan Furman019f0a02018-09-12 11:43:34 -0700588 try:
589 exc = None
590 result = cls._missing_(value)
591 except Exception as e:
592 exc = e
593 result = None
594 if isinstance(result, cls):
595 return result
596 else:
Walter Dörwald323842c2019-07-18 20:37:13 +0200597 ve_exc = ValueError("%r is not a valid %s" % (value, cls.__qualname__))
Ethan Furman019f0a02018-09-12 11:43:34 -0700598 if result is None and exc is None:
599 raise ve_exc
600 elif exc is None:
601 exc = TypeError(
602 'error in %s._missing_: returned %r instead of None or a valid member'
603 % (cls.__name__, result)
604 )
605 exc.__context__ = ve_exc
606 raise exc
Ethan Furmanee47e5c2016-08-31 00:12:15 -0700607
Ethan Furmanc16595e2016-09-10 23:36:59 -0700608 def _generate_next_value_(name, start, count, last_values):
609 for last_value in reversed(last_values):
610 try:
611 return last_value + 1
612 except TypeError:
613 pass
614 else:
Ethan Furmanee47e5c2016-08-31 00:12:15 -0700615 return start
Ethan Furmanc16595e2016-09-10 23:36:59 -0700616
Ethan Furmanee47e5c2016-08-31 00:12:15 -0700617 @classmethod
618 def _missing_(cls, value):
Walter Dörwald323842c2019-07-18 20:37:13 +0200619 raise ValueError("%r is not a valid %s" % (value, cls.__qualname__))
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700620
621 def __repr__(self):
622 return "<%s.%s: %r>" % (
Ethan Furman520ad572013-07-19 19:47:21 -0700623 self.__class__.__name__, self._name_, self._value_)
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700624
625 def __str__(self):
Ethan Furman520ad572013-07-19 19:47:21 -0700626 return "%s.%s" % (self.__class__.__name__, self._name_)
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700627
Ethan Furman388a3922013-08-12 06:51:41 -0700628 def __dir__(self):
Ethan Furman0ae550b2014-10-14 08:58:32 -0700629 added_behavior = [
630 m
631 for cls in self.__class__.mro()
632 for m in cls.__dict__
Ethan Furman354ecf12015-03-11 08:43:12 -0700633 if m[0] != '_' and m not in self._member_map_
Ethan Furman0ae550b2014-10-14 08:58:32 -0700634 ]
Ethan Furmanec5f8eb2014-10-21 13:40:35 -0700635 return (['__class__', '__doc__', '__module__'] + added_behavior)
Ethan Furman388a3922013-08-12 06:51:41 -0700636
Ethan Furmanec15a822013-08-31 19:17:41 -0700637 def __format__(self, format_spec):
638 # mixed-in Enums should use the mixed-in type's __format__, otherwise
639 # we can get strange results with the Enum name showing up instead of
640 # the value
641
thatneat2f19e822019-07-04 11:28:37 -0700642 # pure Enum branch, or branch with __str__ explicitly overridden
643 str_overridden = type(self).__str__ != Enum.__str__
644 if self._member_type_ is object or str_overridden:
Ethan Furmanec15a822013-08-31 19:17:41 -0700645 cls = str
646 val = str(self)
647 # mix-in branch
648 else:
649 cls = self._member_type_
Ethan Furman0081f232014-09-16 17:31:23 -0700650 val = self._value_
Ethan Furmanec15a822013-08-31 19:17:41 -0700651 return cls.__format__(val, format_spec)
652
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700653 def __hash__(self):
Ethan Furman520ad572013-07-19 19:47:21 -0700654 return hash(self._name_)
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700655
Ethan Furmanca1b7942014-02-08 11:36:27 -0800656 def __reduce_ex__(self, proto):
Ethan Furmandc870522014-02-18 12:37:12 -0800657 return self.__class__, (self._value_, )
Ethan Furmanca1b7942014-02-08 11:36:27 -0800658
Ethan Furman33918c12013-09-27 23:02:02 -0700659 # DynamicClassAttribute is used to provide access to the `name` and
660 # `value` properties of enum members while keeping some measure of
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700661 # protection from modification, while still allowing for an enumeration
662 # to have members named `name` and `value`. This works because enumeration
663 # members are not set directly on the enum class -- __getattr__ is
664 # used to look them up.
665
Ethan Furmane03ea372013-09-25 07:14:41 -0700666 @DynamicClassAttribute
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700667 def name(self):
Ethan Furmanc850f342013-09-15 16:59:35 -0700668 """The name of the Enum member."""
Ethan Furman520ad572013-07-19 19:47:21 -0700669 return self._name_
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700670
Ethan Furmane03ea372013-09-25 07:14:41 -0700671 @DynamicClassAttribute
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700672 def value(self):
Ethan Furmanc850f342013-09-15 16:59:35 -0700673 """The value of the Enum member."""
Ethan Furman520ad572013-07-19 19:47:21 -0700674 return self._value_
Ethan Furman6b3d64a2013-06-14 16:55:46 -0700675
676
677class IntEnum(int, Enum):
678 """Enum where members are also (and must be) ints"""
Ethan Furmanf24bb352013-07-18 17:05:39 -0700679
680
Ethan Furman24e837f2015-03-18 17:27:57 -0700681def _reduce_ex_by_name(self, proto):
682 return self.name
683
Ethan Furman65a5a472016-09-01 23:55:19 -0700684class Flag(Enum):
Ethan Furmanee47e5c2016-08-31 00:12:15 -0700685 """Support for flags"""
Ethan Furmanc16595e2016-09-10 23:36:59 -0700686
687 def _generate_next_value_(name, start, count, last_values):
Ethan Furmanee47e5c2016-08-31 00:12:15 -0700688 """
689 Generate the next value when not given.
690
691 name: the name of the member
HongWeipengbb16fb22019-09-21 13:22:54 +0800692 start: the initial start value or None
Ethan Furmanee47e5c2016-08-31 00:12:15 -0700693 count: the number of existing members
694 last_value: the last value assigned or None
695 """
696 if not count:
697 return start if start is not None else 1
Ethan Furmanc16595e2016-09-10 23:36:59 -0700698 for last_value in reversed(last_values):
699 try:
700 high_bit = _high_bit(last_value)
701 break
Ethan Furman3515dcc2016-09-18 13:15:41 -0700702 except Exception:
Ethan Furmanc16595e2016-09-10 23:36:59 -0700703 raise TypeError('Invalid Flag value: %r' % last_value) from None
Ethan Furmanee47e5c2016-08-31 00:12:15 -0700704 return 2 ** (high_bit+1)
705
706 @classmethod
707 def _missing_(cls, value):
708 original_value = value
709 if value < 0:
710 value = ~value
711 possible_member = cls._create_pseudo_member_(value)
Ethan Furmanee47e5c2016-08-31 00:12:15 -0700712 if original_value < 0:
713 possible_member = ~possible_member
714 return possible_member
715
716 @classmethod
717 def _create_pseudo_member_(cls, value):
Ethan Furman3515dcc2016-09-18 13:15:41 -0700718 """
719 Create a composite member iff value contains only members.
720 """
Ethan Furmanee47e5c2016-08-31 00:12:15 -0700721 pseudo_member = cls._value2member_map_.get(value, None)
722 if pseudo_member is None:
Ethan Furman3515dcc2016-09-18 13:15:41 -0700723 # verify all bits are accounted for
724 _, extra_flags = _decompose(cls, value)
725 if extra_flags:
Walter Dörwald323842c2019-07-18 20:37:13 +0200726 raise ValueError("%r is not a valid %s" % (value, cls.__qualname__))
Ethan Furman3515dcc2016-09-18 13:15:41 -0700727 # construct a singleton enum pseudo-member
Ethan Furmanee47e5c2016-08-31 00:12:15 -0700728 pseudo_member = object.__new__(cls)
729 pseudo_member._name_ = None
730 pseudo_member._value_ = value
Ethan Furman28cf6632017-01-24 12:12:06 -0800731 # use setdefault in case another thread already created a composite
732 # with this value
733 pseudo_member = cls._value2member_map_.setdefault(value, pseudo_member)
Ethan Furmanee47e5c2016-08-31 00:12:15 -0700734 return pseudo_member
735
Ethan Furmanee47e5c2016-08-31 00:12:15 -0700736 def __contains__(self, other):
737 if not isinstance(other, self.__class__):
Rahul Jha94306522018-09-10 23:51:04 +0530738 raise TypeError(
739 "unsupported operand type(s) for 'in': '%s' and '%s'" % (
740 type(other).__qualname__, self.__class__.__qualname__))
Ethan Furmanee47e5c2016-08-31 00:12:15 -0700741 return other._value_ & self._value_ == other._value_
742
Ethan Furmanee47e5c2016-08-31 00:12:15 -0700743 def __repr__(self):
744 cls = self.__class__
745 if self._name_ is not None:
746 return '<%s.%s: %r>' % (cls.__name__, self._name_, self._value_)
Ethan Furman3515dcc2016-09-18 13:15:41 -0700747 members, uncovered = _decompose(cls, self._value_)
Ethan Furman27682d22016-09-04 11:39:01 -0700748 return '<%s.%s: %r>' % (
749 cls.__name__,
750 '|'.join([str(m._name_ or m._value_) for m in members]),
751 self._value_,
752 )
Ethan Furmanee47e5c2016-08-31 00:12:15 -0700753
754 def __str__(self):
755 cls = self.__class__
756 if self._name_ is not None:
757 return '%s.%s' % (cls.__name__, self._name_)
Ethan Furman3515dcc2016-09-18 13:15:41 -0700758 members, uncovered = _decompose(cls, self._value_)
Ethan Furmanee47e5c2016-08-31 00:12:15 -0700759 if len(members) == 1 and members[0]._name_ is None:
760 return '%s.%r' % (cls.__name__, members[0]._value_)
761 else:
762 return '%s.%s' % (
763 cls.__name__,
764 '|'.join([str(m._name_ or m._value_) for m in members]),
765 )
766
Ethan Furman25d94bb2016-09-02 16:32:32 -0700767 def __bool__(self):
768 return bool(self._value_)
769
Ethan Furmanee47e5c2016-08-31 00:12:15 -0700770 def __or__(self, other):
771 if not isinstance(other, self.__class__):
772 return NotImplemented
773 return self.__class__(self._value_ | other._value_)
774
775 def __and__(self, other):
776 if not isinstance(other, self.__class__):
777 return NotImplemented
778 return self.__class__(self._value_ & other._value_)
779
780 def __xor__(self, other):
781 if not isinstance(other, self.__class__):
782 return NotImplemented
783 return self.__class__(self._value_ ^ other._value_)
784
785 def __invert__(self):
Ethan Furman3515dcc2016-09-18 13:15:41 -0700786 members, uncovered = _decompose(self.__class__, self._value_)
Serhiy Storchaka81108372017-09-26 00:55:55 +0300787 inverted = self.__class__(0)
788 for m in self.__class__:
789 if m not in members and not (m._value_ & self._value_):
790 inverted = inverted | m
Ethan Furmanee47e5c2016-08-31 00:12:15 -0700791 return self.__class__(inverted)
792
793
Ethan Furman65a5a472016-09-01 23:55:19 -0700794class IntFlag(int, Flag):
Ethan Furmanee47e5c2016-08-31 00:12:15 -0700795 """Support for integer-based Flags"""
796
797 @classmethod
Ethan Furman3515dcc2016-09-18 13:15:41 -0700798 def _missing_(cls, value):
799 if not isinstance(value, int):
Walter Dörwald323842c2019-07-18 20:37:13 +0200800 raise ValueError("%r is not a valid %s" % (value, cls.__qualname__))
Ethan Furman3515dcc2016-09-18 13:15:41 -0700801 new_member = cls._create_pseudo_member_(value)
802 return new_member
803
804 @classmethod
Ethan Furmanee47e5c2016-08-31 00:12:15 -0700805 def _create_pseudo_member_(cls, value):
806 pseudo_member = cls._value2member_map_.get(value, None)
807 if pseudo_member is None:
Ethan Furman3515dcc2016-09-18 13:15:41 -0700808 need_to_create = [value]
809 # get unaccounted for bits
810 _, extra_flags = _decompose(cls, value)
811 # timer = 10
812 while extra_flags:
813 # timer -= 1
814 bit = _high_bit(extra_flags)
815 flag_value = 2 ** bit
816 if (flag_value not in cls._value2member_map_ and
817 flag_value not in need_to_create
818 ):
819 need_to_create.append(flag_value)
820 if extra_flags == -flag_value:
821 extra_flags = 0
822 else:
823 extra_flags ^= flag_value
824 for value in reversed(need_to_create):
825 # construct singleton pseudo-members
826 pseudo_member = int.__new__(cls, value)
827 pseudo_member._name_ = None
828 pseudo_member._value_ = value
Ethan Furman28cf6632017-01-24 12:12:06 -0800829 # use setdefault in case another thread already created a composite
830 # with this value
831 pseudo_member = cls._value2member_map_.setdefault(value, pseudo_member)
Ethan Furmanee47e5c2016-08-31 00:12:15 -0700832 return pseudo_member
833
Ethan Furmanee47e5c2016-08-31 00:12:15 -0700834 def __or__(self, other):
835 if not isinstance(other, (self.__class__, int)):
836 return NotImplemented
Ethan Furman3515dcc2016-09-18 13:15:41 -0700837 result = self.__class__(self._value_ | self.__class__(other)._value_)
838 return result
Ethan Furmanee47e5c2016-08-31 00:12:15 -0700839
840 def __and__(self, other):
841 if not isinstance(other, (self.__class__, int)):
842 return NotImplemented
843 return self.__class__(self._value_ & self.__class__(other)._value_)
844
845 def __xor__(self, other):
846 if not isinstance(other, (self.__class__, int)):
847 return NotImplemented
848 return self.__class__(self._value_ ^ self.__class__(other)._value_)
849
850 __ror__ = __or__
851 __rand__ = __and__
852 __rxor__ = __xor__
853
854 def __invert__(self):
Ethan Furman3515dcc2016-09-18 13:15:41 -0700855 result = self.__class__(~self._value_)
856 return result
Ethan Furmanee47e5c2016-08-31 00:12:15 -0700857
858
859def _high_bit(value):
Ethan Furman04439532016-09-02 15:50:21 -0700860 """returns index of highest bit, or -1 if value is zero or negative"""
Ethan Furman3515dcc2016-09-18 13:15:41 -0700861 return value.bit_length() - 1
Ethan Furmanee47e5c2016-08-31 00:12:15 -0700862
Ethan Furmanf24bb352013-07-18 17:05:39 -0700863def unique(enumeration):
864 """Class decorator for enumerations ensuring unique member values."""
865 duplicates = []
866 for name, member in enumeration.__members__.items():
867 if name != member.name:
868 duplicates.append((name, member.name))
869 if duplicates:
870 alias_details = ', '.join(
871 ["%s -> %s" % (alias, name) for (alias, name) in duplicates])
872 raise ValueError('duplicate values found in %r: %s' %
873 (enumeration, alias_details))
874 return enumeration
Ethan Furman3515dcc2016-09-18 13:15:41 -0700875
876def _decompose(flag, value):
877 """Extract all members from the value."""
878 # _decompose is only called if the value is not named
879 not_covered = value
880 negative = value < 0
Ethan Furman3515dcc2016-09-18 13:15:41 -0700881 members = []
HongWeipeng0b41a922019-11-27 06:36:02 +0800882 for member in flag:
883 member_value = member.value
Ethan Furman3515dcc2016-09-18 13:15:41 -0700884 if member_value and member_value & value == member_value:
885 members.append(member)
886 not_covered &= ~member_value
HongWeipeng0b41a922019-11-27 06:36:02 +0800887 if not negative:
888 tmp = not_covered
889 while tmp:
890 flag_value = 2 ** _high_bit(tmp)
891 if flag_value in flag._value2member_map_:
892 members.append(flag._value2member_map_[flag_value])
893 not_covered &= ~flag_value
894 tmp &= ~flag_value
Ethan Furman3515dcc2016-09-18 13:15:41 -0700895 if not members and value in flag._value2member_map_:
896 members.append(flag._value2member_map_[value])
897 members.sort(key=lambda m: m._value_, reverse=True)
898 if len(members) > 1 and members[0].value == value:
899 # we have the breakdown, don't need the value member itself
900 members.pop(0)
901 return members, not_covered