blob: ebadd9f662a13894013c29c2c0abfb1759c811d8 [file] [log] [blame]
Olivier Deprezf4ef2d02021-04-20 13:36:24 +02001import sys
2from types import MappingProxyType, DynamicClassAttribute
3
4
5__all__ = [
6 'EnumMeta',
7 'Enum', 'IntEnum', 'Flag', 'IntFlag',
8 'auto', 'unique',
9 ]
10
11
12def _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
20def _is_dunder(name):
21 """Returns True if a __dunder__ name, False otherwise."""
22 return (len(name) > 4 and
23 name[:2] == name[-2:] == '__' and
24 name[2] != '_' and
25 name[-3] != '_')
26
27
28def _is_sunder(name):
29 """Returns True if a _sunder_ name, False otherwise."""
30 return (len(name) > 2 and
31 name[0] == name[-1] == '_' and
32 name[1:2] != '_' and
33 name[-2:-1] != '_')
34
35
36def _make_class_unpicklable(cls):
37 """Make the given class un-picklable."""
38 def _break_on_call_reduce(self, proto):
39 raise TypeError('%r cannot be pickled' % self)
40 cls.__reduce_ex__ = _break_on_call_reduce
41 cls.__module__ = '<unknown>'
42
43_auto_null = object()
44class auto:
45 """
46 Instances are replaced with an appropriate value in Enum class suites.
47 """
48 value = _auto_null
49
50
51class _EnumDict(dict):
52 """Track enum member order and ensure member names are not reused.
53
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 = []
61 self._last_values = []
62 self._ignore = []
63 self._auto_called = False
64
65 def __setitem__(self, key, value):
66 """Changes anything not dundered or not a descriptor.
67
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):
75 if key not in (
76 '_order_', '_create_pseudo_member_',
77 '_generate_next_value_', '_missing_', '_ignore_',
78 ):
79 raise ValueError('_names_ are reserved for future Enum use')
80 if key == '_generate_next_value_':
81 # check if members already defined as auto()
82 if self._auto_called:
83 raise TypeError("_generate_next_value_ must be defined before members")
84 setattr(self, '_generate_next_value', value)
85 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, ))
94 elif _is_dunder(key):
95 if key == '__order__':
96 key = '_order_'
97 elif key in self._member_names:
98 # descriptor overwriting an enum?
99 raise TypeError('Attempted to reuse key: %r' % key)
100 elif key in self._ignore:
101 pass
102 elif not _is_descriptor(value):
103 if key in self:
104 # enum overwriting a descriptor?
105 raise TypeError('%r already defined as: %r' % (key, self[key]))
106 if isinstance(value, auto):
107 if value.value == _auto_null:
108 value.value = self._generate_next_value(key, 1, len(self._member_names), self._last_values[:])
109 self._auto_called = True
110 value = value.value
111 self._member_names.append(key)
112 self._last_values.append(value)
113 super().__setitem__(key, value)
114
115
116# 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`
119Enum = None
120
121
122class EnumMeta(type):
123 """Metaclass for Enum"""
124 @classmethod
125 def __prepare__(metacls, cls, bases):
126 # check that previous enum members do not exist
127 metacls._check_for_existing_members(cls, bases)
128 # create the namespace dict
129 enum_dict = _EnumDict()
130 # inherit previous flags and _generate_next_value_ function
131 member_type, first_enum = metacls._get_mixins_(cls, bases)
132 if first_enum is not None:
133 enum_dict['_generate_next_value_'] = getattr(first_enum, '_generate_next_value_', None)
134 return enum_dict
135
136 def __new__(metacls, cls, bases, classdict):
137 # an Enum class is final once enumeration items have been defined; it
138 # cannot be mixed with other types (int, float, etc.) if it has an
139 # inherited __new__ unless a new __new__ is defined (or the resulting
140 # class will fail).
141 #
142 # remove any keys listed in _ignore_
143 classdict.setdefault('_ignore_', []).append('_ignore_')
144 ignore = classdict['_ignore_']
145 for key in ignore:
146 classdict.pop(key, None)
147 member_type, first_enum = metacls._get_mixins_(cls, bases)
148 __new__, save_new, use_args = metacls._find_new_(classdict, member_type,
149 first_enum)
150
151 # save enum items into separate mapping so they don't get baked into
152 # the new class
153 enum_members = {k: classdict[k] for k in classdict._member_names}
154 for name in classdict._member_names:
155 del classdict[name]
156
157 # adjust the sunders
158 _order_ = classdict.pop('_order_', None)
159
160 # check for illegal enum names (any others?)
161 invalid_names = set(enum_members) & {'mro', ''}
162 if invalid_names:
163 raise ValueError('Invalid enum member name: {0}'.format(
164 ','.join(invalid_names)))
165
166 # create a default docstring if one has not been provided
167 if '__doc__' not in classdict:
168 classdict['__doc__'] = 'An enumeration.'
169
170 # create our new Enum type
171 enum_class = super().__new__(metacls, cls, bases, classdict)
172 enum_class._member_names_ = [] # names in definition order
173 enum_class._member_map_ = {} # name->value map
174 enum_class._member_type_ = member_type
175
176 # save DynamicClassAttribute attributes from super classes so we know
177 # if we can take the shortcut of storing members in the class dict
178 dynamic_attributes = {k for c in enum_class.mro()
179 for k, v in c.__dict__.items()
180 if isinstance(v, DynamicClassAttribute)}
181
182 # Reverse value->name map for hashable values.
183 enum_class._value2member_map_ = {}
184
185 # If a custom type is mixed into the Enum, and it does not know how
186 # to pickle itself, pickle.dumps will succeed but pickle.loads will
187 # fail. Rather than have the error show up later and possibly far
188 # from the source, sabotage the pickle protocol for this class so
189 # that pickle.dumps also fails.
190 #
191 # However, if the new class implements its own __reduce_ex__, do not
192 # sabotage -- it's on them to make sure it works correctly. We use
193 # __reduce_ex__ instead of any of the others as it is preferred by
194 # pickle over __reduce__, and it handles all pickle protocols.
195 if '__reduce_ex__' not in classdict:
196 if member_type is not object:
197 methods = ('__getnewargs_ex__', '__getnewargs__',
198 '__reduce_ex__', '__reduce__')
199 if not any(m in member_type.__dict__ for m in methods):
200 _make_class_unpicklable(enum_class)
201
202 # instantiate them, checking for duplicates as we go
203 # we instantiate first instead of checking for duplicates first in case
204 # a custom __new__ is doing something funky with the values -- such as
205 # auto-numbering ;)
206 for member_name in classdict._member_names:
207 value = enum_members[member_name]
208 if not isinstance(value, tuple):
209 args = (value, )
210 else:
211 args = value
212 if member_type is tuple: # special case for tuple enums
213 args = (args, ) # wrap it one more time
214 if not use_args:
215 enum_member = __new__(enum_class)
216 if not hasattr(enum_member, '_value_'):
217 enum_member._value_ = value
218 else:
219 enum_member = __new__(enum_class, *args)
220 if not hasattr(enum_member, '_value_'):
221 if member_type is object:
222 enum_member._value_ = value
223 else:
224 enum_member._value_ = member_type(*args)
225 value = enum_member._value_
226 enum_member._name_ = member_name
227 enum_member.__objclass__ = enum_class
228 enum_member.__init__(*args)
229 # If another member with the same value was already defined, the
230 # new member becomes an alias to the existing one.
231 for name, canonical_member in enum_class._member_map_.items():
232 if canonical_member._value_ == enum_member._value_:
233 enum_member = canonical_member
234 break
235 else:
236 # Aliases don't appear in member names (only in __members__).
237 enum_class._member_names_.append(member_name)
238 # performance boost for any member that would not shadow
239 # a DynamicClassAttribute
240 if member_name not in dynamic_attributes:
241 setattr(enum_class, member_name, enum_member)
242 # now add to _member_map_
243 enum_class._member_map_[member_name] = enum_member
244 try:
245 # This may fail if value is not hashable. We can't add the value
246 # to the map, and by-value lookups for this value will be
247 # linear.
248 enum_class._value2member_map_[value] = enum_member
249 except TypeError:
250 pass
251
252 # double check that repr and friends are not the mixin's or various
253 # things break (such as pickle)
254 # however, if the method is defined in the Enum itself, don't replace
255 # it
256 for name in ('__repr__', '__str__', '__format__', '__reduce_ex__'):
257 if name in classdict:
258 continue
259 class_method = getattr(enum_class, name)
260 obj_method = getattr(member_type, name, None)
261 enum_method = getattr(first_enum, name, None)
262 if obj_method is not None and obj_method is class_method:
263 setattr(enum_class, name, enum_method)
264
265 # replace any other __new__ with our own (as long as Enum is not None,
266 # anyway) -- again, this is to support pickle
267 if Enum is not None:
268 # if the user defined their own __new__, save it before it gets
269 # clobbered in case they subclass later
270 if save_new:
271 enum_class.__new_member__ = __new__
272 enum_class.__new__ = Enum.__new__
273
274 # py3 support for definition order (helps keep py2/py3 code in sync)
275 if _order_ is not None:
276 if isinstance(_order_, str):
277 _order_ = _order_.replace(',', ' ').split()
278 if _order_ != enum_class._member_names_:
279 raise TypeError('member order does not match _order_')
280
281 return enum_class
282
283 def __bool__(self):
284 """
285 classes/types should always be True.
286 """
287 return True
288
289 def __call__(cls, value, names=None, *, module=None, qualname=None, type=None, start=1):
290 """Either returns an existing member, or creates a new enum class.
291
292 This method is used both when an enum class is given a value to match
293 to an enumeration member (i.e. Color(3)) and for the functional API
294 (i.e. Color = Enum('Color', names='RED GREEN BLUE')).
295
296 When used for the functional API:
297
298 `value` will be the name of the new class.
299
300 `names` should be either a string of white-space/comma delimited names
301 (values will start at `start`), or an iterator/mapping of name, value pairs.
302
303 `module` should be set to the module this class is being created in;
304 if it is not set, an attempt to find that module will be made, but if
305 it fails the class will not be picklable.
306
307 `qualname` should be set to the actual location this class can be found
308 at in its module; by default it is set to the global scope. If this is
309 not correct, unpickling will fail in some circumstances.
310
311 `type`, if set, will be mixed in as the first base class.
312
313 """
314 if names is None: # simple value lookup
315 return cls.__new__(cls, value)
316 # otherwise, functional API: we're creating a new Enum type
317 return cls._create_(value, names, module=module, qualname=qualname, type=type, start=start)
318
319 def __contains__(cls, member):
320 if not isinstance(member, Enum):
321 raise TypeError(
322 "unsupported operand type(s) for 'in': '%s' and '%s'" % (
323 type(member).__qualname__, cls.__class__.__qualname__))
324 return isinstance(member, cls) and member._name_ in cls._member_map_
325
326 def __delattr__(cls, attr):
327 # nicer error message when someone tries to delete an attribute
328 # (see issue19025).
329 if attr in cls._member_map_:
330 raise AttributeError(
331 "%s: cannot delete Enum member." % cls.__name__)
332 super().__delattr__(attr)
333
334 def __dir__(self):
335 return (['__class__', '__doc__', '__members__', '__module__'] +
336 self._member_names_)
337
338 def __getattr__(cls, name):
339 """Return the enum member matching `name`
340
341 We use __getattr__ instead of descriptors or inserting into the enum
342 class' __dict__ in order to support `name` and `value` being both
343 properties for enum members (which live in the class' __dict__) and
344 enum members themselves.
345
346 """
347 if _is_dunder(name):
348 raise AttributeError(name)
349 try:
350 return cls._member_map_[name]
351 except KeyError:
352 raise AttributeError(name) from None
353
354 def __getitem__(cls, name):
355 return cls._member_map_[name]
356
357 def __iter__(cls):
358 return (cls._member_map_[name] for name in cls._member_names_)
359
360 def __len__(cls):
361 return len(cls._member_names_)
362
363 @property
364 def __members__(cls):
365 """Returns a mapping of member name->value.
366
367 This mapping lists all enum members, including aliases. Note that this
368 is a read-only view of the internal mapping.
369
370 """
371 return MappingProxyType(cls._member_map_)
372
373 def __repr__(cls):
374 return "<enum %r>" % cls.__name__
375
376 def __reversed__(cls):
377 return (cls._member_map_[name] for name in reversed(cls._member_names_))
378
379 def __setattr__(cls, name, value):
380 """Block attempts to reassign Enum members.
381
382 A simple assignment to the class namespace only changes one of the
383 several possible ways to get an Enum member from the Enum class,
384 resulting in an inconsistent Enumeration.
385
386 """
387 member_map = cls.__dict__.get('_member_map_', {})
388 if name in member_map:
389 raise AttributeError('Cannot reassign members.')
390 super().__setattr__(name, value)
391
392 def _create_(cls, class_name, names, *, module=None, qualname=None, type=None, start=1):
393 """Convenience method to create a new Enum class.
394
395 `names` can be:
396
397 * A string containing member names, separated either with spaces or
398 commas. Values are incremented by 1 from `start`.
399 * An iterable of member names. Values are incremented by 1 from `start`.
400 * An iterable of (member name, value) pairs.
401 * A mapping of member name -> value pairs.
402
403 """
404 metacls = cls.__class__
405 bases = (cls, ) if type is None else (type, cls)
406 _, first_enum = cls._get_mixins_(cls, bases)
407 classdict = metacls.__prepare__(class_name, bases)
408
409 # special processing needed for names?
410 if isinstance(names, str):
411 names = names.replace(',', ' ').split()
412 if isinstance(names, (tuple, list)) and names and isinstance(names[0], str):
413 original_names, names = names, []
414 last_values = []
415 for count, name in enumerate(original_names):
416 value = first_enum._generate_next_value_(name, start, count, last_values[:])
417 last_values.append(value)
418 names.append((name, value))
419
420 # Here, names is either an iterable of (name, value) or a mapping.
421 for item in names:
422 if isinstance(item, str):
423 member_name, member_value = item, names[item]
424 else:
425 member_name, member_value = item
426 classdict[member_name] = member_value
427 enum_class = metacls.__new__(metacls, class_name, bases, classdict)
428
429 # TODO: replace the frame hack if a blessed way to know the calling
430 # module is ever developed
431 if module is None:
432 try:
433 module = sys._getframe(2).f_globals['__name__']
434 except (AttributeError, ValueError, KeyError):
435 pass
436 if module is None:
437 _make_class_unpicklable(enum_class)
438 else:
439 enum_class.__module__ = module
440 if qualname is not None:
441 enum_class.__qualname__ = qualname
442
443 return enum_class
444
445 def _convert_(cls, name, module, filter, source=None):
446 """
447 Create a new Enum subclass that replaces a collection of global constants
448 """
449 # convert all constants from source (or module) that pass filter() to
450 # a new Enum called name, and export the enum and its members back to
451 # module;
452 # also, replace the __reduce_ex__ method so unpickling works in
453 # previous Python versions
454 module_globals = vars(sys.modules[module])
455 if source:
456 source = vars(source)
457 else:
458 source = module_globals
459 # _value2member_map_ is populated in the same order every time
460 # for a consistent reverse mapping of number to name when there
461 # are multiple names for the same number.
462 members = [
463 (name, value)
464 for name, value in source.items()
465 if filter(name)]
466 try:
467 # sort by value
468 members.sort(key=lambda t: (t[1], t[0]))
469 except TypeError:
470 # unless some values aren't comparable, in which case sort by name
471 members.sort(key=lambda t: t[0])
472 cls = cls(name, members, module=module)
473 cls.__reduce_ex__ = _reduce_ex_by_name
474 module_globals.update(cls.__members__)
475 module_globals[name] = cls
476 return cls
477
478 @staticmethod
479 def _check_for_existing_members(class_name, bases):
480 for chain in bases:
481 for base in chain.__mro__:
482 if issubclass(base, Enum) and base._member_names_:
483 raise TypeError("%s: cannot extend enumeration %r" % (class_name, base.__name__))
484
485 @staticmethod
486 def _get_mixins_(class_name, bases):
487 """Returns the type for creating enum members, and the first inherited
488 enum class.
489
490 bases: the tuple of bases that was given to __new__
491
492 """
493 if not bases:
494 return object, Enum
495
496 def _find_data_type(bases):
497 data_types = []
498 for chain in bases:
499 candidate = None
500 for base in chain.__mro__:
501 if base is object:
502 continue
503 elif '__new__' in base.__dict__:
504 if issubclass(base, Enum):
505 continue
506 data_types.append(candidate or base)
507 break
508 elif not issubclass(base, Enum):
509 candidate = base
510 if len(data_types) > 1:
511 raise TypeError('%r: too many data types: %r' % (class_name, data_types))
512 elif data_types:
513 return data_types[0]
514 else:
515 return None
516
517 # ensure final parent class is an Enum derivative, find any concrete
518 # data type, and check that Enum has no members
519 first_enum = bases[-1]
520 if not issubclass(first_enum, Enum):
521 raise TypeError("new enumerations should be created as "
522 "`EnumName([mixin_type, ...] [data_type,] enum_type)`")
523 member_type = _find_data_type(bases) or object
524 if first_enum._member_names_:
525 raise TypeError("Cannot extend enumerations")
526 return member_type, first_enum
527
528 @staticmethod
529 def _find_new_(classdict, member_type, first_enum):
530 """Returns the __new__ to be used for creating the enum members.
531
532 classdict: the class dictionary given to __new__
533 member_type: the data type whose __new__ will be used by default
534 first_enum: enumeration to check for an overriding __new__
535
536 """
537 # now find the correct __new__, checking to see of one was defined
538 # by the user; also check earlier enum classes in case a __new__ was
539 # saved as __new_member__
540 __new__ = classdict.get('__new__', None)
541
542 # should __new__ be saved as __new_member__ later?
543 save_new = __new__ is not None
544
545 if __new__ is None:
546 # check all possibles for __new_member__ before falling back to
547 # __new__
548 for method in ('__new_member__', '__new__'):
549 for possible in (member_type, first_enum):
550 target = getattr(possible, method, None)
551 if target not in {
552 None,
553 None.__new__,
554 object.__new__,
555 Enum.__new__,
556 }:
557 __new__ = target
558 break
559 if __new__ is not None:
560 break
561 else:
562 __new__ = object.__new__
563
564 # if a non-object.__new__ is used then whatever value/tuple was
565 # assigned to the enum member name will be passed to __new__ and to the
566 # new enum member's __init__
567 if __new__ is object.__new__:
568 use_args = False
569 else:
570 use_args = True
571 return __new__, save_new, use_args
572
573
574class Enum(metaclass=EnumMeta):
575 """Generic enumeration.
576
577 Derive from this class to define new enumerations.
578
579 """
580 def __new__(cls, value):
581 # all enum instances are actually created during class construction
582 # without calling this method; this method is called by the metaclass'
583 # __call__ (i.e. Color(3) ), and by pickle
584 if type(value) is cls:
585 # For lookups like Color(Color.RED)
586 return value
587 # by-value search for a matching enum member
588 # see if it's in the reverse mapping (for hashable values)
589 try:
590 return cls._value2member_map_[value]
591 except KeyError:
592 # Not found, no need to do long O(n) search
593 pass
594 except TypeError:
595 # not there, now do long search -- O(n) behavior
596 for member in cls._member_map_.values():
597 if member._value_ == value:
598 return member
599 # still not found -- try _missing_ hook
600 try:
601 exc = None
602 result = cls._missing_(value)
603 except Exception as e:
604 exc = e
605 result = None
606 if isinstance(result, cls):
607 return result
608 else:
609 ve_exc = ValueError("%r is not a valid %s" % (value, cls.__qualname__))
610 if result is None and exc is None:
611 raise ve_exc
612 elif exc is None:
613 exc = TypeError(
614 'error in %s._missing_: returned %r instead of None or a valid member'
615 % (cls.__name__, result)
616 )
617 exc.__context__ = ve_exc
618 raise exc
619
620 def _generate_next_value_(name, start, count, last_values):
621 for last_value in reversed(last_values):
622 try:
623 return last_value + 1
624 except TypeError:
625 pass
626 else:
627 return start
628
629 @classmethod
630 def _missing_(cls, value):
631 return None
632
633 def __repr__(self):
634 return "<%s.%s: %r>" % (
635 self.__class__.__name__, self._name_, self._value_)
636
637 def __str__(self):
638 return "%s.%s" % (self.__class__.__name__, self._name_)
639
640 def __dir__(self):
641 added_behavior = [
642 m
643 for cls in self.__class__.mro()
644 for m in cls.__dict__
645 if m[0] != '_' and m not in self._member_map_
646 ]
647 return (['__class__', '__doc__', '__module__'] + added_behavior)
648
649 def __format__(self, format_spec):
650 # mixed-in Enums should use the mixed-in type's __format__, otherwise
651 # we can get strange results with the Enum name showing up instead of
652 # the value
653
654 # pure Enum branch, or branch with __str__ explicitly overridden
655 str_overridden = type(self).__str__ != Enum.__str__
656 if self._member_type_ is object or str_overridden:
657 cls = str
658 val = str(self)
659 # mix-in branch
660 else:
661 cls = self._member_type_
662 val = self._value_
663 return cls.__format__(val, format_spec)
664
665 def __hash__(self):
666 return hash(self._name_)
667
668 def __reduce_ex__(self, proto):
669 return self.__class__, (self._value_, )
670
671 # DynamicClassAttribute is used to provide access to the `name` and
672 # `value` properties of enum members while keeping some measure of
673 # protection from modification, while still allowing for an enumeration
674 # to have members named `name` and `value`. This works because enumeration
675 # members are not set directly on the enum class -- __getattr__ is
676 # used to look them up.
677
678 @DynamicClassAttribute
679 def name(self):
680 """The name of the Enum member."""
681 return self._name_
682
683 @DynamicClassAttribute
684 def value(self):
685 """The value of the Enum member."""
686 return self._value_
687
688
689class IntEnum(int, Enum):
690 """Enum where members are also (and must be) ints"""
691
692
693def _reduce_ex_by_name(self, proto):
694 return self.name
695
696class Flag(Enum):
697 """Support for flags"""
698
699 def _generate_next_value_(name, start, count, last_values):
700 """
701 Generate the next value when not given.
702
703 name: the name of the member
704 start: the initial start value or None
705 count: the number of existing members
706 last_value: the last value assigned or None
707 """
708 if not count:
709 return start if start is not None else 1
710 for last_value in reversed(last_values):
711 try:
712 high_bit = _high_bit(last_value)
713 break
714 except Exception:
715 raise TypeError('Invalid Flag value: %r' % last_value) from None
716 return 2 ** (high_bit+1)
717
718 @classmethod
719 def _missing_(cls, value):
720 original_value = value
721 if value < 0:
722 value = ~value
723 possible_member = cls._create_pseudo_member_(value)
724 if original_value < 0:
725 possible_member = ~possible_member
726 return possible_member
727
728 @classmethod
729 def _create_pseudo_member_(cls, value):
730 """
731 Create a composite member iff value contains only members.
732 """
733 pseudo_member = cls._value2member_map_.get(value, None)
734 if pseudo_member is None:
735 # verify all bits are accounted for
736 _, extra_flags = _decompose(cls, value)
737 if extra_flags:
738 raise ValueError("%r is not a valid %s" % (value, cls.__qualname__))
739 # construct a singleton enum pseudo-member
740 pseudo_member = object.__new__(cls)
741 pseudo_member._name_ = None
742 pseudo_member._value_ = value
743 # use setdefault in case another thread already created a composite
744 # with this value
745 pseudo_member = cls._value2member_map_.setdefault(value, pseudo_member)
746 return pseudo_member
747
748 def __contains__(self, other):
749 if not isinstance(other, self.__class__):
750 raise TypeError(
751 "unsupported operand type(s) for 'in': '%s' and '%s'" % (
752 type(other).__qualname__, self.__class__.__qualname__))
753 return other._value_ & self._value_ == other._value_
754
755 def __repr__(self):
756 cls = self.__class__
757 if self._name_ is not None:
758 return '<%s.%s: %r>' % (cls.__name__, self._name_, self._value_)
759 members, uncovered = _decompose(cls, self._value_)
760 return '<%s.%s: %r>' % (
761 cls.__name__,
762 '|'.join([str(m._name_ or m._value_) for m in members]),
763 self._value_,
764 )
765
766 def __str__(self):
767 cls = self.__class__
768 if self._name_ is not None:
769 return '%s.%s' % (cls.__name__, self._name_)
770 members, uncovered = _decompose(cls, self._value_)
771 if len(members) == 1 and members[0]._name_ is None:
772 return '%s.%r' % (cls.__name__, members[0]._value_)
773 else:
774 return '%s.%s' % (
775 cls.__name__,
776 '|'.join([str(m._name_ or m._value_) for m in members]),
777 )
778
779 def __bool__(self):
780 return bool(self._value_)
781
782 def __or__(self, other):
783 if not isinstance(other, self.__class__):
784 return NotImplemented
785 return self.__class__(self._value_ | other._value_)
786
787 def __and__(self, other):
788 if not isinstance(other, self.__class__):
789 return NotImplemented
790 return self.__class__(self._value_ & other._value_)
791
792 def __xor__(self, other):
793 if not isinstance(other, self.__class__):
794 return NotImplemented
795 return self.__class__(self._value_ ^ other._value_)
796
797 def __invert__(self):
798 members, uncovered = _decompose(self.__class__, self._value_)
799 inverted = self.__class__(0)
800 for m in self.__class__:
801 if m not in members and not (m._value_ & self._value_):
802 inverted = inverted | m
803 return self.__class__(inverted)
804
805
806class IntFlag(int, Flag):
807 """Support for integer-based Flags"""
808
809 @classmethod
810 def _missing_(cls, value):
811 if not isinstance(value, int):
812 raise ValueError("%r is not a valid %s" % (value, cls.__qualname__))
813 new_member = cls._create_pseudo_member_(value)
814 return new_member
815
816 @classmethod
817 def _create_pseudo_member_(cls, value):
818 pseudo_member = cls._value2member_map_.get(value, None)
819 if pseudo_member is None:
820 need_to_create = [value]
821 # get unaccounted for bits
822 _, extra_flags = _decompose(cls, value)
823 # timer = 10
824 while extra_flags:
825 # timer -= 1
826 bit = _high_bit(extra_flags)
827 flag_value = 2 ** bit
828 if (flag_value not in cls._value2member_map_ and
829 flag_value not in need_to_create
830 ):
831 need_to_create.append(flag_value)
832 if extra_flags == -flag_value:
833 extra_flags = 0
834 else:
835 extra_flags ^= flag_value
836 for value in reversed(need_to_create):
837 # construct singleton pseudo-members
838 pseudo_member = int.__new__(cls, value)
839 pseudo_member._name_ = None
840 pseudo_member._value_ = value
841 # use setdefault in case another thread already created a composite
842 # with this value
843 pseudo_member = cls._value2member_map_.setdefault(value, pseudo_member)
844 return pseudo_member
845
846 def __or__(self, other):
847 if not isinstance(other, (self.__class__, int)):
848 return NotImplemented
849 result = self.__class__(self._value_ | self.__class__(other)._value_)
850 return result
851
852 def __and__(self, other):
853 if not isinstance(other, (self.__class__, int)):
854 return NotImplemented
855 return self.__class__(self._value_ & self.__class__(other)._value_)
856
857 def __xor__(self, other):
858 if not isinstance(other, (self.__class__, int)):
859 return NotImplemented
860 return self.__class__(self._value_ ^ self.__class__(other)._value_)
861
862 __ror__ = __or__
863 __rand__ = __and__
864 __rxor__ = __xor__
865
866 def __invert__(self):
867 result = self.__class__(~self._value_)
868 return result
869
870
871def _high_bit(value):
872 """returns index of highest bit, or -1 if value is zero or negative"""
873 return value.bit_length() - 1
874
875def unique(enumeration):
876 """Class decorator for enumerations ensuring unique member values."""
877 duplicates = []
878 for name, member in enumeration.__members__.items():
879 if name != member.name:
880 duplicates.append((name, member.name))
881 if duplicates:
882 alias_details = ', '.join(
883 ["%s -> %s" % (alias, name) for (alias, name) in duplicates])
884 raise ValueError('duplicate values found in %r: %s' %
885 (enumeration, alias_details))
886 return enumeration
887
888def _decompose(flag, value):
889 """Extract all members from the value."""
890 # _decompose is only called if the value is not named
891 not_covered = value
892 negative = value < 0
893 members = []
894 for member in flag:
895 member_value = member.value
896 if member_value and member_value & value == member_value:
897 members.append(member)
898 not_covered &= ~member_value
899 if not negative:
900 tmp = not_covered
901 while tmp:
902 flag_value = 2 ** _high_bit(tmp)
903 if flag_value in flag._value2member_map_:
904 members.append(flag._value2member_map_[flag_value])
905 not_covered &= ~flag_value
906 tmp &= ~flag_value
907 if not members and value in flag._value2member_map_:
908 members.append(flag._value2member_map_[value])
909 members.sort(key=lambda m: m._value_, reverse=True)
910 if len(members) > 1 and members[0].value == value:
911 # we have the breakdown, don't need the value member itself
912 members.pop(0)
913 return members, not_covered