From 991140eebe4a618354ada9ec592e9a405dab604d Mon Sep 17 00:00:00 2001 From: Rapptz Date: Sun, 9 Jun 2019 00:06:34 -0400 Subject: [PATCH] Replace Enum with an internal one for significant speed improvements. MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit This has been a massive pain point for me personally due to the poor design of the Enum class leading to the common use cases used in the library being significantly slow. Since this Enum is not public facing in terms of *creation*, I can only implement the APIs that are used when *accessing* them. This Enum is a drop-in replacement to the pre-existing enum.Enum class except it comes with significant speed-ups. Since this is a lot to go over, I will let the numbers speak for themselves: In [4]: %timeit enums.try_enum(enums.Status, 'offline') 263 ns ± 34.3 ns per loop (7 runs, 1000000 loops each) In [5]: %timeit NeoStatus.try_value('offline') 134 ns ± 0.859 ns per loop (7 runs, 10000000 loops each) In [6]: %timeit enums.Status.offline 116 ns ± 0.378 ns per loop (7 runs, 10000000 loops each) In [7]: %timeit NeoStatus.offline 31.6 ns ± 0.327 ns per loop (7 runs, 10000000 loops each) In [8]: %timeit enums.Status.offline.value 382 ns ± 15.2 ns per loop (7 runs, 1000000 loops each) In [9]: %timeit NeoStatus.offline.value 65.5 ns ± 0.953 ns per loop (7 runs, 10000000 loops each) In [10]: %timeit str(enums.Status.offline) 630 ns ± 14.8 ns per loop (7 runs, 1000000 loops each) In [11]: %timeit str(NeoStatus.offline) 253 ns ± 3.53 ns per loop (7 runs, 1000000 loops each) In [12]: %timeit enums.Status('offline') 697 ns ± 8.42 ns per loop (7 runs, 1000000 loops each) In [13]: %timeit NeoStatus('offline') 182 ns ± 1.83 ns per loop (7 runs, 10000000 loops each) --- discord/enums.py | 142 ++++++++++++++++++------------ discord/ext/commands/cooldowns.py | 4 +- discord/state.py | 5 +- 3 files changed, 92 insertions(+), 59 deletions(-) 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'))