Allow enums to be compared

discord/enums.py

@@ -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

docs/api.rst

@@ -1591,6 +1591,8 @@ of :class:`enum.Enum`.
 
     .. container:: operations
 
+        .. versionadded:: 2.0
+
         .. describe:: x == y
 
             Checks if two verification levels are equal.
@@ -1633,6 +1635,29 @@ of :class:`enum.Enum`.
 
     Specifies whether a :class:`Guild` has notifications on for all messages or mentions only by default.
 
+    .. container:: operations
+
+        .. versionadded:: 2.0
+
+        .. describe:: x == y
+
+            Checks if two notification levels are equal.
+        .. describe:: x != y
+
+            Checks if two notification levels are not equal.
+        .. describe:: x > y
+
+            Checks if a notification level is higher than another.
+        .. describe:: x < y
+
+            Checks if a notification level is lower than another.
+        .. describe:: x >= y
+
+            Checks if a notification level is higher or equal to another.
+        .. describe:: x <= y
+
+            Checks if a notification level is lower or equal to another.
+
     .. attribute:: all_messages
 
         Members receive notifications for every message regardless of them being mentioned.
@@ -1648,6 +1673,8 @@ of :class:`enum.Enum`.
 
     .. container:: operations
 
+        .. versionadded:: 2.0
+
         .. describe:: x == y
 
             Checks if two content filter levels are equal.
@@ -2532,6 +2559,27 @@ of :class:`enum.Enum`.
 
     .. versionadded:: 2.0
 
+    .. container:: operations
+
+        .. describe:: x == y
+
+            Checks if two NSFW levels are equal.
+        .. describe:: x != y
+
+            Checks if two NSFW levels are not equal.
+        .. describe:: x > y
+
+            Checks if a NSFW level is higher than another.
+        .. describe:: x < y
+
+            Checks if a NSFW level is lower than another.
+        .. describe:: x >= y
+
+            Checks if a NSFW level is higher or equal to another.
+        .. describe:: x <= y
+
+            Checks if a NSFW level is lower or equal to another.
+
     .. attribute:: default
 
         The guild has not been categorised yet.