Replace Enum with an internal one for significant speed improvements.
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)
This commit is contained in:
Rapptz
parent
8bf0482af5
commit
991140eebe
142
discord/enums.py
142
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 '<enum %r>' % 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
|
||||
|
||||
@ -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
|
||||
|
||||
@ -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'))
|
||||
|
||||
Loading…
x
Reference in New Issue
Block a user