diff --git a/cogs/teto_image_cog.py b/cogs/teto_image_cog.py index 774bba4..abe9b9d 100644 --- a/cogs/teto_image_cog.py +++ b/cogs/teto_image_cog.py @@ -1,14 +1,15 @@ import discord from discord.ext import commands import aiohttp +from typing import Union class TetoImageView(discord.ui.View): def __init__(self, cog): - super().__init__() + super().__init__(timeout=180.0) self.cog = cog @discord.ui.button(label='Show Another Image', style=discord.ButtonStyle.primary) - async def show_another_image(self, interaction: discord.Interaction, button: discord.ui.Button): + async def show_another_image(self, interaction: discord.Interaction, _: discord.ui.Button): await self.cog.get_teto_image(interaction) class TetoImageCog(commands.Cog): @@ -17,37 +18,75 @@ class TetoImageCog(commands.Cog): self.teto_url = "https://slipstreamm.dev/teto" self.footer_url = "https://slipstreamm.dev/teto-web" + async def fetch_teto_image(self): + """Fetches a random Teto image and returns the URL.""" + async with aiohttp.ClientSession() as session: + headers = {'User-Agent': 'Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/91.0.4472.124 Safari/537.36'} + async with session.get(self.teto_url, headers=headers, allow_redirects=False) as response: + if response.status == 302: + image_path = response.headers.get('Location') + if image_path: + return f"https://slipstreamm.dev{image_path}" + return None + @commands.hybrid_command(name='tetoimage', description='Get a random image of Kasane Teto.') - async def get_teto_image(self, ctx): + async def get_teto_image(self, ctx_or_interaction: Union[commands.Context, discord.Interaction]): """Gets a random image of Kasane Teto.""" - async with ctx.typing(): - try: - async with aiohttp.ClientSession() as session: - headers = {'User-Agent': 'Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/91.0.4472.124 Safari/537.36'} - async with session.get(self.teto_url, headers=headers, allow_redirects=False) as response: - if response.status == 302: - image_path = response.headers.get('Location') - if image_path: - image_url = f"https://slipstreamm.dev{image_path}" - embed = discord.Embed( - title="Random Teto Image", - description=f"Website: {self.footer_url}", - color=discord.Color.red() - ) - if image_url.startswith("http") or image_url.startswith("https"): - embed.set_image(url=image_url) - view = TetoImageView(self) - await ctx.send(embed=embed, view=view) - else: - await ctx.send("The retrieved URL was not valid.") - else: - await ctx.send("Could not get the image URL from the redirect.") + # Determine if this is a Context or Interaction + is_interaction = isinstance(ctx_or_interaction, discord.Interaction) + + # Handle typing indicator differently based on object type + if is_interaction: + await ctx_or_interaction.response.defer(thinking=True) + else: + async with ctx_or_interaction.typing(): + pass + + try: + image_url = await self.fetch_teto_image() + + if image_url: + embed = discord.Embed( + title="Random Teto Image", + description=f"Website: {self.footer_url}", + color=discord.Color.red() + ) + + if image_url.startswith("http") or image_url.startswith("https"): + embed.set_image(url=image_url) + view = TetoImageView(self) + + # Send response differently based on object type + if is_interaction: + if ctx_or_interaction.response.is_done(): + await ctx_or_interaction.followup.send(embed=embed, view=view) else: - await ctx.send(f"Unexpected response status: {response.status}") - except aiohttp.ClientConnectorError as e: - await ctx.send(f"Could not connect to the image service: {e}") - except Exception as e: - await ctx.send(f"An error occurred: {e}") + await ctx_or_interaction.response.send_message(embed=embed, view=view) + else: + await ctx_or_interaction.send(embed=embed, view=view) + else: + error_msg = "The retrieved URL was not valid." + await self._send_error(ctx_or_interaction, error_msg) + else: + error_msg = "Could not get the image URL from the redirect." + await self._send_error(ctx_or_interaction, error_msg) + + except aiohttp.ClientConnectorError as e: + error_msg = f"Could not connect to the image service: {e}" + await self._send_error(ctx_or_interaction, error_msg) + except Exception as e: + error_msg = f"An error occurred: {e}" + await self._send_error(ctx_or_interaction, error_msg) + + async def _send_error(self, ctx_or_interaction, error_msg): + """Helper method to send error messages based on context type.""" + if isinstance(ctx_or_interaction, discord.Interaction): + if ctx_or_interaction.response.is_done(): + await ctx_or_interaction.followup.send(error_msg) + else: + await ctx_or_interaction.response.send_message(error_msg) + else: + await ctx_or_interaction.send(error_msg) async def setup(bot): await bot.add_cog(TetoImageCog(bot))