Merge branch 'twitter_update' into 'main'

Twitter command style, backend, updates and bugfixes
Mainly focused on retweet filtering

See merge request stark-industries/j.a.r.v.i.s.!28

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

jarvis/cogs/starboard.py

@@ -178,7 +178,7 @@ class StarboardCog(commands.Cog):
         if not isinstance(message, SlashMessage):
             if message.startswith("https://"):
                 message = message.split("/")[-1]
-            message = await channel.fetch_message(message.id)
+            message = await channel.fetch_message(message)
 
         exists = Star.objects(
             message=message.id,

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 = list(filter(lambda x: x.id > twitter.last_tweet, twitter_data[twitter.handle]))
-                tweets = sorted(tweets, key=lambda x: x.id)
                 if tweets:
-                    guild = await self.bot.fetch_guild(twitter.guild)
+                    tweets = sorted(tweets, key=lambda x: x.id)
-                    channels = await guild.fetch_channels()
+                    if twitter.guild not in self._guild_cache:
-                    channel = find(lambda x: x.id == twitter.channel, channels)
+                        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 <t:{timestamp}:f>: {url}")
+                        verb = "re" if tweet.retweeted_status else ""
+                        await channel.send(f"`@{twitter.handle}` {verb}tweeted this at <t:{timestamp}:f>: {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="handle", description="Twitter account", option_type=COptionType.STRING, required=True),
-            create_option(name="channel", description="Channel to post tweets into", option_type=7, 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."""

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()
+        if not self.__db_loaded:
-        for item in config:
+            db = self.__mongo[self.mongo["database"]]
-            setattr(self, item.key, item.value)
+            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:

jarvis/db/models.py

@@ -125,7 +125,7 @@ class Mute(Document):
     active = BooleanField(default=True)
     user = SnowflakeField(required=True)
     admin = SnowflakeField(required=True)
-    duration = IntField(min_value=1, max_value=300, default=10)
+    duration = IntField(min_value=-1, max_value=300, default=10)
     guild = SnowflakeField(required=True)
     reason = StringField(max_length=100, required=True)
     created_at = DateTimeField(default=datetime.utcnow)
@@ -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)