From 0175fce442cf1ffd6e677a973f4fdd1288c23767 Mon Sep 17 00:00:00 2001 From: Zevaryx Date: Sat, 19 Feb 2022 20:10:05 -0700 Subject: [PATCH] Migrate twitter --- jarvis/__init__.py | 6 ++-- jarvis/cogs/twitter.py | 63 ++++++++++++++++++++++++++++-------------- poetry.lock | 4 +-- pyproject.toml | 2 +- 4 files changed, 48 insertions(+), 27 deletions(-) diff --git a/jarvis/__init__.py b/jarvis/__init__.py index ebc6693..4b743aa 100644 --- a/jarvis/__init__.py +++ b/jarvis/__init__.py @@ -47,8 +47,8 @@ class Jarvis(Snake): self, ctx: Context, error: Exception, *args: list, **kwargs: dict ) -> None: """Lepton on_command_error override.""" - guild = await jarvis.get_guild(DEFAULT_GUILD) - channel = guild.get_channel(DEFAULT_ERROR_CHANNEL) + guild = await jarvis.fetch_guild(DEFAULT_GUILD) + channel = await guild.fetch_channel(DEFAULT_ERROR_CHANNEL) error_time = datetime.utcnow().strftime("%d-%m-%Y %H:%M-%S.%f UTC") timestamp = int(datetime.now().timestamp()) timestamp = f"" @@ -104,7 +104,7 @@ async def on_ready() -> None: def run() -> None: """Run J.A.R.V.I.S.""" - connect(**jconfig.mongo["connect"], testing=jconfig.mongo["database"] == "jarvis") + connect(**jconfig.mongo["connect"], testing=jconfig.mongo["database"] != "jarvis") jconfig.get_db_config() for extension in utils.get_extensions(): diff --git a/jarvis/cogs/twitter.py b/jarvis/cogs/twitter.py index 7d16b92..dc7d875 100644 --- a/jarvis/cogs/twitter.py +++ b/jarvis/cogs/twitter.py @@ -2,8 +2,8 @@ import asyncio import tweepy -from bson import ObjectId from dis_snek import InteractionContext, Permissions, Scale, Snake +from dis_snek.client.utils.misc_utils import get 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 ( @@ -14,7 +14,7 @@ from dis_snek.models.snek.application_commands import ( ) from dis_snek.models.snek.command import check from jarvis_core.db import q -from jarvis_core.db.models import Twitter +from jarvis_core.db.models import TwitterAccount, TwitterFollow from jarvis import jconfig from jarvis.utils.permissions import admin_or_permissions @@ -68,7 +68,7 @@ class TwitterCog(Scale): return try: - account = (await asyncio.to_thread(self.api.get_user(screen_name=handle)))[0] + account = await asyncio.to_thread(self.api.get_user, screen_name=handle) latest_tweet = (await asyncio.to_thread(self.api.user_timeline, screen_name=handle))[0] except Exception: await ctx.send( @@ -76,42 +76,54 @@ class TwitterCog(Scale): ) return - count = len([i async for i in Twitter.find(guild=ctx.guild.id)]) + count = len([i async for i in TwitterFollow.find(q(guild=ctx.guild.id))]) if count >= 12: await ctx.send("Cannot follow more than 12 Twitter accounts", ephemeral=True) return - exists = Twitter.find_one(q(twitter_id=account.id, guild=ctx.guild.id)) + exists = await TwitterFollow.find_one(q(twitter_id=account.id, guild=ctx.guild.id)) if exists: await ctx.send("Twitter account already being followed in this guild", ephemeral=True) return - t = Twitter( - handle=account.screen_name, + ta = await TwitterAccount.find_one(q(twitter_id=account.id)) + if not ta: + ta = TwitterAccount( + handle=account.screen_name, + twitter_id=account.id, + last_tweet=latest_tweet.id, + ) + await ta.commit() + + tf = TwitterFollow( twitter_id=account.id, guild=ctx.guild.id, channel=channel.id, admin=ctx.author.id, - last_tweet=latest_tweet.id, retweets=retweets, ) - await t.commit() + await tf.commit() await ctx.send(f"Now following `@{handle}` in {channel.mention}") @slash_command(name="twitter", sub_cmd_name="unfollow", description="Unfollow Twitter accounts") @check(admin_or_permissions(Permissions.MANAGE_GUILD)) async def _twitter_unfollow(self, ctx: InteractionContext) -> None: - twitters = Twitter.objects(guild=ctx.guild.id) + t = TwitterFollow.find(q(guild=ctx.guild.id)) + twitters = [] + async for twitter in t: + twitters.append(twitter) if not twitters: await ctx.send("You need to follow a Twitter account first", ephemeral=True) return options = [] - handlemap = {str(x.id): x.handle for x in twitters} + handlemap = {} for twitter in twitters: - option = SelectOption(label=twitter.handle, value=str(twitter.id)) + account = await TwitterAccount.find_one(q(twitter_id=twitter.twitter_id)) + handlemap[str(twitter.twitter_id)] = account.handle + option = SelectOption(label=account.handle, value=str(twitter.twitter_id)) options.append(option) select = Select( @@ -119,7 +131,7 @@ class TwitterCog(Scale): ) components = [ActionRow(select)] - block = "\n".join(x.handle for x in twitters) + block = "\n".join(x for x in handlemap.values()) message = await ctx.send( content=f"You are following the following accounts:\n```\n{block}\n```\n\n" "Please choose accounts to unfollow", @@ -133,7 +145,8 @@ class TwitterCog(Scale): timeout=60 * 5, ) for to_delete in context.context.values: - _ = Twitter.objects(guild=ctx.guild.id, id=ObjectId(to_delete)).delete() + follow = get(twitters, guild=ctx.guild.id, twitter_id=int(to_delete)) + await follow.delete() for row in components: for component in row.components: component.disabled = True @@ -164,14 +177,20 @@ class TwitterCog(Scale): @check(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) + t = TwitterFollow.find(q(guild=ctx.guild.id)) + twitters = [] + async for twitter in t: + twitters.append(twitter) if not twitters: await ctx.send("You need to follow a Twitter account first", ephemeral=True) return options = [] + handlemap = {} for twitter in twitters: - option = SelectOption(label=twitter.handle, value=str(twitter.id)) + account = await TwitterAccount.find_one(q(twitter_id=twitter.id)) + handlemap[str(twitter.twitter_id)] = account.handle + option = SelectOption(label=account.handle, value=str(twitter.twitter_id)) options.append(option) select = Select( @@ -179,7 +198,7 @@ class TwitterCog(Scale): ) components = [ActionRow(select)] - block = "\n".join(x.handle for x in twitters) + block = "\n".join(x for x in handlemap.values()) 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", @@ -193,11 +212,13 @@ class TwitterCog(Scale): timeout=60 * 5, ) - handlemap = {str(x.id): x.handle for x in twitters} + handlemap = {} for to_update in context.context.values: - t = await Twitter.find_one(q(guild=ctx.guild.id, id=ObjectId(to_update)))() - t.retweets = retweets - t.save() + account = await TwitterAccount.find_one(q(twitter_id=int(to_update))) + handlemap[str(twitter.twitter_id)] = account.handle + t = get(twitters, guild=ctx.guild.id, twitter_id=int(to_update)) + t.update(q(retweets=True)) + await t.commit() for row in components: for component in row.components: diff --git a/poetry.lock b/poetry.lock index af0717b..e90ef29 100644 --- a/poetry.lock +++ b/poetry.lock @@ -214,7 +214,7 @@ plugins = ["setuptools"] [[package]] name = "jarvis-core" -version = "0.2.1" +version = "0.4.1" description = "" category = "main" optional = false @@ -232,7 +232,7 @@ umongo = "^3.1.0" type = "git" url = "https://git.zevaryx.com/stark-industries/jarvis/jarvis-core.git" reference = "main" -resolved_reference = "0e627eae725abb1e6f3766c5dc94bd80d0ac6702" +resolved_reference = "20f67444579d36d508def2b47ea826af9bbdd2d7" [[package]] name = "jedi" diff --git a/pyproject.toml b/pyproject.toml index a897256..98287ba 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -1,6 +1,6 @@ [tool.poetry] name = "jarvis" -version = "2.0.0a0" +version = "2.0.0a1" description = "J.A.R.V.I.S. admin bot" authors = ["Zevaryx "]