feat: Implement own IntEnum & StrEnum (#1201)

Damego · web-flow · commit f052b44fe392 · 2023-01-21T21:02:56.000+05:00
* feat: Implement own `IntEnum` &amp; `StrEnum`

* feat: add logging

* Update team.py
diff --git a/interactions/api/gateway/client.py b/interactions/api/gateway/client.py
@@ -17,7 +17,6 @@
     wait_for,
 )
 from contextlib import suppress
-from enum import IntEnum
 from sys import platform, version_info
 from time import perf_counter
 from typing import TYPE_CHECKING, Any, Dict, List, Optional, Tuple, Type, Union
@@ -26,7 +25,7 @@
 from aiohttp import ClientWebSocketResponse, WSMessage, WSMsgType
 
 from ...base import __version__, get_logger
-from ...client.enums import ComponentType, InteractionType, OptionType
+from ...client.enums import ComponentType, IntEnum, InteractionType, OptionType
 from ...client.models import Option
 from ...utils.missing import MISSING
 from ..dispatch import Listener
diff --git a/interactions/api/models/audit_log.py b/interactions/api/models/audit_log.py
@@ -1,8 +1,8 @@
 # versionadded declared in docs gen file
 
-from enum import IntEnum
 from typing import TYPE_CHECKING, List, Optional, TypeVar
 
+from ...client.enums import IntEnum
 from ...utils.attrs_utils import DictSerializerMixin, convert_list, define, field
 from .channel import Channel
 from .misc import Snowflake
diff --git a/interactions/api/models/channel.py b/interactions/api/models/channel.py
@@ -1,6 +1,5 @@
 from asyncio import Task, create_task, get_running_loop, sleep
 from datetime import datetime, timedelta, timezone
-from enum import IntEnum
 from inspect import isawaitable
 from math import inf
 from typing import (
@@ -17,6 +16,7 @@
 )
 from warnings import warn
 
+from ...client.enums import IntEnum
 from ...utils.abc.base_context_managers import BaseAsyncContextManager
 from ...utils.abc.base_iterators import DiscordPaginationIterator
 from ...utils.attrs_utils import (
diff --git a/interactions/api/models/flags.py b/interactions/api/models/flags.py
@@ -1,4 +1,6 @@
-from enum import Enum, IntFlag
+from enum import IntFlag
+
+from ...client.enums import StrEnum
 
 __all__ = ("Intents", "AppFlags", "StatusType", "UserFlags", "Permissions", "MessageFlags")
 
@@ -176,7 +178,7 @@ class AppFlags(IntFlag):
     APPLICATION_COMMAND_BADGE = 1 << 23
 
 
-class StatusType(str, Enum):
+class StatusType(StrEnum):
     """
     An enumerable object representing Discord status icons that a user may have.
     """
diff --git a/interactions/api/models/guild.py b/interactions/api/models/guild.py
@@ -1,5 +1,4 @@
 from datetime import datetime
-from enum import Enum, IntEnum
 from inspect import isawaitable
 from math import inf
 from typing import (
@@ -16,6 +15,7 @@
 )
 from warnings import warn
 
+from ...client.enums import IntEnum, StrEnum
 from ...utils.abc.base_iterators import DiscordPaginationIterator
 from ...utils.attrs_utils import (
     ClientSerializerMixin,
@@ -125,7 +125,7 @@ class InviteTargetType(IntEnum):
     EMBEDDED_APPLICATION = 2
 
 
-class GuildFeatures(Enum):
+class GuildFeatures(StrEnum):
     ANIMATED_BANNER = "ANIMATED_BANNER"
     ANIMATED_ICON = "ANIMATED_ICON"
     BANNER = "BANNER"
diff --git a/interactions/api/models/message.py b/interactions/api/models/message.py
@@ -1,9 +1,9 @@
 import contextlib
 from datetime import datetime
-from enum import IntEnum
 from io import BytesIO
 from typing import TYPE_CHECKING, List, Optional, Union
 
+from ...client.enums import IntEnum
 from ...client.models.component import ActionRow, Button, SelectMenu
 from ...utils.attrs_utils import (
     ClientSerializerMixin,
diff --git a/interactions/api/models/misc.py b/interactions/api/models/misc.py
@@ -7,14 +7,14 @@
 
 import datetime
 from base64 import b64encode
-from enum import Enum, IntEnum
 from io import FileIO, IOBase
 from logging import Logger
 from math import floor
 from os.path import basename
 from typing import List, Optional, Union
 
 from ...base import get_logger
+from ...client.enums import IntEnum, StrEnum
 from ...utils.attrs_utils import DictSerializerMixin, convert_int, convert_list, define, field
 from ...utils.missing import MISSING
 from ..error import LibraryException
@@ -413,7 +413,7 @@ def filename(self) -> str:
         return self._name.split("/")[-1].split(".")[0]
 
 
-class AllowedMentionType(str, Enum):
+class AllowedMentionType(StrEnum):
     """
     .. versionadded:: 4.3.2
 
diff --git a/interactions/api/models/presence.py b/interactions/api/models/presence.py
@@ -1,7 +1,7 @@
 import time
-from enum import IntEnum
 from typing import Any, List, Optional
 
+from ...client.enums import IntEnum
 from ...utils.attrs_utils import DictSerializerMixin, convert_list, define, field
 from .emoji import Emoji
 from .flags import StatusType
diff --git a/interactions/api/models/team.py b/interactions/api/models/team.py
@@ -1,8 +1,7 @@
 from datetime import datetime
-from enum import IntEnum
 from typing import Any, Dict, List, Optional, Union
 
-from ...client.enums import Locale
+from ...client.enums import IntEnum, Locale
 from ...utils.attrs_utils import (
     ClientSerializerMixin,
     DictSerializerMixin,
diff --git a/interactions/api/models/webhook.py b/interactions/api/models/webhook.py
@@ -1,9 +1,9 @@
 # versionadded is specified in docs gen file
 
 from datetime import datetime
-from enum import IntEnum
 from typing import TYPE_CHECKING, List, Optional, Union
 
+from ...client.enums import IntEnum
 from ...utils.attrs_utils import ClientSerializerMixin, define, field
 from ...utils.missing import MISSING
 from ..error import LibraryException
diff --git a/interactions/client/enums.py b/interactions/client/enums.py
@@ -1,6 +1,14 @@
-from enum import Enum, IntEnum
+import logging
+from enum import Enum
+from typing import Any, Type
+
+from ..base import get_logger
+
+log: logging.Logger = get_logger("enums")
 
 __all__ = (
+    "IntEnum",
+    "StrEnum",
     "ApplicationCommandType",
     "InteractionType",
     "InteractionCallbackType",
@@ -13,6 +21,35 @@
 )
 
 
+def _cursed_enum(cls: Type[Enum], obj: type, value: Any) -> Enum:
+    log.info(f"Enum class {cls.__name__} received an unexpected value `{value}`.")
+
+    new = obj.__new__(cls)  # type: ignore
+    new._name_ = f"UNKNOWN: {value}"
+    new._value_ = value
+
+    return cls._value2member_map_.setdefault(value, new)
+
+
+class IntEnum(int, Enum):
+    """Enum where members must be ints"""
+
+    @classmethod
+    def _missing_(cls, value: int) -> Enum:
+        return _cursed_enum(cls, int, value)
+
+
+class StrEnum(str, Enum):
+    """Enum where members must be strings"""
+
+    @classmethod
+    def _missing_(cls, value: str) -> Enum:
+        return _cursed_enum(cls, str, value)
+
+    def __str__(self):
+        return self.value
+
+
 class ApplicationCommandType(IntEnum):
     """
     An enumerable object representing the types of application commands.
@@ -168,7 +205,7 @@ class TextStyleType(IntEnum):
     PARAGRAPH = 2
 
 
-class Locale(str, Enum):
+class Locale(StrEnum):
     """
     .. versionadded:: 4.2.0
 
diff --git a/interactions/ext/error.py b/interactions/ext/error.py
@@ -1,4 +1,4 @@
-from enum import Enum
+from ..client.enums import StrEnum
 
 __all__ = (
     "ErrorType",
@@ -9,7 +9,7 @@
 )
 
 
-class ErrorType(str, Enum):
+class ErrorType(StrEnum):
     """
     An enumerable object representing the type of error responses raised.
 
diff --git a/interactions/ext/version.py b/interactions/ext/version.py
@@ -1,8 +1,8 @@
-from enum import Enum
 from hashlib import md5
 from string import ascii_lowercase
 from typing import List, Optional, Union
 
+from ..client.enums import StrEnum
 from .error import IncorrectAlphanumeric, TooManyAuthors
 
 __all__ = (
@@ -12,7 +12,7 @@
 )
 
 
-class VersionAlphanumericType(str, Enum):
+class VersionAlphanumericType(StrEnum):
     ALPHA = "alpha"
     BETA = "beta"
     RELEASE_CANDIDATE = "rc"
diff --git a/interactions/utils/get.py b/interactions/utils/get.py
@@ -1,7 +1,6 @@
 # versionadded declared in docs gen file
 
 from asyncio import sleep
-from enum import Enum
 from inspect import isawaitable
 from logging import getLogger
 from sys import version_info
@@ -28,6 +27,7 @@ class GenericAlias:
 from ..api.models.message import Message
 from ..api.models.misc import Snowflake
 from ..api.models.role import Role
+from ..client.enums import StrEnum
 
 log = getLogger("get")
 
@@ -43,7 +43,7 @@ class GenericAlias:
 )
 
 
-class Force(str, Enum):
+class Force(StrEnum):
     """
     An enumerable object representing the force types for the get method.
 
diff --git a/interactions/utils/get.pyi b/interactions/utils/get.pyi
@@ -1,8 +1,7 @@
-from enum import Enum
 from typing import Awaitable, Coroutine, List, Literal, Optional, Type, TypeVar, Union, overload
 
-from interactions.client.bot import Client
-
+from ..client.bot import Client
+from ..client.enums import StrEnum
 from ..api.http.client import HTTPClient
 from ..api.models.channel import Channel
 from ..api.models.guild import Guild
@@ -19,7 +18,7 @@ _T = TypeVar("_T")
 
 __all__: tuple
 
-class Force(str, Enum):
+class Force(StrEnum):
     """
     An enum representing the force methods for the get method
     """
