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