Allow enums to be compared

This commit is contained in:
Sebastian Law
2021-08-23 23:28:39 -07:00
committed by GitHub
parent f586f4dfbd
commit 835432d161
2 changed files with 60 additions and 8 deletions

View File

@ -58,13 +58,17 @@ __all__ = (
)
def _create_value_cls(name):
def _create_value_cls(name, comparable):
cls = namedtuple('_EnumValue_' + name, 'name value')
cls.__repr__ = lambda self: f'<{name}.{self.name}: {self.value!r}>'
cls.__str__ = lambda self: f'{name}.{self.name}'
if comparable:
cls.__le__ = lambda self, other: isinstance(other, self.__class__) and self.value <= other.value
cls.__ge__ = lambda self, other: isinstance(other, self.__class__) and self.value >= other.value
cls.__lt__ = lambda self, other: isinstance(other, self.__class__) and self.value < other.value
cls.__gt__ = lambda self, other: isinstance(other, self.__class__) and self.value > other.value
return cls
def _is_descriptor(obj):
return hasattr(obj, '__get__') or hasattr(obj, '__set__') or hasattr(obj, '__delete__')
@ -76,12 +80,12 @@ class EnumMeta(type):
_enum_member_map_: ClassVar[Dict[str, Any]]
_enum_value_map_: ClassVar[Dict[Any, Any]]
def __new__(cls, name, bases, attrs):
def __new__(cls, name, bases, attrs, *, comparable: bool = False):
value_mapping = {}
member_mapping = {}
member_names = []
value_cls = _create_value_cls(name)
value_cls = _create_value_cls(name, comparable)
for key, value in list(attrs.items()):
is_descriptor = _is_descriptor(value)
if key[0] == '_' and not is_descriptor:
@ -252,7 +256,7 @@ class SpeakingState(Enum):
return self.value
class VerificationLevel(Enum):
class VerificationLevel(Enum, comparable=True):
none = 0
low = 1
medium = 2
@ -263,7 +267,7 @@ class VerificationLevel(Enum):
return self.name
class ContentFilter(Enum):
class ContentFilter(Enum, comparable=True):
disabled = 0
no_role = 1
all_members = 2
@ -296,7 +300,7 @@ class DefaultAvatar(Enum):
return self.name
class NotificationLevel(Enum):
class NotificationLevel(Enum, comparable=True):
all_messages = 0
only_mentions = 1
@ -578,7 +582,7 @@ class StagePrivacyLevel(Enum):
guild_only = 2
class NSFWLevel(Enum):
class NSFWLevel(Enum, comparable=True):
default = 0
explicit = 1
safe = 2