Migrate twitter, closes #104

This commit is contained in:
Zeva Rose 2022-02-03 20:20:12 -07:00
parent e1917e870e
commit dd8b3be19a

View file

@ -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))