Browse Source

Make enum code work with typecheckers

This also makes it so invalid enum values fall back to a proxy type
that still works similar to actual values.
pull/6658/head
Rapptz 4 years ago
parent
commit
008432c76a
  1. 35
      discord/enums.py

35
discord/enums.py

@ -24,6 +24,7 @@ DEALINGS IN THE SOFTWARE.
import types import types
from collections import namedtuple from collections import namedtuple
from typing import Any, TYPE_CHECKING, Type, TypeVar
__all__ = ( __all__ = (
'Enum', 'Enum',
@ -90,6 +91,7 @@ class EnumMeta(type):
attrs['_enum_value_map_'] = value_mapping attrs['_enum_value_map_'] = value_mapping
attrs['_enum_member_map_'] = member_mapping attrs['_enum_member_map_'] = member_mapping
attrs['_enum_member_names_'] = member_names attrs['_enum_member_names_'] = member_names
attrs['_enum_value_cls_'] = value_cls
actual_cls = super().__new__(cls, name, bases, attrs) actual_cls = super().__new__(cls, name, bases, attrs)
value_cls._actual_enum_cls_ = actual_cls value_cls._actual_enum_cls_ = actual_cls
return actual_cls return actual_cls
@ -133,14 +135,16 @@ class EnumMeta(type):
except AttributeError: except AttributeError:
return False return False
class Enum(metaclass=EnumMeta): if TYPE_CHECKING:
@classmethod from enum import Enum
def try_value(cls, value): else:
try: class Enum(metaclass=EnumMeta):
return cls._enum_value_map_[value] @classmethod
except (KeyError, TypeError): def try_value(cls, value):
return value try:
return cls._enum_value_map_[value]
except (KeyError, TypeError):
return value
class ChannelType(Enum): class ChannelType(Enum):
text = 0 text = 0
@ -422,13 +426,20 @@ class InteractionType(Enum):
ping = 1 ping = 1
application_command = 2 application_command = 2
def try_enum(cls, val): _T = TypeVar('_T')
def create_unknown_value(cls: Type[_T], val: Any) -> _T:
value_cls = cls._enum_value_cls_ # type: ignore
name = f'unknown_{val}'
return value_cls(name=name, value=val)
def try_enum(cls: Type[_T], val: Any) -> _T:
"""A function that tries to turn the value into enum ``cls``. """A function that tries to turn the value into enum ``cls``.
If it fails it returns the value instead. If it fails it returns a proxy invalid value instead.
""" """
try: try:
return cls._enum_value_map_[val] return cls._enum_value_map_[val] # type: ignore
except (KeyError, TypeError, AttributeError): except (KeyError, TypeError, AttributeError):
return val return create_unknown_value(cls, val)

Loading…
Cancel
Save