🔀 Merge pull request #329 from Lunarmagpie/sharding

🐛 ✨ Added VERY basic sharding support and fixed the gateway

Author: zunda-arrow

docs/README.md

@@ -17,7 +17,7 @@
 [![Code style: black](https://img.shields.io/badge/code%20style-black-000000.svg)](https://github.com/psf/black)
 
 # <img src="../assets/svg/pincer.svg" height="24px" alt="Pincer Logo"> Pincer
-The snappy asynchronous discord api wrapper API wrapper written with aiohttp & websockets.
+The snappy asynchronous discord api wrapper API wrapper written with aiohttp.
 
 | :exclamation: | The package is currently within Pre-Alpha phase |
 | ------------- | :---------------------------------------------- |

docs/api/pincer.rst

@@ -89,6 +89,10 @@ Exceptions
 
 .. autoexception:: UnavailableGuildError()
 
+.. autoexception:: TimeoutError()
+
+.. autoexception:: GatewayConnectionError()
+
 .. autoexception:: HTTPError()
 
 .. autoexception:: NotModifiedError()
@@ -149,6 +153,8 @@ Exception Hierarchy
                 - :exc:`InvalidTokenError`
                 - :exc:`HeartbeatError`
             - :exc:`UnavailableGuildError`
+            - :exc:`TimeoutError`
+            - :exc:`GatewayConnectionError`
             - :exc:`HTTPError`
                 - :exc:`NotModifiedError`
                 - :exc:`BadRequestError`

pincer/_config.py

@@ -8,23 +8,23 @@
 @dataclass(repr=False)
 class GatewayConfig:
     """This file is to make maintaining the library and its gateway
-    configuration easier.
+    configuration easier. Leave compression blank for no compression.
     """
-    socket_base_url: str = "wss://gateway.discord.gg/"
+    MAX_RETRIES: int = 5
     version: int = 9
     encoding: str = "json"
-    compression: Optional[str] = "zlib-stream"
+    compression: str = "zlib-stream"
 
     @classmethod
-    def uri(cls) -> str:
+    def make_uri(cls, uri) -> str:
         """
         Returns
         -------
         :class:`str`:
             The GatewayConfig's uri.
         """
         return (
-            f"{cls.socket_base_url}"
+            f"{uri}"
             f"?v={cls.version}"
             f"&encoding={cls.encoding}"
         ) + f"&compress={cls.compression}" * cls.compressed()

pincer/client.py

@@ -4,25 +4,26 @@
 from __future__ import annotations
 
 import logging
-from asyncio import iscoroutinefunction, run, ensure_future
+from asyncio import iscoroutinefunction, ensure_future, create_task, get_event_loop
 from collections import defaultdict
+from functools import partial
 from importlib import import_module
 from inspect import isasyncgenfunction
 from typing import (
     Any,
     Dict,
     List,
     Optional,
+    Iterable,
     Tuple,
     Union,
     overload,
-    AsyncIterator,
     TYPE_CHECKING
 )
 from . import __package__
 from .commands import ChatCommandHandler
 from .core import HTTPClient
-from .core.gateway import Dispatcher
+from .core.gateway import GatewayInfo, Gateway
 from .exceptions import (
     InvalidEventName,
     TooManySetupArguments,
@@ -49,7 +50,7 @@
 from .utils.conversion import construct_client_dict, remove_none
 from .utils.event_mgr import EventMgr
 from .utils.extraction import get_index
-from .utils.insertion import should_pass_cls
+from .utils.insertion import should_pass_cls, should_pass_gateway
 from .utils.signature import get_params
 from .utils.types import CheckFunction
 from .utils.types import Coro
@@ -137,12 +138,10 @@ def decorator(func: Coro):
                 "already been registered"
             )
 
-        async def wrapper(cls, payload: GatewayDispatch):
+        async def wrapper(cls, gateway: Gateway, payload: GatewayDispatch):
             _log.debug("`%s` middleware has been invoked", call)
 
-            return await (
-                func(cls, payload) if should_pass_cls(func) else func(payload)
-            )
+            return await func(cls, gateway, payload)
 
         _events[call] = wrapper
         return wrapper
@@ -154,7 +153,7 @@ async def wrapper(cls, payload: GatewayDispatch):
     event_middleware(event)(middleware_)
 
 
-class Client(Dispatcher):
+class Client:
     """The client is the main instance which is between the programmer
     and the discord API.
 
@@ -200,24 +199,27 @@ def __init__(
         if isinstance(intents, Iterable):
             intents = sum(intents)
 
-        super().__init__(
-            token,
-            handlers={
-                # Gets triggered on all events
-                -1: self.payload_event_handler,
-                # Use this event handler for opcode 0.
-                0: self.event_handler,
-            },
-            intents=intents or Intents.all(),
-            reconnect=reconnect,
-        )
+        if intents is None:
+            intents = Intents.all()
+
+        self.intents = intents
+        self.reconnect = reconnect
+        self.token = token
 
         self.bot: Optional[User] = None
         self.received_message = received or "Command arrived successfully!"
         self.http = HTTPClient(token)
         self.throttler = throttler
         self.event_mgr = EventMgr()
 
+        async def get_gateway():
+            return GatewayInfo.from_dict(
+                await self.http.get("gateway/bot")
+            )
+
+        loop = get_event_loop()
+        self.gateway: GatewayInfo = loop.run_until_complete(get_gateway())
+
         # The guild and channel value is only registered if the Client has the GUILDS
         # intent.
         self.guilds: Dict[Snowflake, Optional[Guild]] = {}
@@ -236,7 +238,6 @@ def chat_commands(self) -> List[str]:
             cmd.app.name for cmd in ChatCommandHandler.register.values()
         ]
 
-
     @property
     def guild_ids(self) -> List[Snowflake]:
         """
@@ -341,7 +342,6 @@ def get_event_coro(name: str) -> List[Optional[Coro]]:
             ]
         )
 
-
     def load_cog(self, path: str, package: Optional[str] = None):
         """Load a cog from a string path, setup method in COG may
         optionally have a first argument which will contain the client!
@@ -461,7 +461,7 @@ async def unload_cog(self, path: str):
         await ChatCommandHandler(self).remove_commands(to_remove)
 
     @staticmethod
-    def execute_event(calls: List[Coro], *args, **kwargs):
+    def execute_event(calls: List[Coro], gateway: Gateway, *args, **kwargs):
         """Invokes an event.
 
         Parameters
@@ -484,19 +484,86 @@ def execute_event(calls: List[Coro], *args, **kwargs):
                     *remove_none(args),
                 )
 
+            if should_pass_gateway(call):
+                call_args = (call_args[0], gateway, *call_args[1:])
+
             ensure_future(call(*call_args, **kwargs))
 
     def run(self):
-        """Start the event listener."""
-        self.start_loop()
+        """Start the bot."""
+        loop = get_event_loop()
+        ensure_future(self.start_shard(0, 1), loop=loop)
+        loop.run_forever()
+
+    def run_autosharded(self):
+        """
+        Runs the bot with the amount of shards specified by the Discord gateway.
+        """
+        num_shards = self.gateway.shards
+        return self.run_shards(range(num_shards), num_shards)
+
+    def run_shards(self, shards: Iterable, num_shards: int):
+        """
+        Runs shards that you specify.
+
+        shards: Iterable
+            The shards to run.
+        num_shards: int
+            The total amount of shards.
+        """
+        loop = get_event_loop()
+
+        for shard in shards:
+            ensure_future(self.start_shard(shard, num_shards), loop=loop)
+
+        loop.run_forever()
+
+    async def start_shard(
+        self,
+        shard: int,
+        num_shards: int
+    ):
+        """|coro|
+        Starts a shard
+        This should not be run most of the time. ``run_shards`` and ``run_autosharded``
+        will likely do what you want.
+
+        shard : int
+            The number of the shard to start.
+        num_shards : int
+            The total number of shards.
+        """
+
+        gateway = Gateway(
+            self.token,
+            intents=self.intents,
+            url=self.gateway.url,
+            shard=shard,
+            num_shards=num_shards
+        )
+        await gateway.init_session()
+
+        gateway.append_handlers({
+            # Gets triggered on all events
+            -1: partial(self.payload_event_handler, gateway),
+            # Use this event handler for opcode 0.
+            0: partial(self.event_handler, gateway)
+        })
+
+        create_task(gateway.start_loop())
 
     def __del__(self):
         """Ensure close of the http client."""
         if hasattr(self, "http"):
-            run(self.http.close())
+            create_task(self.http.close())
 
     async def handle_middleware(
-        self, payload: GatewayDispatch, key: str, *args, **kwargs
+        self,
+        payload: GatewayDispatch,
+        key: str,
+        gateway: Gateway,
+        *args,
+        **kwargs
     ) -> Tuple[Optional[Coro], List[Any], Dict[str, Any]]:
         """|coro|
 
@@ -527,7 +594,7 @@ async def handle_middleware(
         next_call, arguments, params = ware, [], {}
 
         if iscoroutinefunction(ware):
-            extractable = await ware(self, payload, *args, **kwargs)
+            extractable = await ware(self, gateway, payload, *args, **kwargs)
 
             if not isinstance(extractable, tuple):
                 raise RuntimeError(
@@ -544,11 +611,16 @@ async def handle_middleware(
             return (next_call, ret_object)
 
         return await self.handle_middleware(
-            payload, next_call, *arguments, **params
+            payload, next_call, gateway, *arguments, **params
         )
 
     async def execute_error(
-        self, error: Exception, name: str = "on_error", *args, **kwargs
+        self,
+        error: Exception,
+        gateway: Gateway,
+        name: str = "on_error",
+        *args,
+        **kwargs
     ):
         """|coro|
 
@@ -567,11 +639,16 @@ async def execute_error(
             if ``call := self.get_event_coro(name)`` is :data:`False`
         """
         if calls := self.get_event_coro(name):
-            self.execute_event(calls, error, *args, **kwargs)
+            self.execute_event(calls, gateway, error, *args, **kwargs)
         else:
             raise error
 
-    async def process_event(self, name: str, payload: GatewayDispatch):
+    async def process_event(
+        self,
+        name: str,
+        payload: GatewayDispatch,
+        gateway: Gateway
+    ):
         """|coro|
 
         Processes and invokes an event and its middleware
@@ -587,16 +664,20 @@ async def process_event(self, name: str, payload: GatewayDispatch):
             what specifically happened.
         """
         try:
-            key, args = await self.handle_middleware(payload, name)
+            key, args = await self.handle_middleware(payload, name, gateway)
             self.event_mgr.process_events(key, args)
 
             if calls := self.get_event_coro(key):
-                self.execute_event(calls, args)
+                self.execute_event(calls, gateway, args)
 
         except Exception as e:
-            await self.execute_error(e)
+            await self.execute_error(e, gateway)
 
-    async def event_handler(self, _, payload: GatewayDispatch):
+    async def event_handler(
+        self,
+        gateway: Gateway,
+        payload: GatewayDispatch
+    ):
         """|coro|
 
         Handles all payload events with opcode 0.
@@ -611,9 +692,13 @@ async def event_handler(self, _, payload: GatewayDispatch):
             required data for the client to know what event it is and
             what specifically happened.
         """
-        await self.process_event(payload.event_name.lower(), payload)
+        await self.process_event(payload.event_name.lower(), payload, gateway)
 
-    async def payload_event_handler(self, _, payload: GatewayDispatch):
+    async def payload_event_handler(
+        self,
+        gateway: Gateway,
+        payload: GatewayDispatch
+    ):
         """|coro|
 
         Special event which activates the on_payload event.
@@ -628,7 +713,7 @@ async def payload_event_handler(self, _, payload: GatewayDispatch):
             required data for the client to know what event it is and
             what specifically happened.
         """
-        await self.process_event("payload", payload)
+        await self.process_event("payload", payload, gateway)
 
     @overload
     async def create_guild(
@@ -907,7 +992,6 @@ async def get_webhook(
         """
         return await Webhook.from_id(self, id, token)
 
-
     async def get_current_user(self) -> User:
         """|coro|
         The user object of the requester's account.
@@ -923,7 +1007,7 @@ async def get_current_user(self) -> User:
         """
         return User.from_dict(
             construct_client_dict(
-                self, 
+                self,
                 await self.http.get("users/@me")
             )
         )