e557f69c83
In 3.5.0 and 3.5.1 asyncio.ensure_future requires a Future or a coroutine otherwise a TypeError is raised. The issue is that the websockets.connect call is an awaitable rather than a coroutine. asyncio.ensure_future did not gain support for awaitables until 3.5.2. This patch allows 3.5.0 and 3.5.1 to connect regardless of their python version.
328 lines
12 KiB
Python
328 lines
12 KiB
Python
# -*- coding: utf-8 -*-
|
|
|
|
"""
|
|
The MIT License (MIT)
|
|
|
|
Copyright (c) 2015-2017 Rapptz
|
|
|
|
Permission is hereby granted, free of charge, to any person obtaining a
|
|
copy of this software and associated documentation files (the "Software"),
|
|
to deal in the Software without restriction, including without limitation
|
|
the rights to use, copy, modify, merge, publish, distribute, sublicense,
|
|
and/or sell copies of the Software, and to permit persons to whom the
|
|
Software is furnished to do so, subject to the following conditions:
|
|
|
|
The above copyright notice and this permission notice shall be included in
|
|
all copies or substantial portions of the Software.
|
|
|
|
THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS
|
|
OR IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
|
|
FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
|
|
AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
|
|
LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING
|
|
FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER
|
|
DEALINGS IN THE SOFTWARE.
|
|
"""
|
|
|
|
from .state import AutoShardedConnectionState
|
|
from .client import Client
|
|
from .gateway import *
|
|
from .errors import ClientException, InvalidArgument
|
|
from . import compat
|
|
from .enums import Status
|
|
|
|
import asyncio
|
|
import logging
|
|
import websockets
|
|
import itertools
|
|
|
|
log = logging.getLogger(__name__)
|
|
|
|
class Shard:
|
|
def __init__(self, ws, client):
|
|
self.ws = ws
|
|
self._client = client
|
|
self.loop = self._client.loop
|
|
self._current = compat.create_future(self.loop)
|
|
self._current.set_result(None) # we just need an already done future
|
|
|
|
@property
|
|
def id(self):
|
|
return self.ws.shard_id
|
|
|
|
@asyncio.coroutine
|
|
def poll(self):
|
|
try:
|
|
yield from self.ws.poll_event()
|
|
except ResumeWebSocket as e:
|
|
log.info('Got a request to RESUME the websocket at Shard ID %s.', self.id)
|
|
coro = DiscordWebSocket.from_client(self._client, resume=True,
|
|
shard_id=self.id,
|
|
session=self.ws.session_id,
|
|
sequence=self.ws.sequence)
|
|
self.ws = yield from asyncio.wait_for(coro, timeout=180.0, loop=self.loop)
|
|
|
|
def get_future(self):
|
|
if self._current.done():
|
|
self._current = compat.create_task(self.poll(), loop=self.loop)
|
|
|
|
return self._current
|
|
|
|
@asyncio.coroutine
|
|
def _ensure_coroutine_connect(gateway, loop):
|
|
# In 3.5+ websockets.connect does not return a coroutine, but an awaitable.
|
|
# The problem is that in 3.5.0 and in some cases 3.5.1, asyncio.ensure_future and
|
|
# by proxy, asyncio.wait_for, do not accept awaitables, but rather futures or coroutines.
|
|
# By wrapping it up into this function we ensure that it's in a coroutine and not an awaitable
|
|
# even for 3.5.0 users.
|
|
ws = yield from websockets.connect(gateway, loop=loop, klass=DiscordWebSocket)
|
|
return ws
|
|
|
|
class AutoShardedClient(Client):
|
|
"""A client similar to :class:`Client` except it handles the complications
|
|
of sharding for the user into a more manageable and transparent single
|
|
process bot.
|
|
|
|
When using this client, you will be able to use it as-if it was a regular
|
|
:class:`Client` with a single shard when implementation wise internally it
|
|
is split up into multiple shards. This allows you to not have to deal with
|
|
IPC or other complicated infrastructure.
|
|
|
|
It is recommended to use this client only if you have surpassed at least
|
|
1000 guilds.
|
|
|
|
If no :attr:`shard_count` is provided, then the library will use the
|
|
Bot Gateway endpoint call to figure out how many shards to use.
|
|
|
|
If a ``shard_ids`` parameter is given, then those shard IDs will be used
|
|
to launch the internal shards. Note that :attr:`shard_count` must be provided
|
|
if this is used. By default, when omitted, the client will launch shards from
|
|
0 to ``shard_count - 1``.
|
|
|
|
Attributes
|
|
------------
|
|
shard_ids: Optional[List[int]]
|
|
An optional list of shard_ids to launch the shards with.
|
|
"""
|
|
def __init__(self, *args, loop=None, **kwargs):
|
|
kwargs.pop('shard_id', None)
|
|
self.shard_ids = kwargs.pop('shard_ids', None)
|
|
super().__init__(*args, loop=loop, **kwargs)
|
|
|
|
if self.shard_ids is not None:
|
|
if self.shard_count is None:
|
|
raise ClientException('When passing manual shard_ids, you must provide a shard_count.')
|
|
elif not isinstance(self.shard_ids, (list, tuple)):
|
|
raise ClientException('shard_ids parameter must be a list or a tuple.')
|
|
|
|
self.connection = AutoShardedConnectionState(dispatch=self.dispatch, chunker=self._chunker,
|
|
syncer=self._syncer, http=self.http, loop=self.loop, **kwargs)
|
|
|
|
# instead of a single websocket, we have multiple
|
|
# the key is the shard_id
|
|
self.shards = {}
|
|
|
|
def _get_websocket(guild_id):
|
|
i = (guild_id >> 22) % self.shard_count
|
|
return self.shards[i].ws
|
|
|
|
self.connection._get_websocket = _get_websocket
|
|
self._still_sharding = True
|
|
|
|
@asyncio.coroutine
|
|
def _chunker(self, guild, *, shard_id=None):
|
|
try:
|
|
guild_id = guild.id
|
|
shard_id = shard_id or guild.shard_id
|
|
except AttributeError:
|
|
guild_id = [s.id for s in guild]
|
|
|
|
payload = {
|
|
'op': 8,
|
|
'd': {
|
|
'guild_id': guild_id,
|
|
'query': '',
|
|
'limit': 0
|
|
}
|
|
}
|
|
|
|
ws = self.shards[shard_id].ws
|
|
yield from ws.send_as_json(payload)
|
|
|
|
@asyncio.coroutine
|
|
def request_offline_members(self, *guilds):
|
|
"""|coro|
|
|
|
|
Requests previously offline members from the guild to be filled up
|
|
into the :attr:`Guild.members` cache. This function is usually not
|
|
called. It should only be used if you have the ``fetch_offline_members``
|
|
parameter set to ``False``.
|
|
|
|
When the client logs on and connects to the websocket, Discord does
|
|
not provide the library with offline members if the number of members
|
|
in the guild is larger than 250. You can check if a guild is large
|
|
if :attr:`Guild.large` is ``True``.
|
|
|
|
Parameters
|
|
-----------
|
|
\*guilds
|
|
An argument list of guilds to request offline members for.
|
|
|
|
Raises
|
|
-------
|
|
InvalidArgument
|
|
If any guild is unavailable or not large in the collection.
|
|
"""
|
|
if any(not g.large or g.unavailable for g in guilds):
|
|
raise InvalidArgument('An unavailable or non-large guild was passed.')
|
|
|
|
_guilds = sorted(guilds, key=lambda g: g.shard_id)
|
|
for shard_id, sub_guilds in itertools.groupby(_guilds, key=lambda g: g.shard_id):
|
|
sub_guilds = list(sub_guilds)
|
|
yield from self.connection.request_offline_members(sub_guilds, shard_id=shard_id)
|
|
|
|
@asyncio.coroutine
|
|
def pending_reads(self, shard):
|
|
try:
|
|
while self._still_sharding:
|
|
yield from shard.poll()
|
|
except asyncio.CancelledError:
|
|
pass
|
|
|
|
@asyncio.coroutine
|
|
def launch_shard(self, gateway, shard_id):
|
|
try:
|
|
ws = yield from asyncio.wait_for(_ensure_coroutine_connect(gateway, self.loop), loop=self.loop, timeout=180.0)
|
|
except Exception as e:
|
|
log.info('Failed to connect for shard_id: %s. Retrying...', shard_id)
|
|
yield from asyncio.sleep(5.0, loop=self.loop)
|
|
return (yield from self.launch_shard(gateway, shard_id))
|
|
|
|
ws.token = self.http.token
|
|
ws._connection = self.connection
|
|
ws._dispatch = self.dispatch
|
|
ws.gateway = gateway
|
|
ws.shard_id = shard_id
|
|
ws.shard_count = self.shard_count
|
|
|
|
try:
|
|
# OP HELLO
|
|
yield from asyncio.wait_for(ws.poll_event(), loop=self.loop, timeout=180.0)
|
|
yield from asyncio.wait_for(ws.identify(), loop=self.loop, timeout=180.0)
|
|
except asyncio.TimeoutError:
|
|
log.info('Timed out when connecting for shard_id: %s. Retrying...', shard_id)
|
|
yield from asyncio.sleep(5.0, loop=self.loop)
|
|
return (yield from self.launch_shard(gateway, shard_id))
|
|
|
|
# keep reading the shard while others connect
|
|
self.shards[shard_id] = ret = Shard(ws, self)
|
|
compat.create_task(self.pending_reads(ret), loop=self.loop)
|
|
yield from asyncio.sleep(5.0, loop=self.loop)
|
|
|
|
@asyncio.coroutine
|
|
def launch_shards(self):
|
|
if self.shard_count is None:
|
|
self.shard_count, gateway = yield from self.http.get_bot_gateway()
|
|
else:
|
|
gateway = yield from self.http.get_gateway()
|
|
|
|
self.connection.shard_count = self.shard_count
|
|
|
|
shard_ids = self.shard_ids if self.shard_ids else range(self.shard_count)
|
|
|
|
for shard_id in shard_ids:
|
|
yield from self.launch_shard(gateway, shard_id)
|
|
|
|
self._still_sharding = False
|
|
|
|
@asyncio.coroutine
|
|
def _connect(self):
|
|
yield from self.launch_shards()
|
|
|
|
while True:
|
|
pollers = [shard.get_future() for shard in self.shards.values()]
|
|
done, pending = yield from asyncio.wait(pollers, loop=self.loop, return_when=asyncio.FIRST_COMPLETED)
|
|
for f in done:
|
|
# we wanna re-raise to the main Client.connect handler if applicable
|
|
f.result()
|
|
|
|
@asyncio.coroutine
|
|
def close(self):
|
|
"""|coro|
|
|
|
|
Closes the connection to discord.
|
|
"""
|
|
if self.is_closed():
|
|
return
|
|
|
|
self._closed.set()
|
|
|
|
for vc in self.voice_clients:
|
|
try:
|
|
yield from vc.disconnect()
|
|
except:
|
|
pass
|
|
|
|
to_close = [shard.ws.close() for shard in self.shards.values()]
|
|
yield from asyncio.wait(to_close, loop=self.loop)
|
|
yield from self.http.close()
|
|
|
|
@asyncio.coroutine
|
|
def change_presence(self, *, game=None, status=None, afk=False, shard_id=None):
|
|
"""|coro|
|
|
|
|
Changes the client's presence.
|
|
|
|
The game parameter is a Game object (not a string) that represents
|
|
a game being played currently.
|
|
|
|
Parameters
|
|
----------
|
|
game: Optional[:class:`Game`]
|
|
The game being played. None if no game is being played.
|
|
status: Optional[:class:`Status`]
|
|
Indicates what status to change to. If None, then
|
|
:attr:`Status.online` is used.
|
|
afk: bool
|
|
Indicates if you are going AFK. This allows the discord
|
|
client to know how to handle push notifications better
|
|
for you in case you are actually idle and not lying.
|
|
shard_id: Optional[int]
|
|
The shard_id to change the presence to. If not specified
|
|
or ``None``, then it will change the presence of every
|
|
shard the bot can see.
|
|
|
|
Raises
|
|
------
|
|
InvalidArgument
|
|
If the ``game`` parameter is not :class:`Game` or None.
|
|
"""
|
|
|
|
if status is None:
|
|
status = 'online'
|
|
status_enum = Status.online
|
|
elif status is Status.offline:
|
|
status = 'invisible'
|
|
status_enum = Status.offline
|
|
else:
|
|
status_enum = status
|
|
status = str(status)
|
|
|
|
if shard_id is None:
|
|
for shard in self.shards.values():
|
|
yield from shard.ws.change_presence(game=game, status=status, afk=afk)
|
|
|
|
guilds = self.connection.guilds
|
|
else:
|
|
shard = self.shards[shard_id]
|
|
yield from shard.ws.change_presence(game=game, status=status, afk=afk)
|
|
guilds = [g for g in self.connection.guilds if g.shard_id == shard_id]
|
|
|
|
for guild in guilds:
|
|
me = guild.me
|
|
if me is None:
|
|
continue
|
|
|
|
me.game = game
|
|
me.status = status_enum
|