diff --git a/jarvis/cogs/twitter.py b/jarvis/cogs/twitter.py index 65fe265..f69369f 100644 --- a/jarvis/cogs/twitter.py +++ b/jarvis/cogs/twitter.py @@ -4,18 +4,16 @@ import logging import tweepy from bson import ObjectId -from discord import TextChannel -from discord.ext import commands -from discord.ext.tasks import loop -from discord.utils import find -from discord_slash import SlashContext, cog_ext -from discord_slash.model import SlashCommandOptionType as COptionType -from discord_slash.utils.manage_commands import create_choice, create_option -from discord_slash.utils.manage_components import ( - create_actionrow, - create_select, - create_select_option, - wait_for_component, +from dis_snek import InteractionContext, Permissions, Scale, Snake +from dis_snek.ext.tasks.task import Task +from dis_snek.ext.tasks.triggers import IntervalTrigger +from dis_snek.models.discord.channel import GuildText +from dis_snek.models.discord.components import ActionRow, Select, SelectOption +from dis_snek.models.snek.application_commands import ( + OptionTypes, + SlashCommandChoice, + slash_command, + slash_option, ) from jarvis.config import get_config @@ -25,10 +23,10 @@ from jarvis.utils.permissions import admin_or_permissions logger = logging.getLogger("discord") -class TwitterCog(commands.Cog): +class TwitterCog(Scale): """J.A.R.V.I.S. Twitter Cog.""" - def __init__(self, bot: commands.Bot): + def __init__(self, bot: Snake): self.bot = bot config = get_config() auth = tweepy.AppAuthHandler( @@ -39,14 +37,15 @@ class TwitterCog(commands.Cog): self._guild_cache = {} self._channel_cache = {} - @loop(seconds=30) + @Task.create(trigger=IntervalTrigger(minutes=1)) async def _tweets(self) -> None: twitters = Twitter.objects(active=True) handles = Twitter.objects.distinct("handle") twitter_data = {} for handle in handles: try: - twitter_data[handle] = self.api.user_timeline(screen_name=handle) + data = await asyncio.to_thread(self.api.user_timeline, screen_name=handle) + twitter_data[handle] = data except Exception as e: logger.error(f"Error with fetching: {e}") for twitter in twitters: @@ -55,16 +54,15 @@ class TwitterCog(commands.Cog): filter(lambda x: x.id > twitter.last_tweet, twitter_data[twitter.handle]) ) if tweets: + guild_id = twitter.guild + channel_id = twitter.channel tweets = sorted(tweets, key=lambda x: x.id) - if twitter.guild not in self._guild_cache: - self._guild_cache[twitter.guild] = await self.bot.fetch_guild(twitter.guild) + if guild_id not in self._guild_cache: + self._guild_cache[guild_id] = await self.bot.get_guild(guild_id) guild = self._guild_cache[twitter.guild] - if twitter.channel not in self._channel_cache: - channels = await guild.fetch_channels() - self._channel_cache[twitter.channel] = find( - lambda x: x.id == twitter.channel, channels - ) - channel = self._channel_cache[twitter.channel] + if channel_id not in self._channel_cache: + self._channel_cache[channel_id] = await guild.fetch_channel(channel_id) + channel = self._channel_cache[channel_id] for tweet in tweets: retweet = "retweeted_status" in tweet.__dict__ if retweet and not twitter.retweets: @@ -81,46 +79,37 @@ class TwitterCog(commands.Cog): except Exception as e: logger.error(f"Error with tweets: {e}") - @cog_ext.cog_subcommand( - base="twitter", - base_description="Twitter commands", - name="follow", - description="Follow a Twitter account", + @slash_command(name="twitter", sub_cmd_name="follow", description="Follow a Twitter acount") + @slash_option( + name="handle", description="Twitter account", option_type=OptionTypes.STRING, required=True + ) + @slash_option( + name="channel", + description="Channel to post tweets to", + option_type=OptionTypes.CHANNEL, + required=True, + ) + @slash_option( + name="retweets", + description="Mirror re-tweets?", + option_type=OptionTypes.STRING, + required=False, options=[ - create_option( - name="handle", - description="Twitter account", - option_type=COptionType.STRING, - required=True, - ), - create_option( - name="channel", - description="Channel to post tweets into", - option_type=COptionType.CHANNEL, - required=True, - ), - create_option( - name="retweets", - description="Mirror re-tweets?", - option_type=COptionType.STRING, - required=False, - choices=[ - create_choice(name="Yes", value="Yes"), - create_choice(name="No", value="No"), - ], - ), + SlashCommandChoice(name="Yes", value="Yes"), + SlashCommandChoice(name="No", value="No"), ], ) - @admin_or_permissions(manage_guild=True) + @admin_or_permissions(Permissions.MANAGE_GUILD) async def _twitter_follow( - self, ctx: SlashContext, handle: str, channel: TextChannel, retweets: str = "Yes" + self, ctx: InteractionContext, handle: str, channel: GuildText, retweets: str = "Yes" ) -> None: + handle = handle.lower() retweets = retweets == "Yes" if len(handle) > 15: await ctx.send("Invalid Twitter handle", hidden=True) return - if not isinstance(channel, TextChannel): + if not isinstance(channel, GuildText): await ctx.send("Channel must be a text channel", hidden=True) return @@ -155,13 +144,9 @@ class TwitterCog(commands.Cog): await ctx.send(f"Now following `@{handle}` in {channel.mention}") - @cog_ext.cog_subcommand( - base="twitter", - name="unfollow", - description="Unfollow Twitter accounts", - ) - @admin_or_permissions(manage_guild=True) - async def _twitter_unfollow(self, ctx: SlashContext) -> None: + @slash_command(name="twitter", sub_cmd_name="unfollow", description="Unfollow Twitter accounts") + @admin_or_permissions(Permissions.MANAGE_GUILD) + async def _twitter_unfollow(self, ctx: InteractionContext) -> None: twitters = Twitter.objects(guild=ctx.guild.id) if not twitters: await ctx.send("You need to follow a Twitter account first", hidden=True) @@ -170,14 +155,14 @@ class TwitterCog(commands.Cog): options = [] handlemap = {str(x.id): x.handle for x in twitters} for twitter in twitters: - option = create_select_option(label=twitter.handle, value=str(twitter.id)) + option = SelectOption(label=twitter.handle, value=str(twitter.id)) options.append(option) - select = create_select( + select = Select( options=options, custom_id="to_delete", min_values=1, max_values=len(twitters) ) - components = [create_actionrow(select)] + components = [ActionRow(select)] block = "\n".join(x.handle for x in twitters) message = await ctx.send( content=f"You are following the following accounts:\n```\n{block}\n```\n\n" @@ -186,19 +171,19 @@ class TwitterCog(commands.Cog): ) try: - context = await wait_for_component( - self.bot, + context = await self.bot.wait_for_component( check=lambda x: ctx.author.id == x.author.id, messages=message, timeout=60 * 5, ) - for to_delete in context.selected_options: + for to_delete in context.context.values: _ = Twitter.objects(guild=ctx.guild.id, id=ObjectId(to_delete)).delete() for row in components: for component in row["components"]: component["disabled"] = True - block = "\n".join(handlemap[x] for x in context.selected_options) - await context.edit_origin( + + block = "\n".join(handlemap[x] for x in context.context.values) + await context.context.edit_origin( content=f"Unfollowed the following:\n```\n{block}\n```", components=components ) except asyncio.TimeoutError: @@ -207,25 +192,21 @@ class TwitterCog(commands.Cog): component["disabled"] = True await message.edit(components=components) - @cog_ext.cog_subcommand( - base="twitter", + @slash_command( + name="twitter", sub_cmd_name="retweets", description="Modify followed Twitter accounts" + ) + @slash_option( name="retweets", - description="Modify followed Twitter accounts", + description="Mirror re-tweets?", + option_type=OptionTypes.STRING, + required=False, options=[ - create_option( - name="retweets", - description="Mirror re-tweets?", - option_type=COptionType.STRING, - required=True, - choices=[ - create_choice(name="Yes", value="Yes"), - create_choice(name="No", value="No"), - ], - ), + SlashCommandChoice(name="Yes", value="Yes"), + SlashCommandChoice(name="No", value="No"), ], ) - @admin_or_permissions(manage_guild=True) - async def _twitter_modify(self, ctx: SlashContext, retweets: str) -> None: + @admin_or_permissions(Permissions.MANAGE_GUILD) + async def _twitter_modify(self, ctx: InteractionContext, retweets: str) -> None: retweets = retweets == "Yes" twitters = Twitter.objects(guild=ctx.guild.id) if not twitters: @@ -234,14 +215,14 @@ class TwitterCog(commands.Cog): options = [] for twitter in twitters: - option = create_select_option(label=twitter.handle, value=str(twitter.id)) + option = SelectOption(label=twitter.handle, value=str(twitter.id)) options.append(option) - select = create_select( + select = Select( options=options, custom_id="to_update", min_values=1, max_values=len(twitters) ) - components = [create_actionrow(select)] + components = [ActionRow(select)] block = "\n".join(x.handle for x in twitters) message = await ctx.send( content=f"You are following the following accounts:\n```\n{block}\n```\n\n" @@ -250,22 +231,24 @@ class TwitterCog(commands.Cog): ) try: - context = await wait_for_component( - self.bot, + context = await self.bot.wait_for_component( check=lambda x: ctx.author.id == x.author.id, messages=message, timeout=60 * 5, ) + handlemap = {str(x.id): x.handle for x in twitters} - for to_update in context.selected_options: + for to_update in context.context.values: t = Twitter.objects(guild=ctx.guild.id, id=ObjectId(to_update)).first() t.retweets = retweets t.save() + for row in components: for component in row["components"]: component["disabled"] = True - block = "\n".join(handlemap[x] for x in context.selected_options) - await context.edit_origin( + + block = "\n".join(handlemap[x] for x in context.context.values) + await context.context.edit_origin( content=( f"{'Unfollowed' if not retweets else 'Followed'} " "retweets from the following:" @@ -280,6 +263,6 @@ class TwitterCog(commands.Cog): await message.edit(components=components) -def setup(bot: commands.Bot) -> None: +def setup(bot: Snake) -> None: """Add TwitterCog to J.A.R.V.I.S.""" bot.add_cog(TwitterCog(bot))