diff --git a/screenshot/common/__init__.py b/screenshot/common/__init__.py index 52b85ad..4462494 100644 --- a/screenshot/common/__init__.py +++ b/screenshot/common/__init__.py @@ -28,7 +28,7 @@ import asyncio import contextlib from datetime import datetime, timedelta, timezone -from typing import TYPE_CHECKING, AsyncGenerator, Dict, Literal, TypeVar +from typing import TYPE_CHECKING, Any, AsyncGenerator, Coroutine, Dict, Literal, TypeVar from redbot.core import commands @@ -194,7 +194,7 @@ async def get_screenshot_bytes_from_url( wait: int = 10, ) -> bytes: async with self.driver() as driver: - return await asyncio.to_thread( + thread: Coroutine[Any, Any, bytes] = asyncio.to_thread( lambda: self.take_screenshot_with_url( driver, url=url, @@ -203,3 +203,4 @@ async def get_screenshot_bytes_from_url( wait=wait, ) ) + return await asyncio.wait_for(thread, timeout=30.0) diff --git a/screenshot/core.py b/screenshot/core.py index 69b0dee..a46657f 100644 --- a/screenshot/core.py +++ b/screenshot/core.py @@ -193,6 +193,13 @@ async def screenshot(self, ctx: commands.Context, url: URLConverter, *, flags: s allowed_mentions=discord.AllowedMentions(replied_user=False), ) raise commands.CheckFailure() + except asyncio.TimeoutError: + await ctx.send( + "Timed out waiting for the website to load.", + reference=ctx.message.to_reference(fail_if_not_exists=False), + allowed_mentions=discord.AllowedMentions(replied_user=False), + ) + raise commands.CheckFailure() except commands.UserFeedbackCheckFailure as error: if message := error.message: await ctx.send(