jarvis/jarvis-bot
jarvis/jarvis-bot
1
0
Fork 0
Code Issues Pull requests Projects Releases Packages Wiki Activity Actions

Migrate twitter, closes #104

Browse source
This commit is contained in:
Zeva Rose 2022-02-03 20:20:12 -07:00
parent e1917e870e
commit dd8b3be19a
1 changed files with 76 additions and 93 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

169
jarvis/cogs/twitter.py
View file

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