From 4cd33e24e322fc1323807e46eed7742b2f13ca8c Mon Sep 17 00:00:00 2001 From: kennnyshiwa <44236678+kennnyshiwa@users.noreply.github.com> Date: Fri, 1 May 2020 12:15:55 -0400 Subject: [PATCH] [rndstatus] update botstats and other minor improvements (#106) * [rndstauts] update botstats and other minor improvements * fix rndstatus * fix rndstatus again Co-authored-by: Kennnyshiwa --- rndstatus/rndstatus.py | 194 ++++++++++++++++++++++++----------------- 1 file changed, 112 insertions(+), 82 deletions(-) diff --git a/rndstatus/rndstatus.py b/rndstatus/rndstatus.py index c44fdc3..d5cc545 100644 --- a/rndstatus/rndstatus.py +++ b/rndstatus/rndstatus.py @@ -1,13 +1,20 @@ import re import discord from redbot.core import Config, commands, checks +from redbot.core.utils import AsyncIter from random import choice as rndchoice +from collections import defaultdict, Counter, Sequence import time +import contextlib +import asyncio +import logging + + +log = logging.getLogger("red.aikaterna-cogs.rndstatus") class RndStatus(commands.Cog): """Cycles random statuses or displays bot stats. - If a custom status is already set, it won't change it until it's back to none. [p]set game""" @@ -15,6 +22,10 @@ class RndStatus(commands.Cog): self.bot = bot self.last_change = None self.config = Config.get_conf(self, 2752521001, force_registration=True) + self._user_count = 0 + + self.user_task = asyncio.create_task(self._get_user_count()) + self.presence_task = asyncio.create_task(self.maybe_update_presence()) default_global = { "botstats": False, @@ -32,6 +43,24 @@ class RndStatus(commands.Cog): } self.config.register_global(**default_global) + def cog_unload(self): + self.user_task.cancel() + self.presence_task.cancel() + + async def _get_user_count(self, ): + await self.bot.wait_until_ready() + with contextlib.suppress(asyncio.CancelledError): + self._user_count = len(self.bot.users) + while True: + temp_data = defaultdict(set) + async for s in AsyncIter(self.bot.guilds): + if s.unavailable: + continue + async for m in AsyncIter(s.members): + temp_data["Unique Users"].add(m.id) + self._user_count = len(temp_data["Unique Users"]) + await asyncio.sleep(30) + @commands.group(autohelp=True) @commands.guild_only() @checks.is_owner() @@ -42,7 +71,6 @@ class RndStatus(commands.Cog): @rndstatus.command(name="set") async def _set(self, ctx, *statuses: str): """Sets Red's random statuses. - Accepts multiple statuses. Must be enclosed in double quotes in case of multiple words. Example: @@ -60,10 +88,10 @@ class RndStatus(commands.Cog): async def _streamer(self, ctx: commands.Context, *, streamer=None): """Set the streamer needed for streaming statuses. """ - + saved_streamer = await self.config.streamer() - if streamer == None: - return await ctx.send("Current Streamer: " + saved_streamer) + if streamer is None: + return await ctx.send(f"Current Streamer: {saved_streamer}" ) await self.config.streamer.set(streamer) await ctx.send( "Done. Redo this command with no parameters to see the current streamer." @@ -74,24 +102,23 @@ class RndStatus(commands.Cog): """Toggle for a bot stats status instead of random messages.""" botstats = await self.config.botstats() await self.config.botstats.set(not botstats) - await ctx.send("Botstats toggle: {}.".format(not botstats)) - if not botstats == False: + await ctx.send(f"Botstats toggle: {not botstats}.") + if botstats is not False: await self.bot.change_presence(activity=None) @rndstatus.command() async def delay(self, ctx, seconds: int): """Sets interval of random status switch. - Must be 20 or superior.""" if seconds < 20: seconds = 20 await self.config.delay.set(seconds) - await ctx.send(f"Interval set to {str(seconds)} seconds.") + await ctx.send(f"Interval set to {seconds} seconds.") @rndstatus.command() async def type(self, ctx, type: int): """Define the rndstatus type. - + Type list: 0 = Playing 1 = Streaming @@ -103,90 +130,93 @@ class RndStatus(commands.Cog): else: await ctx.send("Type must be between 0 and 3.") - @commands.Cog.listener() - async def on_message(self, message): - try: - current_game = str(message.guild.me.activity.name) - except AttributeError: - current_game = None - statuses = await self.config.statuses() - botstats = await self.config.botstats() - prefix = await self.bot.get_valid_prefixes() - streamer = await self.config.streamer() - url = "https://www.twitch.tv/" + streamer + async def maybe_update_presence(self): + await self.bot.wait_until_ready() + pattern = re.compile(rf"<@!?{self.bot.user.id}>") + delay = 0 + while True: + try: + cog_settings = await self.config.all() + guilds = self.bot.guilds + guild = next(g for g in guilds if not g.unavailable) + try: + current_game = str(guild.me.activity.name) + except AttributeError: + current_game = None + statuses = cog_settings["statuses"] + botstats = cog_settings["botstats"] + streamer = cog_settings["streamer"] + _type = cog_settings["type"] + delay = cog_settings["delay"] - if botstats: - me = self.bot.user - pattern = re.compile(rf"<@!?{me.id}>") - clean_prefix = pattern.sub(f"@{me.name}", prefix[0]) - total_users = sum(len(s.members) for s in self.bot.guilds) - servers = str(len(self.bot.guilds)) - botstatus = f"{clean_prefix}help | {total_users} users | {servers} servers" - if self.last_change == None: - type = await self.config.type() - if type == 1: - await self.bot.change_presence( - activity=discord.Streaming(name=botstatus, url=url) - ) - else: - await self.bot.change_presence( - activity=discord.Activity(name=botstatus, type=type) - ) - self.last_change = int(time.perf_counter()) - if message.author.id == self.bot.user.id: - return - delay = await self.config.delay() - if abs(self.last_change - int(time.perf_counter())) < int(delay): - return - if (current_game != str(botstatus)) or current_game == None: - type = await self.config.type() - if type == 1: - return await self.bot.change_presence( - activity=discord.Streaming(name=botstatus, url=url) - ) - else: - return await self.bot.change_presence( - activity=discord.Activity(name=botstatus, type=type) - ) + url = f"https://www.twitch.tv/{streamer}" + prefix = await self.bot.get_valid_prefixes() - if self.last_change == None: - if len(statuses) > 0 and (current_game in statuses or current_game == None): - new_status = self.random_status(message, statuses) - self.last_change = int(time.perf_counter()) - type = await self.config.type() - if type == 1: - await self.bot.change_presence( - activity=discord.Streaming(name=new_status, url=url) - ) - else: - await self.bot.change_presence( - activity=discord.Activity(name=new_status, type=type) - ) - - if message.author.id != self.bot.user.id: - delay = await self.config.delay() - if abs(self.last_change - int(time.perf_counter())) >= int(delay): - self.last_change = int(time.perf_counter()) - new_status = self.random_status(message, statuses) - if current_game != new_status: - if current_game in statuses or current_game == None: - type = await self.config.type() - if type == 1: + if botstats: + me = self.bot.user + clean_prefix = pattern.sub(f"@{me.name}", prefix[0]) + total_users = self._user_count + servers = str(len(self.bot.guilds)) + botstatus = f"{clean_prefix}help | {total_users} users | {servers} servers" + if self.last_change is None: + if _type == 1: + await self.bot.change_presence( + activity=discord.Streaming(name=botstatus, url=url) + ) + else: + await self.bot.change_presence( + activity=discord.Activity(name=botstatus, type=_type) + ) + self.last_change = int(time.perf_counter()) + if abs(self.last_change - int(time.perf_counter())) < int(delay): + return + if (current_game != str(botstatus)) or current_game is None: + if _type == 1: + return await self.bot.change_presence( + activity=discord.Streaming(name=botstatus, url=url) + ) + else: + return await self.bot.change_presence( + activity=discord.Activity(name=botstatus, type=_type) + ) + if self.last_change is None: + if len(statuses) > 0 and (current_game in statuses or current_game is None): + new_status = self.random_status(guild, statuses) + self.last_change = int(time.perf_counter()) + if _type == 1: await self.bot.change_presence( activity=discord.Streaming(name=new_status, url=url) ) else: await self.bot.change_presence( - activity=discord.Activity(name=new_status, type=type) + activity=discord.Activity(name=new_status, type=_type) ) + if abs(self.last_change - int(time.perf_counter())) >= int(delay): + self.last_change = int(time.perf_counter()) + new_status = self.random_status(guild, statuses) + if current_game != new_status: + if current_game in statuses or current_game is None: + if _type == 1: + await self.bot.change_presence( + activity=discord.Streaming(name=new_status, url=url) + ) + else: + await self.bot.change_presence( + activity=discord.Activity(name=new_status, type=_type) + ) + except asyncio.CancelledError: + break + except Exception as e: + log.exception(e, exc_info=e) + await asyncio.sleep(delay) - def random_status(self, msg, statuses): + def random_status(self, guild, statuses): try: - current = str(msg.guild.me.activity.name) + current = str(guild.me.activity.name) except AttributeError: current = None try: - new = str(msg.guild.me.activity.name) + new = str(guild.me.activity.name) except AttributeError: new = None if len(statuses) > 1: @@ -196,4 +226,4 @@ class RndStatus(commands.Cog): new = statuses[0] else: new = None - return new + return new \ No newline at end of file