From 09f38dc955e52603e8118d49134defef501f4d99 Mon Sep 17 00:00:00 2001 From: Zevaryx Date: Tue, 26 Oct 2021 10:44:38 -0600 Subject: [PATCH] Add Tweet filtering for retweets --- jarvis/__init__.py | 2 +- jarvis/cogs/twitter.py | 122 ++++++++++++++++++++++++++++++++++++----- jarvis/config.py | 18 ++++-- jarvis/db/models.py | 1 + 4 files changed, 123 insertions(+), 20 deletions(-) diff --git a/jarvis/__init__.py b/jarvis/__init__.py index fe0d348..15b64f1 100644 --- a/jarvis/__init__.py +++ b/jarvis/__init__.py @@ -41,7 +41,7 @@ jarvis = commands.Bot( slash = SlashCommand(jarvis, sync_commands=False, sync_on_cog_reload=True) jarvis_self = Process() -__version__ = "1.11.1" +__version__ = "1.11.2" @jarvis.event diff --git a/jarvis/cogs/twitter.py b/jarvis/cogs/twitter.py index 7877e6a..63cf4f4 100644 --- a/jarvis/cogs/twitter.py +++ b/jarvis/cogs/twitter.py @@ -9,7 +9,8 @@ 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.utils.manage_commands import create_option +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, @@ -33,22 +34,35 @@ class TwitterCog(commands.Cog): auth = tweepy.AppAuthHandler(config.twitter["consumer_key"], config.twitter["consumer_secret"]) self.api = tweepy.API(auth) self._tweets.start() + self._guild_cache = {} + self._channel_cache = {} @loop(seconds=30) async def _tweets(self) -> None: twitters = Twitter.objects(active=True) + handles = Twitter.objects.distinct("handle") + twitter_data = {} + for handle in handles: + twitter_data[handle] = self.api.user_timeline(screen_name=handle) for twitter in twitters: try: - tweets = self.api.user_timeline(screen_name=twitter.handle, since_id=twitter.last_tweet) - tweets = sorted(tweets, key=lambda x: x.id) + tweets = list(filter(lambda x: x.id > twitter.last_tweet, twitter_data[twitter.handle])) if tweets: - guild = await self.bot.fetch_guild(twitter.guild) - channels = await guild.fetch_channels() - channel = find(lambda x: x.id == twitter.channel, channels) + 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) + 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] for tweet in tweets: + if tweet.retweeted_status and not twitter.retweets: + continue timestamp = int(tweet.created_at.timestamp()) url = f"https://twitter.com/{twitter.handle}/status/{tweet.id}" - await channel.send(f"`@{twitter.handle}` tweeted this at : {url}") + verb = "re" if tweet.retweeted_status else "" + await channel.send(f"`@{twitter.handle}` {verb}tweeted this at : {url}") newest = max(tweets, key=lambda x: x.id) twitter.last_tweet = newest.id twitter.save() @@ -61,12 +75,27 @@ class TwitterCog(commands.Cog): name="follow", description="Follow a Twitter account", options=[ - create_option(name="handle", description="Twitter account", option_type=3, required=True), - create_option(name="channel", description="Channel to post tweets into", option_type=7, required=True), + 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")], + ), ], ) @admin_or_permissions(manage_guild=True) - async def _twitter_follow(self, ctx: SlashContext, handle: str, channel: TextChannel) -> None: + async def _twitter_follow( + self, ctx: SlashContext, handle: str, channel: TextChannel, retweets: str = "Yes" + ) -> None: + retweets = retweets == "Yes" if len(handle) > 15: await ctx.send("Invalid Twitter handle", hidden=True) return @@ -92,7 +121,12 @@ class TwitterCog(commands.Cog): return t = Twitter( - handle=handle, guild=ctx.guild.id, channel=channel.id, admin=ctx.author.id, last_tweet=latest_tweet.id + handle=handle, + guild=ctx.guild.id, + channel=channel.id, + admin=ctx.author.id, + last_tweet=latest_tweet.id, + retweets=retweets, ) t.save() @@ -121,7 +155,9 @@ class TwitterCog(commands.Cog): components = [create_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```", components=components + content=f"You are following the following accounts:\n```\n{block}\n```\n\n" + "Please choose accounts to unfollow", + components=components, ) try: @@ -133,7 +169,7 @@ class TwitterCog(commands.Cog): for row in components: for component in row["components"]: component["disabled"] = True - block = "\n".join(x.handle for x in context.selected_options) + block = "\n".join(x for x in context.selected_options) await context.edit_origin(content=f"Unfollowed the following:\n```\n{block}\n```", components=components) except asyncio.TimeoutError: for row in components: @@ -141,6 +177,66 @@ class TwitterCog(commands.Cog): component["disabled"] = True await message.edit(components=components) + @cog_ext.cog_subcommand( + base="twitter", + name="retweets", + description="Modify followed Twitter accounts", + 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")], + ), + ], + ) + @admin_or_permissions(manage_guild=True) + async def _twitter_modify(self, ctx: SlashContext, retweets: str) -> None: + retweets = retweets == "Yes" + twitters = Twitter.objects(guild=ctx.guild.id) + if not twitters: + await ctx.send("You need to follow a Twitter account first", hidden=True) + return + + options = [] + for twitter in twitters: + option = create_select_option(label=twitter.handle, value=str(twitter.id)) + options.append(option) + + select = create_select(options=options, custom_id="to_update", min_values=1, max_values=len(twitters)) + + components = [create_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" + f"Please choose which accounts to {'un' if not retweets else ''}follow retweets from", + components=components, + ) + + try: + context = await wait_for_component( + self.bot, check=lambda x: ctx.author.id == x.author.id, messages=message, timeout=60 * 5 + ) + for to_update in context.selected_options: + t = Twitter.objects(guild=ctx.guild.id, id=ObjectId()) + t.retweets = retweets + t.save() + for row in components: + for component in row["components"]: + component["disabled"] = True + block = "\n".join(x for x in context.selected_options) + await context.edit_origin( + content=f"{'Unfollowed' if not retweets else 'Followed'} retweets from the following:" + f"\n```\n{block}\n```", + components=components, + ) + except asyncio.TimeoutError: + for row in components: + for component in row["components"]: + component["disabled"] = True + await message.edit(components=components) + def setup(bot: commands.Bot) -> None: """Add TwitterCog to J.A.R.V.I.S.""" diff --git a/jarvis/config.py b/jarvis/config.py index df2c5fb..0338b29 100644 --- a/jarvis/config.py +++ b/jarvis/config.py @@ -1,8 +1,7 @@ """Load the config for J.A.R.V.I.S.""" +from pymongo import MongoClient from yaml import load -from jarvis.db.models import Config as DBConfig - try: from yaml import CLoader as Loader except ImportError: @@ -47,12 +46,17 @@ class Config(object): self.max_messages = max_messages self.gitlab_token = gitlab_token self.twitter = twitter + self.__db_loaded = False + self.__mongo = MongoClient(**self.mongo["connect"]) def get_db_config(self) -> None: """Load the database config objects.""" - config = DBConfig.objects() - for item in config: - setattr(self, item.key, item.value) + if not self.__db_loaded: + db = self.__mongo[self.mongo["database"]] + items = db.config.find() + for item in items: + setattr(self, item["key"], item["value"]) + self.__db_loaded = True @classmethod def from_yaml(cls, y: dict) -> "Config": @@ -68,7 +72,9 @@ def get_config(path: str = "config.yaml") -> Config: with open(path) as f: raw = f.read() y = load(raw, Loader=Loader) - return Config.from_yaml(y) + config = Config.from_yaml(y) + config.get_db_config() + return config def reload_config() -> None: diff --git a/jarvis/db/models.py b/jarvis/db/models.py index a4ba6bc..605f3f7 100644 --- a/jarvis/db/models.py +++ b/jarvis/db/models.py @@ -226,6 +226,7 @@ class Twitter(Document): channel = SnowflakeField(required=True) guild = SnowflakeField(required=True) last_tweet = SnowflakeField(required=True) + retweets = BooleanField(default=True) admin = SnowflakeField(required=True) created_at = DateTimeField(default=datetime.utcnow)