diff --git a/discord/enums.py b/discord/enums.py index 5e47f8a93..2162a614b 100644 --- a/discord/enums.py +++ b/discord/enums.py @@ -24,9 +24,13 @@ FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE SOFTWARE. """ -from enum import Enum +import types +import inspect +from collections import namedtuple +from operator import attrgetter __all__ = ( + 'Enum', 'ChannelType', 'MessageType', 'VoiceRegion', @@ -48,20 +52,88 @@ __all__ = ( 'Theme', ) -def fast_lookup(cls): - # NOTE: implies hashable - try: - lookup = cls._value2member_map_ - except AttributeError: - lookup = { - member.value: member - for member in cls.__members__ - } - finally: - cls.__fast_value_lookup__ = lookup - return cls +def _create_value_cls(name): + cls = namedtuple('_EnumValue_' + name, 'name value') + cls.__repr__ = lambda self: '<%s.%s: %r>' % (name, self.name, self.value) + cls.__str__ = lambda self: '%s.%s' % (name, self.name) + return cls + +class EnumMeta(type): + def __new__(cls, name, bases, attrs): + value_mapping = {} + member_mapping = {} + member_names = [] + + value_cls = _create_value_cls(name) + for key, value in list(attrs.items()): + is_function = isinstance(value, types.FunctionType) + if key[0] == '_' and not is_function: + continue + + if is_function: + setattr(value_cls, key, value) + del attrs[key] + continue + + if inspect.ismethoddescriptor(value): + continue + + try: + new_value = value_mapping[value] + except KeyError: + new_value = value_cls(name=key, value=value) + value_mapping[value] = new_value + member_names.append(key) + + member_mapping[key] = new_value + attrs[key] = new_value + + attrs['_enum_value_map_'] = value_mapping + attrs['_enum_member_map_'] = member_mapping + attrs['_enum_member_names_'] = member_names + actual_cls = super().__new__(cls, name, bases, attrs) + value_cls._actual_enum_cls_ = actual_cls + return actual_cls + + def __iter__(cls): + return (cls._enum_member_map_[name] for name in cls._enum_member_names_) + + def __repr__(cls): + return '' % cls.__name__ + + @property + def __members__(cls): + return types.MappingProxyType(cls._enum_member_map_) + + def __call__(cls, value): + try: + return cls._enum_value_map_[value] + except (KeyError, TypeError): + raise ValueError("%r is not a valid %s" % (value, cls.__name__)) + + def __getitem__(cls, key): + return cls._enum_member_map_[key] + + def __setattr__(cls, name, value): + raise TypeError('Enums are immutable.') + + def __instancecheck__(self, instance): + # isinstance(x, Y) + # -> __instancecheck__(Y, x) + try: + return instance._actual_enum_cls_ is self + except AttributeError: + return False + +class Enum(metaclass=EnumMeta): + @classmethod + def try_value(cls, value): + try: + return cls._enum_value_map_[value] + except (KeyError, TypeError): + return value + -@fast_lookup class ChannelType(Enum): text = 0 private = 1 @@ -74,7 +146,6 @@ class ChannelType(Enum): def __str__(self): return self.name -@fast_lookup class MessageType(Enum): default = 0 recipient_add = 1 @@ -89,7 +160,6 @@ class MessageType(Enum): premium_guild_tier_2 = 10 premium_guild_tier_3 = 11 -@fast_lookup class VoiceRegion(Enum): us_west = 'us-west' us_east = 'us-east' @@ -115,7 +185,6 @@ class VoiceRegion(Enum): def __str__(self): return self.value -@fast_lookup class SpeakingState(Enum): none = 0 voice = 1 @@ -128,7 +197,6 @@ class SpeakingState(Enum): def __int__(self): return self.value -@fast_lookup class VerificationLevel(Enum): none = 0 low = 1 @@ -141,7 +209,6 @@ class VerificationLevel(Enum): def __str__(self): return self.name -@fast_lookup class ContentFilter(Enum): disabled = 0 no_role = 1 @@ -150,13 +217,11 @@ class ContentFilter(Enum): def __str__(self): return self.name -@fast_lookup class UserContentFilter(Enum): disabled = 0 friends = 1 all_messages = 2 -@fast_lookup class FriendFlags(Enum): noone = 0 mutual_guilds = 1 @@ -164,12 +229,10 @@ class FriendFlags(Enum): guild_and_friends = 3 everyone = 4 -@fast_lookup class Theme(Enum): light = 'light' dark = 'dark' -@fast_lookup class Status(Enum): online = 'online' offline = 'offline' @@ -181,7 +244,6 @@ class Status(Enum): def __str__(self): return self.value -@fast_lookup class DefaultAvatar(Enum): blurple = 0 grey = 1 @@ -193,25 +255,21 @@ class DefaultAvatar(Enum): def __str__(self): return self.name -@fast_lookup class RelationshipType(Enum): friend = 1 blocked = 2 incoming_request = 3 outgoing_request = 4 -@fast_lookup class NotificationLevel(Enum): all_messages = 0 only_mentions = 1 -@fast_lookup class AuditLogActionCategory(Enum): create = 1 delete = 2 update = 3 -@fast_lookup class AuditLogAction(Enum): guild_update = 1 channel_create = 10 @@ -294,7 +352,6 @@ class AuditLogAction(Enum): elif v < 80: return 'message' -@fast_lookup class UserFlags(Enum): staff = 1 partner = 2 @@ -305,7 +362,6 @@ class UserFlags(Enum): hypesquad_balance = 256 early_supporter = 512 -@fast_lookup class ActivityType(Enum): unknown = -1 playing = 0 @@ -313,13 +369,11 @@ class ActivityType(Enum): listening = 2 watching = 3 -@fast_lookup class HypeSquadHouse(Enum): bravery = 1 brilliance = 2 balance = 3 -@fast_lookup class PremiumType(Enum): nitro_classic = 1 nitro = 2 @@ -330,27 +384,7 @@ def try_enum(cls, val): If it fails it returns the value instead. """ - # For some ungodly reason, `cls(x)` is *really* slow - # For most use cases it's about 750ns per call - # Internally this is dispatched like follows: - # cls(x) - # cls.__new__(cls, x) - # cls._value2member_map[x] - # if above fails ^ - # find it in cls._member_map.items() - - # Accessing the _value2member_map directly gives the biggest - # boost to performance, from 750ns to 130ns - - # Now, the weird thing is that regular dict access is approx 31ns - # So there's a slowdown in the attribute access somewhere in the - # __getattr__ chain that I can't do much about - - # Since this relies on internals the enums have an internal shim - # decorator that defines an alias for my own purposes or creates - # it for me under __fast_value_lookup__ - try: - return cls.__fast_value_lookup__[val] - except (KeyError, AttributeError): + return cls._enum_value_map_[val] + except (KeyError, TypeError, AttributeError): return val diff --git a/discord/ext/commands/cooldowns.py b/discord/ext/commands/cooldowns.py index 873c832c9..5cd305f6f 100644 --- a/discord/ext/commands/cooldowns.py +++ b/discord/ext/commands/cooldowns.py @@ -24,7 +24,7 @@ FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE SOFTWARE. """ -import enum +from discord.enums import Enum import time __all__ = ( @@ -33,7 +33,7 @@ __all__ = ( 'CooldownMapping', ) -class BucketType(enum.Enum): +class BucketType(Enum): default = 0 user = 1 guild = 2 diff --git a/discord/state.py b/discord/state.py index bbee44a8c..c4b4ea6c4 100644 --- a/discord/state.py +++ b/discord/state.py @@ -28,7 +28,6 @@ import asyncio from collections import deque, namedtuple, OrderedDict import copy import datetime -import enum import itertools import logging import math @@ -45,11 +44,11 @@ from .channel import * from .raw_models import * from .member import Member from .role import Role -from .enums import ChannelType, try_enum, Status +from .enums import ChannelType, try_enum, Status, Enum from . import utils from .embeds import Embed -class ListenerType(enum.Enum): +class ListenerType(Enum): chunk = 0 Listener = namedtuple('Listener', ('type', 'future', 'predicate'))