Merge pull request #1323 from python-discord/snakes-cleanup

Snakes cleanup

3 files changed, 38 insertions, 49 deletions

diff --git a/bot/exts/fun/snakes/_snakes_cog.py b/bot/exts/fun/snakes/_snakes_cog.py
index 05ab7bfd..82496629 100644
--- a/bot/exts/fun/snakes/_snakes_cog.py
+++ b/bot/exts/fun/snakes/_snakes_cog.py
@@ -11,8 +11,8 @@ from functools import partial
 from io import BytesIO
 from typing import Any
 
-import async_timeout
 from PIL import Image, ImageDraw, ImageFont
+from aiohttp import ClientTimeout
 from discord import Colour, Embed, File, Member, Message, Reaction
 from discord.errors import HTTPException
 from discord.ext.commands import Cog, CommandError, Context, bot_has_permissions, group
@@ -241,8 +241,11 @@ class Snakes(Cog):
         # Draw the text onto the final image
         draw = ImageDraw.Draw(full_image)
         for line in textwrap.wrap(description, 36):
-            draw.text([margin + 4, offset], line, font=CARD["font"])
-            offset += CARD["font"].getsize(line)[1]
+            draw.text((margin + 4, offset), line, font=CARD["font"])
+
+            _left, top, _right, bottom = CARD["font"].getbbox(line)
+            # Height of the text + 4px spacing
+            offset += bottom - top + 4
 
         # Get the image contents as a BufferIO object
         buffer = BytesIO()
@@ -279,8 +282,8 @@ class Snakes(Cog):
         if params is None:
             params = {}
 
-        async with async_timeout.timeout(10), self.bot.http_session.get(url, params=params) as response:
-            return await response.json()
+        async with self.bot.http_session.get(url, params=params, timeout=ClientTimeout(total=10)) as response:
+                return await response.json()
 
     def _get_random_long_message(self, messages: list[str], retries: int = 10) -> str:
         """
@@ -986,21 +989,28 @@ class Snakes(Cog):
         """
         # Get the snake data we need
         if not name:
-            name_obj = await self._get_snake_name()
-            name = name_obj["scientific"]
-            content = await self._get_snek(name)
+            for _ in range(3):
+                name_obj = await self._get_snake_name()
+                name = name_obj["scientific"]
+                content = await self._get_snek(name)
 
+                if len(content["image_list"]) > 0:
+                    break
         elif isinstance(name, dict):
             content = name
-
         else:
             content = await self._get_snek(name)
 
+        try:
+            image_url = content["image_list"][0]
+        except IndexError:
+            await ctx.send("No images found for this snake.")
+            return
+
         # Make the card
         async with ctx.typing():
-
             stream = BytesIO()
-            async with async_timeout.timeout(10), self.bot.http_session.get(content["image_list"][0]) as response:
+            async with self.bot.http_session.get(image_url, timeout=ClientTimeout(total=10)) as response:
                 stream.write(await response.read())
 
             stream.seek(0)
diff --git a/bot/utils/__init__.py b/bot/utils/__init__.py
index 2a4acd96..f1ae0e75 100644
--- a/bot/utils/__init__.py
+++ b/bot/utils/__init__.py
@@ -7,6 +7,7 @@ from datetime import UTC, datetime
 
 import discord
 from discord.ext.commands import BadArgument, Context
+from pydis_core.utils.scheduling import create_task
 
 from bot.constants import Client, Month
 from bot.utils.pagination import LinePaginator
@@ -55,52 +56,30 @@ async def disambiguate(
     choices = (f"{index}: {entry}" for index, entry in enumerate(entries, start=1))
 
     def check(message: discord.Message) -> bool:
-        return (
-            message.content.isdecimal()
-            and message.author == ctx.author
-            and message.channel == ctx.channel
-        )
+        return message.author == ctx.author and message.channel == ctx.channel
 
-    try:
-        if embed is None:
-            embed = discord.Embed()
+    if embed is None:
+        embed = discord.Embed()
 
-        coro1 = ctx.bot.wait_for("message", check=check, timeout=timeout)
-        coro2 = LinePaginator.paginate(
+    # Run the paginator in the background, this means it will continue to work after the user has
+    # selected an option, and stopping it wont prevent the user from selecting an option, both
+    # of which are fine and make implementation simpler.
+    create_task(
+        LinePaginator.paginate(
             choices, ctx, embed=embed, max_lines=entries_per_page,
-            empty=empty, max_size=6000, timeout=9000
+            empty=empty, max_size=6000, timeout=timeout
         )
+    )
 
-        # wait_for timeout will go to except instead of the wait_for thing as I expected
-        futures = [asyncio.ensure_future(coro1), asyncio.ensure_future(coro2)]
-        done, pending = await asyncio.wait(futures, return_when=asyncio.FIRST_COMPLETED, loop=ctx.bot.loop)
-
-        # :yert:
-        result = next(iter(done)).result()
-
-        # Pagination was canceled - result is None
-        if result is None:
-            for coro in pending:
-                coro.cancel()
-            raise BadArgument("Canceled.")
-
-        # Pagination was not initiated, only one page
-        if result.author == ctx.bot.user:
-            # Continue the wait_for
-            result = await next(iter(pending))
-
-        # Love that duplicate code
-        for coro in pending:
-            coro.cancel()
+    try:
+        message = await ctx.bot.wait_for("message", check=check, timeout=timeout)
     except asyncio.TimeoutError:
         raise BadArgument("Timed out.")
 
-    # Guaranteed to not error because of isdecimal() in check
-    index = int(result.content)
-
     try:
+        index = int(message.content)
         return entries[index - 1]
-    except IndexError:
+    except (ValueError, IndexError):
         raise BadArgument("Invalid choice.")
 
 
diff --git a/bot/utils/pagination.py b/bot/utils/pagination.py
index c6cd7d12..58115fd6 100644
--- a/bot/utils/pagination.py
+++ b/bot/utils/pagination.py
@@ -90,7 +90,7 @@ class LinePaginator(Paginator):
         cls, lines: Iterable[str], ctx: Context,
         embed: Embed, prefix: str = "", suffix: str = "",
         max_lines: int | None = None, max_size: int = 500, empty: bool = True,
-        restrict_to_user: User = None, timeout: int = 300, footer_text: str | None = None,
+        restrict_to_user: User = None, timeout: float = 300, footer_text: str | None = None,
         url: str | None = None, exception_on_empty_embed: bool = False
     ) -> None:
         """
@@ -308,7 +308,7 @@ class ImagePaginator(Paginator):
 
     @classmethod
     async def paginate(cls, pages: list[tuple[str, str]], ctx: Context, embed: Embed,
-                       prefix: str = "", suffix: str = "", timeout: int = 300,
+                       prefix: str = "", suffix: str = "", timeout: float = 300,
                        exception_on_empty_embed: bool = False) -> None:
         """
         Use a paginator and set of reactions to provide pagination over a set of title/image pairs.