juni/aikaterna-cogs
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

[Seen] Improves DisK I/O for bots with larger server and member base (#53)

Browse Source 
* Removes 3.0 compatibility
Improves Disk IO by writting to config every 60 seconds
Also listen to `on_typing`, `on_message_edit`, `on_reaction_remove`, `on_reaction_add`

* Change from `on_message` to `on_message_without_command`
This commit is contained in:
Draper
2019-08-12 15:59:33 +01:00
committed by aikaterna
parent 3a7d27fbb3
commit fdecc629c5
2 changed files with 156 additions and 35 deletions
Download Patch File Download Diff File Expand all files Collapse all files

6
seen/__init__.py
View File

@@ -1,5 +1,7 @@
from .seen import Seen
def setup(bot):
bot.add_cog(Seen(bot))
async def setup(bot):
cog = Seen(bot)
await cog.initialize()
bot.add_cog(cog)

185
seen/seen.py
View File

@@ -1,42 +1,100 @@
import asyncio
import contextlib
import datetime
from typing import Union
import discord
import time
from datetime import datetime
from redbot.core import Config, commands
_SCHEMA_VERSION = 2
BaseCog = getattr(commands, "Cog", object)
listener = getattr(commands.Cog, "listener", None) # Trusty + Sinbad
if listener is None:
class Seen(commands.Cog):
"""Shows last time a user was seen in chat."""
def listener(name=None):
return lambda x: x
class Seen(BaseCog):
"""Shows last time a user was seen in chat"""
def __init__(self, bot):
self.bot = bot
self.config = Config.get_conf(self, 2784481001, force_registration=True)
default_member = {"member_seen": []}
default_global = dict(schema_version=1)
default_member = dict(seen=None)
self.config.register_global(**default_global)
self.config.register_member(**default_member)
self._cache = {}
self._task = self.bot.loop.create_task(self._save_to_config())
async def initialize(self):
asyncio.ensure_future(
self._migrate_config(
from_version=await self.config.schema_version(), to_version=_SCHEMA_VERSION
)
)
async def _migrate_config(self, from_version: int, to_version: int):
if from_version == to_version:
return
elif from_version < to_version:
all_guild_data = await self.config.all_members()
users_data = {}
for guild_id, guild_data in all_guild_data.items():
for user_id, user_data in guild_data.items():
for _, v in user_data.items():
if not v:
v = None
if user_id not in users_data:
users_data[guild_id][user_id] = {"seen": v}
else:
if (v and not users_data[guild_id][user_id]["seen"]) or (
v
and users_data[guild_id][user_id]["seen"]
and v > users_data[guild_id][user_id]["seen"]
):
users_data[guild_id][user_id] = {"seen": v}
group = self.config._get_base_group(self.config.MEMBER) # Bulk update to new scope
async with group.all() as new_data:
for guild_id, member_data in users_data.items():
new_data[guild_id] = member_data
# new schema is now in place
await self.config.schema_version.set(_SCHEMA_VERSION)
# migration done, now let's delete all the old stuff
await self.config.clear_all_members()
@commands.guild_only()
@commands.command(name="seen")
@commands.bot_has_permissions(embed_links=True)
async def _seen(self, ctx, author: discord.Member):
"""Shows last time a user was seen in chat"""
member_seen = await self.config.member(author).member_seen()
now = int(time.time())
try:
time_elapsed = int(now - member_seen)
except TypeError:
"""Shows last time a user was seen in chat."""
member_seen_config = await self.config.member(author).seen()
member_seen_cache = self._cache.get(author.guild.id, {}).get(author.id, None)
if not member_seen_cache and not member_seen_config:
embed = discord.Embed(
colour=discord.Color.red(), title="I haven't seen that user yet."
)
return await ctx.send(embed=embed)
if not member_seen_cache:
member_seen = member_seen_config
elif not member_seen_config:
member_seen = member_seen_cache
elif member_seen_cache > member_seen_config:
member_seen = member_seen_cache
elif member_seen_config > member_seen_cache:
member_seen = member_seen_config
else:
member_seen = member_seen_cache or member_seen_config
now = int(time.time())
time_elapsed = int(now - member_seen)
output = self._dynamic_time(time_elapsed)
if output[2] < 1:
ts = "just now"
else:
@@ -54,27 +112,88 @@ class Seen(BaseCog):
elif output[2] > 1:
ts += "{} minutes ago".format(output[2])
em = discord.Embed(colour=discord.Color.green())
avatar = author.avatar_url if author.avatar else author.default_avatar_url
avatar = author.avatar_url or author.default_avatar_url
em.set_author(name="{} was seen {}".format(author.display_name, ts), icon_url=avatar)
await ctx.send(embed=em)
def _dynamic_time(self, time_elapsed):
@staticmethod
def _dynamic_time(time_elapsed):
m, s = divmod(time_elapsed, 60)
h, m = divmod(m, 60)
d, h = divmod(h, 24)
return (d, h, m)
return d, h, m
@listener()
async def on_message(self, message):
if (
not isinstance(message.channel, discord.abc.PrivateChannel)
and self.bot.user.id != message.author.id
):
prefixes = await self.bot.get_prefix(message)
if not any(message.content.startswith(n) for n in prefixes):
author = message.author
ts = int(time.time())
try:
await self.config.member(author).member_seen.set(ts)
except AttributeError:
pass
@commands.Cog.listener()
async def on_message_without_command(self, message):
if message.guild:
if message.guild.id not in self._cache:
self._cache[message.guild.id] = {}
self._cache[message.guild.id][message.author.id] = int(time.time())
@commands.Cog.listener()
async def on_typing(
self,
channel: discord.abc.Messageable,
user: Union[discord.User, discord.Member],
when: datetime.datetime,
):
if user.guild:
if user.guild.id not in self._cache:
self._cache[user.guild.id] = {}
self._cache[user.guild.id][user.id] = int(time.time())
@commands.Cog.listener()
async def on_message_edit(self, before: discord.Message, after: discord.Message):
if after.guild:
if after.guild.id not in self._cache:
self._cache[after.guild.id] = {}
self._cache[after.guild.id][after.author.id] = int(time.time())
@commands.Cog.listener()
async def on_reaction_remove(
self, reaction: discord.Reaction, user: Union[discord.Member, discord.User]
):
if user.guild:
if user.guild.id not in self._cache:
self._cache[user.guild.id] = {}
self._cache[user.guild.id][user.id] = int(time.time())
@commands.Cog.listener()
async def on_reaction_add(
self, reaction: discord.Reaction, user: Union[discord.Member, discord.User]
):
if user.guild:
if user.guild.id not in self._cache:
self._cache[user.guild.id] = {}
self._cache[user.guild.id][user.id] = int(time.time())
def cog_unload(self):
self.bot.loop.create_task(self._clean_up())
async def _clean_up(self):
if self._task:
self._task.cancel()
if self._cache:
group = self.config._get_base_group(self.config.MEMBER) # Bulk update to config
async with group.all() as new_data:
for guild_id, member_data in self._cache.items():
if str(guild_id) not in new_data:
new_data[str(guild_id)] = {}
for member_id, seen in member_data.items():
new_data[str(guild_id)][str(member_id)] = {"seen": seen}
async def _save_to_config(self):
await self.bot.wait_until_ready()
with contextlib.suppress(asyncio.CancelledError):
while True:
users_data = self._cache.copy()
self._cache = {}
group = self.config._get_base_group(self.config.MEMBER) # Bulk update to config
async with group.all() as new_data:
for guild_id, member_data in users_data.items():
if str(guild_id) not in new_data:
new_data[str(guild_id)] = {}
for member_id, seen in member_data.items():
new_data[str(guild_id)][str(member_id)] = {"seen": seen}
await asyncio.sleep(60)