diff --git a/bot.py b/bot.py index dc52a70..170e7e1 100644 --- a/bot.py +++ b/bot.py @@ -24,8 +24,6 @@ logging.basicConfig(level=logging.INFO, filename='bonfire.log') @bot.event async def on_ready(): - if not hasattr(bot, 'uptime'): - bot.uptime = pendulum.utcnow() if not hasattr(bot, 'owner'): appinfo = await bot.application_info() bot.owner = appinfo.owner @@ -33,7 +31,7 @@ async def on_ready(): @bot.event async def on_message(message): - if message.author.bot or utils.should_ignore(message): + if message.author.bot or utils.should_ignore(bot, message): return await bot.process_commands(message) @@ -49,9 +47,8 @@ async def process_command(ctx): server = ctx.message.guild command = ctx.command - command_usage = await utils.get_content('command_usage', key=command.qualified_name) - if command_usage is None: - command_usage = {'command': command.qualified_name} + command_usage = await bot.db.actual_load('command_usage', key=command.qualified_name) or \ + {'command': command.qualified_name} # Add one to the total usage for this command, basing it off 0 to start with (obviously) total_usage = command_usage.get('total_usage', 0) + 1 @@ -71,8 +68,7 @@ async def process_command(ctx): command_usage['server_usage'] = total_server_usage # Save all the changes - if not await utils.update_content('command_usage', command_usage, command.qualified_name): - await utils.add_content('command_usage', command_usage) + bot.db.save('command_usage', command_usage) @bot.event @@ -129,4 +125,7 @@ if __name__ == '__main__': for e in utils.extensions: bot.load_extension(e) + + bot.db = utils.DB() + bot.uptime = pendulum.utcnow() bot.run(utils.bot_token) diff --git a/cogs/admin.py b/cogs/admin.py index eadd2b5..489c4c7 100644 --- a/cogs/admin.py +++ b/cogs/admin.py @@ -5,7 +5,6 @@ from . import utils import discord import re -import rethinkdb as r valid_perms = [p for p in dir(discord.Permissions) if isinstance(getattr(discord.Permissions, p), property)] @@ -51,7 +50,7 @@ class Administration: EXAMPLE: !ignore #general RESULT: Bonfire will ignore commands sent in the general channel""" - key = str(ctx.message.guild.id) + key = ctx.message.guild.id converter = commands.converter.MemberConverter() member = None @@ -66,9 +65,7 @@ class Administration: await ctx.send("{} does not appear to be a member or channel!".format(member_or_channel)) return - settings = await utils.get_content('server_settings', key) - if settings is None: - settings = {} + settings = self.bot.db.load('server_settings', key=key, pluck='ignored') or {} ignored = settings.get('ignored', {'members': [], 'channels': []}) if member: if str(member.id) in ignored['members']: @@ -80,7 +77,7 @@ class Administration: else: ignored['members'].append(str(member.id)) fmt = "Ignoring {}".format(member.display_name) - elif channel: + else: if str(channel.id) in ignored['channels']: await ctx.send("I am already ignoring {}!".format(channel.mention)) return @@ -88,8 +85,12 @@ class Administration: ignored['channels'].append(str(channel.id)) fmt = "Ignoring {}".format(channel.mention) - update = {'ignored': ignored} - await utils.update_content('server_settings', update, key) + entry = { + 'ignored': ignored, + 'server_id': str(key) + } + + self.bot.db.save('server_settings', entry) await ctx.send(fmt) @commands.command() @@ -115,9 +116,7 @@ class Administration: await ctx.send("{} does not appear to be a member or channel!".format(member_or_channel)) return - settings = await utils.get_content('server_settings', key) - if settings is None: - settings = {} + settings = self.bot.db.load('server_settings', key=key) or {} ignored = settings.get('ignored', {'members': [], 'channels': []}) if member: if str(member.id) not in ignored['members']: @@ -126,7 +125,7 @@ class Administration: ignored['members'].remove(str(member.id)) fmt = "I am no longer ignoring {}".format(member.display_name) - elif channel: + else: if str(channel.id) not in ignored['channels']: await ctx.send("I'm not even ignoring {}!".format(channel.mention)) return @@ -134,8 +133,12 @@ class Administration: ignored['channels'].remove(str(channel.id)) fmt = "I am no longer ignoring {}".format(channel.mention) - update = {'ignored': ignored} - await utils.update_content('server_settings', update, key) + entry = { + 'ignored': ignored, + 'server_id': str(key) + } + + self.bot.db.save('server_settings', entry) await ctx.send(fmt) @commands.command(aliases=['alerts']) @@ -147,11 +150,12 @@ class Administration: EXAMPLE: !alerts #alerts RESULT: No more alerts spammed in #general!""" - key = str(ctx.message.guild.id) - entry = {'server_id': key, - 'notification_channel': str(channel.id)} - if not await utils.update_content('server_settings', entry, key): - await utils.add_content('server_settings', entry) + entry = { + 'server_id': str(ctx.message.guild.id), + 'notifications_channel': str(channel.id) + } + + self.bot.db.save('server_settings', entry) await ctx.send("I have just changed this server's 'notifications' channel" "\nAll notifications will now go to `{}`".format(channel)) @@ -168,12 +172,13 @@ class Administration: # So we base this channel on it's own and not from alerts # When mod logging becomes available, that will be kept to it's own channel if wanted as well on_off = True if re.search("(on|yes|true)", on_off.lower()) else False - key = str(ctx.message.guild.id) - entry = {'server_id': key, - 'join_leave': on_off} - if not await utils.update_content('server_settings', entry, key): - await utils.add_content('server_settings', entry) + entry = { + 'server_id': str(ctx.message.guild.id), + 'join_leave': on_off + } + + self.bot.db.save('server_settings',entry) fmt = "notify" if on_off else "not notify" await ctx.send("This server will now {} if someone has joined or left".format(fmt)) @@ -196,17 +201,15 @@ class Administration: else: key = str(ctx.message.guild.id) - entry = {'server_id': key, - 'nsfw_channels': [str(ctx.message.channel.id)]} - update = {'nsfw_channels': r.row['nsfw_channels'].append(str(ctx.message.channel.id))} + channels = self.bot.db.load('server_settings', key=key, pluck='nsfw_channels') or [] + channels.append(str(ctx.message.channel.id)) - server_settings = await utils.get_content('server_settings', key) - if server_settings and 'nsfw_channels' in server_settings.keys(): - await utils.update_content('server_settings', update, key) - elif server_settings: - await utils.update_content('server_settings', entry, key) - else: - await utils.add_content('server_settings', entry) + entry = { + 'server_id': key, + 'nsfw_channels': channels + } + + self.bot.db.save('server_settings', entry) await ctx.send("This channel has just been registered as 'nsfw'! Have fun you naughties ;)") @@ -217,27 +220,24 @@ class Administration: EXAMPLE: !nsfw remove RESULT: ;(""" - + channel = str(ctx.message.channel.id) if type(ctx.message.channel) is discord.DMChannel: key = 'DMs' else: key = str(ctx.message.guild.id) - server_settings = await utils.get_content('server_settings', key) - channel = str(ctx.message.channel.id) - try: - channels = server_settings.get('nsfw_channels', None) - if channel in channels: - channels.remove(channel) + channels = self.bot.db.load('server_settings', key=key, pluck='nsfw_channels') or [] + if channel in channels: + channels.remove(channel) - entry = {'nsfw_channels': channels} - await utils.update_content('server_settings', entry, key) - await ctx.send("This channel has just been unregistered as a nsfw channel") - return - except (TypeError, IndexError): - pass - - await ctx.send("This channel is not registered as a 'nsfw' channel!") + entry = { + 'server_id': key, + 'nsfw_channels': channels + } + self.bot.db.save('server_settings', entry) + await ctx.send("This channel has just been unregistered as a nsfw channel") + else: + await ctx.send("This channel is not registerred as a nsfw channel!") @commands.group(invoke_without_command=True) @commands.guild_only() @@ -259,11 +259,7 @@ class Administration: await ctx.send("That is not a valid command!") return - server_settings = await utils.get_content('server_settings', str(ctx.message.guild.id)) - try: - server_perms = server_settings['permissions'] - except (TypeError, IndexError, KeyError): - server_perms = {} + server_perms = self.bot.db.load('server_settings', key=ctx.message.guild.id, pluck='permissions') or {} perms_value = server_perms.get(cmd.qualified_name) if perms_value is None: @@ -351,12 +347,12 @@ class Administration: await ctx.send("This command cannot have custom permissions setup!") return - key = str(ctx.message.guild.id) - entry = {'server_id': key, - 'permissions': {cmd.qualified_name: perm_value}} + entry = { + 'server_id': str(ctx.message.guild.id), + 'permissions': {cmd.qualified_name: perm_value} + } - if not await utils.update_content('server_settings', entry, key): - await utils.add_content('server_settings', entry) + self.bot.db.save('server_settings', entry) await ctx.send("I have just added your custom permissions; " "you now need to have `{}` permissions to use the command `{}`".format(permissions, command)) @@ -377,8 +373,12 @@ class Administration: "That command does not exist! You can't have custom permissions on a non-existant command....") return - update = {'permissions': {cmd.qualified_name: None}} - await utils.update_content('server_settings', update, str(ctx.message.guild.id)) + entry = { + 'server_id': str(ctx.message.guild.id), + 'permissions': {cmd.qualified_name: None} + } + + self.bot.db.save('server_settings', entry) await ctx.send("I have just removed the custom permissions for {}!".format(cmd)) @commands.command() @@ -396,11 +396,12 @@ class Administration: if prefix.lower().strip() == "none": prefix = None - entry = {'server_id': key, - 'prefix': prefix} + entry = { + 'server_id': key, + 'prefix': prefix + } - if not await utils.update_content('server_settings', entry, key): - await utils.add_content('server_settings', entry) + self.bot.db.save('server_settings', entry) if prefix is None: fmt = "I have just cleared your custom prefix, the default prefix will have to be used now" @@ -417,14 +418,9 @@ class Administration: EXAMPLE: !rules 5 RESULT: Rule 5 is printed""" - server_settings = await utils.get_content('server_settings', str(ctx.message.guild.id)) - if server_settings is None: - await ctx.send("This server currently has no rules on it! I see you like to live dangerously...") - return + rules = self.bot.db.load('server_settings', key=ctx.message.guild.id, pluck='rules') - rules = server_settings.get('rules') - - if not rules or len(rules) == 0: + if rules is None: await ctx.send("This server currently has no rules on it! I see you like to live dangerously...") return @@ -452,17 +448,15 @@ class Administration: EXAMPLE: !rules add No fun allowed in this server >:c RESULT: No more fun...unless they break the rules!""" key = str(ctx.message.guild.id) - entry = {'server_id': key, - 'rules': [rule]} - update = {'rules': r.row['rules'].append(rule)} + rules = self.bot.db.load('server_settings', key=key, pluck='rules') or [] + rules.append(rule) - server_settings = await utils.get_content('server_settings', key) - if server_settings and 'rules' in server_settings.keys(): - await utils.update_content('server_settings', update, key) - elif server_settings: - await utils.update_content('server_settings', entry, key) - else: - await utils.add_content('server_settings', entry) + entry = { + 'server_id': key, + 'rules': rules + } + + self.bot.db.save('server_settings', entry) await ctx.send("I have just saved your new rule, use the rules command to view this server's current rules") @@ -475,11 +469,55 @@ class Administration: EXAMPLE: !rules delete 5 RESULT: Freedom from opression!""" - update = {'rules': r.row['rules'].delete_at(rule - 1)} - if not await utils.update_content('server_settings', update, str(ctx.message.guild.id)): - await ctx.send("That is not a valid rule number, try running the command again.") - else: + key = str(ctx.message.guild.id) + rules = self.bot.db.load('server_settings', key=key, pluck='rules') or [] + try: + rules.pop(rule - 1) + entry = { + 'server_id': key, + 'rules': rules + } + self.bot.db.save('server_settings', entry) await ctx.send("I have just removed that rule from your list of rules!") + except IndexError: + await ctx.send("That is not a valid rule number, try running the command again.") + + @commands.command() + @commands.guild_only() + @utils.custom_perms(manage_guild=True) + async def queuetype(self, ctx, new_type=None): + """Switches the song queue type for music + Choices are `user` or `song` queue + The `user` queue rotates off of a wait list, where people join the waitlist and the next song in their + playlist is the one that is played. + + The `song` queue rotates based on songs themselves, where people add a song to the server's playlist, + and these are rotated through. + + EXAMPLE: !queuetype user + RESULT: !queuetype """ + key = str(ctx.message.guild.id) + + if new_type is None: + cur_type = self.bot.db.load('server_settings', key=key, pluck='queue_type') or 'song' + await ctx.send("Current queue type is {}".format(cur_type)) + return + + new_type = new_type.lower().strip() + if new_type not in ['user', 'song']: + await ctx.send("Queue choices are either `user` or `song`. " + "Run `{}help queuetype` if you need more information".format(ctx.prefix)) + else: + entry = { + 'server_id': key, + 'queue_type': new_type + } + self.bot.db.save('server_settings', entry) + state = self.bot.get_cog('Music').voice_states.get(ctx.message.guild.id) + if state: + if new_type == "user" and not state.user_queue or new_type == "song" and state.user_queue: + state.switch_queue_type() + await ctx.send("Current queue type is now `{}`".format(new_type)) def setup(bot): diff --git a/cogs/dj.py b/cogs/dj.py new file mode 100644 index 0000000..c0b4b44 --- /dev/null +++ b/cogs/dj.py @@ -0,0 +1,65 @@ +from .voice_utilities import * +import discord + + +class DJEvents: + """A simple class to save our DJ objects, once someone is detected to have joined a channel, + their DJ information will automatically update""" + + def __init__(self, bot): + self.bot = bot + self.djs = {} + + async def on_ready(self): + for channel in [c for c in self.bot.get_all_channels() if isinstance(c, discord.VoiceChannel)]: + for member in [m for m in channel.members if not m.bot]: + if member.id not in self.djs: + dj = DJ(member, self.bot) + self.bot.loop.create_task(dj.resolve_playlist()) + self.djs[member.id] = dj + + async def on_voice_state_update(self, member, _, after): + if member and not member.bot and member.id not in self.djs: + dj = DJ(member, self.bot) + self.bot.loop.create_task(dj.resolve_playlist()) + self.djs[member.id] = dj + # Alternatively, if the bot has joined the channel and we never detected the members that are in the channel + # This most likely means the bot has just started up, lets get these user's ready too + if member and member.id == member.guild.me.id and after and after.channel: + for m in after.channel.members: + if not m.bot and m.id not in self.djs: + dj = DJ(m, self.bot) + self.bot.loop.create_task(dj.resolve_playlist()) + self.djs[m.id] = dj + + +class DJ(Playlist): + def __init__(self, member, bot): + super().__init__(bot) + self.member = member + self.playlists = [] + + async def get_next_entry(self, predownload_next=True): + if not self.entries: + return None + else: + entry = self.entries[0] + self.entries.rotate(-1) + return await entry.get_ready_future() + + async def resolve_playlist(self): + self.playlists = self.bot.db.load('user_playlists', key=self.member.id, pluck='playlists') or [] + self.clear() + + for pl in self.playlists: + if pl['active']: + for song in pl['songs']: + try: + await self.add_entry(song['url']) + except ExtractionError: + # For now, just silently ignore this + pass + + +def setup(bot): + bot.add_cog(DJEvents(bot)) diff --git a/cogs/events.py b/cogs/events.py index a6b1f99..46a85e5 100644 --- a/cogs/events.py +++ b/cogs/events.py @@ -44,10 +44,10 @@ class StatsUpdate: async with self.session.post(url, data=payload, headers=headers) as resp: log.info('bots.discord.pw statistics returned {} for {}'.format(resp.status, payload)) - async def on_server_join(self, server): + async def on_guild_join(self, _): self.bot.loop.create_task(self.update()) - async def on_server_leave(self, server): + async def on_guild_leave(self, _): self.bot.loop.create_task(self.update()) async def on_ready(self): @@ -55,12 +55,12 @@ class StatsUpdate: async def on_member_join(self, member): guild = member.guild - server_settings = await config.get_content('server_settings', str(guild.id)) + server_settings = self.bot.db.load('server_settings', key=str(guild.id)) try: join_leave_on = server_settings['join_leave'] if join_leave_on: - channel_id = server_settings.get('notification_channel') or member.guild.id + channel_id = server_settings.get('notifications_channel') or member.guild.id else: return except (IndexError, TypeError, KeyError): @@ -74,12 +74,12 @@ class StatsUpdate: async def on_member_remove(self, member): guild = member.guild - server_settings = await config.get_content('server_settings', str(guild.id)) + server_settings = self.bot.db.load('server_settings', key=str(guild.id)) try: join_leave_on = server_settings['join_leave'] if join_leave_on: - channel_id = server_settings.get('notification_channel') or member.guild.id + channel_id = server_settings.get('notifications_channel') or member.guild.id else: return except (IndexError, TypeError, KeyError): diff --git a/cogs/hangman.py b/cogs/hangman.py index 59f400c..bdf1ebb 100644 --- a/cogs/hangman.py +++ b/cogs/hangman.py @@ -8,6 +8,7 @@ import re import random import asyncio + class Game: def __init__(self, word): self.word = word @@ -142,12 +143,17 @@ class Hangman: return try: - msg = await ctx.message.author.send("Please respond with a phrase you would like to use for your hangman game in **{}**\n\nPlease keep phrases less than 20 characters".format(ctx.message.guild.name)) + msg = await ctx.message.author.send( + "Please respond with a phrase you would like to use for your hangman game in **{}**\n\nPlease keep phrases less than 20 characters".format( + ctx.message.guild.name)) except discord.Forbidden: - await ctx.send("I can't message you {}! Please allow DM's so I can message you and ask for the hangman phrase you want to use!".format(ctx.message.author.display_name)) + await ctx.send( + "I can't message you {}! Please allow DM's so I can message you and ask for the hangman phrase you want to use!".format( + ctx.message.author.display_name)) return - await ctx.send("I have DM'd you {}, please respond there with the phrase you would like to setup".format(ctx.message.author.display_name)) + await ctx.send("I have DM'd you {}, please respond there with the phrase you would like to setup".format( + ctx.message.author.display_name)) def check(m): return m.channel == msg.channel and len(m.content) < 20 @@ -155,12 +161,14 @@ class Hangman: try: msg = await self.bot.wait_for('message', check=check, timeout=60) except asyncio.TimeoutError: - await ctx.send("You took too long! Please look at your DM's next to as that's where I'm asking for the phrase you want to use") + await ctx.send( + "You took too long! Please look at your DM's next to as that's where I'm asking for the phrase you want to use") return forbidden_phrases = ['stop', 'delete', 'remove', 'end', 'create', 'start'] if msg.content in forbidden_phrases: - await ctx.send("Detected forbidden hangman phrase; current forbidden phrases are: \n{}".format("\n".join(forbidden_phrases))) + await ctx.send("Detected forbidden hangman phrase; current forbidden phrases are: \n{}".format( + "\n".join(forbidden_phrases))) return game = self.create(msg.content, ctx) diff --git a/cogs/images.py b/cogs/images.py index c13ef77..460a685 100644 --- a/cogs/images.py +++ b/cogs/images.py @@ -112,7 +112,7 @@ class Images: query = ' '.join(value for value in search if not re.search('&?filter_id=[0-9]+', value)) params = {'q': query} - nsfw = await utils.channel_is_nsfw(ctx.message.channel) + nsfw = await utils.channel_is_nsfw(ctx.message.channel, self.bot.db) # If this is a nsfw channel, we just need to tack on 'explicit' to the terms # Also use the custom filter that I have setup, that blocks some certain tags # If the channel is not nsfw, we don't need to do anything, as the default filter blocks explicit @@ -185,7 +185,7 @@ class Images: params = {'limit': 320, 'tags': tags} - nsfw = await utils.channel_is_nsfw(ctx.message.channel) + nsfw = await utils.channel_is_nsfw(ctx.message.channel, self.bot.db) # e621 by default does not filter explicit content, so tack on # safe/explicit based on if this channel is nsfw or not diff --git a/cogs/interaction.py b/cogs/interaction.py index ade1bc9..0f45f7a 100644 --- a/cogs/interaction.py +++ b/cogs/interaction.py @@ -117,7 +117,6 @@ class Interaction: return False return True - def start_battle(self, player1, player2): battles = self.battles.get(player1.guild.id, []) entry = { @@ -159,9 +158,9 @@ class Interaction: fmt = random.SystemRandom().choice(hugs) await ctx.send(fmt.format(user.display_name)) - @commands.group(invoke_without_command=True) + @commands.group(invoke_without_command=True, enabled=False) @commands.guild_only() - @commands.cooldown(1, 180, BucketType.user) + @commands.cooldown(1, 20, BucketType.user) @utils.custom_perms(send_messages=True) async def battle(self, ctx, player2: discord.Member): """Challenges the mentioned user to a battle @@ -226,10 +225,10 @@ class Interaction: # All we need to do is change what order the challengers are printed/added as a paramater if random.SystemRandom().randint(0, 1): await ctx.send(fmt.format(battleP1.mention, battleP2.mention)) - await utils.update_records('battle_records', battleP1, battleP2) + await utils.update_records('battle_records', self.bot.db, battleP1, battleP2) else: await ctx.send(fmt.format(battleP2.mention, battleP1.mention)) - await utils.update_records('battle_records', battleP2, battleP1) + await utils.update_records('battle_records', self.bot.db, battleP2, battleP1) @commands.command() @commands.guild_only() @@ -257,9 +256,9 @@ class Interaction: self.battling_off(player2=battleP2) await ctx.send("{} has chickened out! What a loser~".format(battleP2.mention)) - @commands.command() + @commands.command(enabled=False) @commands.guild_only() - @commands.cooldown(1, 180, BucketType.user) + @commands.cooldown(1, 10, BucketType.user) @utils.custom_perms(send_messages=True) async def boop(self, ctx, boopee: discord.Member = None, *, message=""): """Boops the mentioned person @@ -284,20 +283,15 @@ class Interaction: return key = str(booper.id) - boops = await utils.get_content('boops', key) - if boops is not None: - boops = boops['boops'] - # If the booper has never booped the member provided, assure it's 0 - amount = boops.get(str(boopee.id), 0) + 1 - boops[str(boopee.id)] = amount - - await utils.update_content('boops', {'boops': boops}, key) - else: - entry = {'member_id': str(booper.id), - 'boops': {str(boopee.id): 1}} - - await utils.add_content('boops', entry) - amount = 1 + boops = self.bot.db.load('boops', key=key, pluck='boops') or {} + amount = boops.get(str(boopee.id), 0) + 1 + entry = { + 'member_id': str(booper.id), + 'boops': { + str(boopee.id): amount + } + } + self.bot.db.save('boops', entry) fmt = "{0.mention} has just booped {1.mention}{3}! That's {2} times now!" await ctx.send(fmt.format(booper, boopee, amount, message)) diff --git a/cogs/links.py b/cogs/links.py index 755d6ab..6d59269 100644 --- a/cogs/links.py +++ b/cogs/links.py @@ -5,9 +5,7 @@ from . import utils from bs4 import BeautifulSoup as bs import discord -import random import re -import math class Links: @@ -29,7 +27,7 @@ class Links: url = "https://www.google.com/search" # Turn safe filter on or off, based on whether or not this is a nsfw channel - nsfw = await utils.channel_is_nsfw(ctx.message.channel) + nsfw = await utils.channel_is_nsfw(ctx.message.channel, self.bot.db) safe = 'off' if nsfw else 'on' params = {'q': query, diff --git a/cogs/misc.py b/cogs/misc.py index 9403f89..efa0066 100644 --- a/cogs/misc.py +++ b/cogs/misc.py @@ -3,9 +3,6 @@ from discord.ext import commands from . import utils -from bs4 import BeautifulSoup as bs -import subprocess -import glob import random import re import calendar @@ -42,7 +39,7 @@ class Miscallaneous: if command is None: for cmd in utils.get_all_commands(self.bot): - if not await cmd.can_run(ctx): + if not await cmd.can_run(ctx) or not cmd.enabled: continue cog = cmd.cog_name @@ -82,7 +79,7 @@ class Miscallaneous: except utils.CannotPaginate as e: await ctx.send(str(e)) else: - # Get the description for a command + # Get the description for a command description = command.help if description is not None: # Split into examples, results, and the description itself based on the string @@ -108,7 +105,6 @@ class Miscallaneous: await ctx.send(embed=embed) - @commands.command() @utils.custom_perms(send_messages=True) async def say(self, ctx, *, msg: str): @@ -123,49 +119,6 @@ class Miscallaneous: except: pass - @commands.command() - @utils.custom_perms(send_messages=True) - async def motd(self, ctx, *, date=None): - """This command can be used to print the current MOTD (Message of the day) - This will most likely not be updated every day, however messages will still be pushed to this every now and then - - EXAMPLE: !motd - RESULT: 'This is an example message of the day!'""" - if date is None: - motd = await utils.get_content('motd') - try: - # Lets set this to the first one in the list first - latest_motd = motd[0] - for entry in motd: - d = pendulum.parse(entry['date']) - - # Check if the date for this entry is newer than our currently saved latest entry - if d > pendulum.parse(latest_motd['date']): - latest_motd = entry - - date = latest_motd['date'] - motd = latest_motd['motd'] - # This will be hit if we do not have any entries for motd - except TypeError: - await ctx.send("No message of the day!") - else: - fmt = "Last updated: {}\n\n{}".format(date, motd) - await ctx.send(fmt) - else: - try: - motd = await utils.get_content('motd', str(pendulum.parse(date).date())) - date = motd['date'] - motd = motd['motd'] - fmt = "Message of the day for {}:\n\n{}".format(date, motd) - await ctx.send(fmt) - # This one will be hit if we return None for that day - except TypeError: - await ctx.send("No message of the day for {}!".format(date)) - # This will be hit if pendulum fails to parse the date passed - except ValueError: - now = pendulum.utcnow().to_date_string() - await ctx.send("Invalid date format! Try like {}".format(now)) - @commands.command() @utils.custom_perms(send_messages=True) async def calendar(self, ctx, month: str = None, year: int = None): @@ -255,7 +208,7 @@ class Miscallaneous: if hasattr(self.bot, 'uptime'): embed.add_field(name='Uptime', value=(pendulum.utcnow() - self.bot.uptime).in_words()) - memory_usage = self.process.memory_full_info().uss / 1024**2 + memory_usage = self.process.memory_full_info().uss / 1024 ** 2 cpu_usage = self.process.cpu_percent() / psutil.cpu_count() embed.add_field(name='Memory Usage', value='{:.2f} MiB'.format(memory_usage)) embed.add_field(name='CPU Usage', value='{}%'.format(cpu_usage)) @@ -300,18 +253,15 @@ class Miscallaneous: await ctx.send("Use this URL to add me to a server that you'd like!\n{}" .format(discord.utils.oauth_url(app_info.id, perms))) - @commands.command() + @commands.command(enabled=False) @utils.custom_perms(send_messages=True) async def joke(self, ctx): """Prints a random riddle EXAMPLE: !joke RESULT: An absolutely terrible joke.""" - joke = await utils.request('http://tambal.azurewebsites.net/joke/random') - if joke is not None and 'joke' in joke: - await ctx.send(joke.get('joke')) - else: - await ctx.send("Sorry, I'm not feeling funny right now...try later") + # Currently disabled until I can find a free API + pass @commands.command() @utils.custom_perms(send_messages=True) diff --git a/cogs/music.py b/cogs/music.py index 74e4ff2..a4667a0 100644 --- a/cogs/music.py +++ b/cogs/music.py @@ -11,7 +11,7 @@ import asyncio import time import re import logging -import traceback +from collections import deque log = logging.getLogger() @@ -20,13 +20,16 @@ if not discord.opus.is_loaded(): class VoiceState: - def __init__(self, guild, bot): + def __init__(self, guild, bot, user_queue=False): self.guild = guild self.songs = Playlist(bot) + self.djs = deque() + self.dj = None self.current = None self.required_skips = 0 self.skip_votes = set() - self.audio_player = bot.loop.create_task(self.audio_player_task()) + self.user_queue = user_queue + self.loop = bot.loop self._volume = 50 @property @@ -50,40 +53,38 @@ class VoiceState: else: return self.voice.is_playing() or self.voice.is_paused() + def switch_queue_type(self): + self.songs.clear() + self.djs.clear() + self.dj = None + self.user_queue = not self.user_queue + self.skip() + + def get_dj(self, member): + for x in self.djs: + if x.member.id == member.id: + return x + def skip(self): self.skip_votes.clear() if self.playing: self.voice.stop() - def after(self): - self.current = None - - async def audio_player_task(self): - while True: - if self.playing: - await asyncio.sleep(1) - continue - song = self.songs.peek() - if song is None: - await asyncio.sleep(1) - continue - - try: - self.current = await self.songs.get_next_entry() - embed = self.current.to_embed() - embed.title = "Now playing!" - await song.channel.send(embed=embed) - except ExtractionError as e: - error = str(e).partition(" ")[2] - await song.channel.send("Failed to download {}!\nError: {}".format(song.title, error)) - continue - except discord.Forbidden: - pass - except: - await song.channel.send("Failed to download {}!".format(song.title)) - log.error(traceback.format_exc()) - continue + def after(self, _=None): + if self.user_queue: + self.djs.append(self.dj) + fut = asyncio.run_coroutine_threadsafe(self.play_next_song(), self.loop) + fut.result() + async def play_next_song(self): + self.skip_votes.clear() + try: + await self.next_song() + except ExtractionError: + # For now lets just silently continue in the queue + # Implementation to the music notifications channel will change what we do here + return await self.play_next_song() + if self.current: source = FFmpegPCMAudio( self.current.filename, before_options='-nostdin', @@ -92,6 +93,31 @@ class VoiceState: source = PCMVolumeTransformer(source, volume=self.volume) self.voice.play(source, after=self.after) self.current.start_time = time.time() + else: + # If we're here what we can assume is the following took place: + # 1) The queue type is `user` + # 2) Someone joined for the first time, starting off the queue + # 3) They don't have a song in their playlist ready yet + # So what we'll do here is just call this again a few seconds later + await asyncio.sleep(2) + return await self.play_next_song() + + async def next_song(self): + if not self.user_queue: + self.current = await self.songs.get_next_entry() + else: + try: + self.dj = self.djs.popleft() + except IndexError: + self.current = None + else: + self.current = await self.dj.get_next_entry() + self.djs.rotate(-1) + if self.current is None: + self.djs.remove(self.dj) + await self.next_song() + else: + self.current.requester = self.dj.member class Music: @@ -106,16 +132,6 @@ class Music: self.downloader = down self.bot.downloader = down - def __unload(self): - # If this is unloaded, cancel all players and disconnect from all channels - for state in self.voice_states.values(): - try: - state.audio_player.cancel() - if state.voice: - self.bot.loop.create_task(state.voice.disconnect()) - except: - pass - async def queue_embed_task(self, state, channel, author): index = 0 message = None @@ -136,7 +152,10 @@ class Music: while True: # Get the current queue (It might change while we're doing this) # So do this in the while loop - queue = state.songs.entries + if state.user_queue: + queue = state.djs + else: + queue = state.songs.entries count = len(queue) # This means the last song was removed if count == 0: @@ -144,8 +163,13 @@ class Music: break # Get the current entry entry = queue[index] + dj = None + if state.user_queue: + dj = entry + entry = entry.peek() # Get the entry's embed embed = entry.to_embed() + # Set the embed's title to indicate the amount of things in the queue count = len(queue) embed.title = "Current Queue [{}/{}]".format(index + 1, count) @@ -201,33 +225,45 @@ class Music: elif '\u2b06' in reaction.emoji: # A second check just to make sure, as well as ensuring index is higher than 0 if author.guild_permissions.kick_members and index > 0: - if entry != queue[index]: + if dj and dj != queue[index]: + fmt = "`Error: Position of this entry has changed, cannot complete your action`" + elif not dj and entry != queue[index]: fmt = "`Error: Position of this entry has changed, cannot complete your action`" else: # Remove the current entry del queue[index] # Add it one position higher - queue.insert(index - 1, entry) + if state.user_queue: + queue.insert(index - 1, dj) + else: + queue.insert(index - 1, entry) # Lets move the index to look at the new place of the entry index -= 1 # If down is clicked elif '\u2b07' in reaction.emoji: # A second check just to make sure, as well as ensuring index is lower than last if author.guild_permissions.kick_members and index < (count - 1): - if entry != queue[index]: + if dj and dj != queue[index]: + fmt = "`Error: Position of this entry has changed, cannot complete your action`" + elif not dj and entry != queue[index]: fmt = "`Error: Position of this entry has changed, cannot complete your action`" else: # Remove the current entry del queue[index] # Add it one position lower - queue.insert(index + 1, entry) + if state.user_queue: + queue.insert(index + 1, dj) + else: + queue.insert(index + 1, entry) # Lets move the index to look at the new place of the entry index += 1 # If x is clicked elif '\u274c' in reaction.emoji: # A second check just to make sure if author.guild_permissions.kick_members or author == entry.requester: - if entry != queue[index]: + if dj and dj != queue[index]: + fmt = "`Error: Position of this entry has changed, cannot complete your action`" + elif not dj and entry != queue[index]: fmt = "`Error: Position of this entry has changed, cannot complete your action`" else: # Simply remove the entry in place @@ -261,16 +297,19 @@ class Music: async def add_entry(self, song, ctx): state = self.voice_states[ctx.message.guild.id] - entry, _ = await state.songs.add_entry(song, ctx) + entry, _ = await state.songs.add_entry(song) + if not state.playing: + await state.play_next_song() + entry.requester = ctx.message.author return entry - async def join_channel(self, channel): + async def join_channel(self, channel, text_channel): state = self.voice_states.get(channel.guild.id) log.info("Joining channel {} in guild {}".format(channel.id, channel.guild.id)) # Send a message letting the channel know we are attempting to join try: - msg = await channel.send("Trying to join channel {}...".format(channel.name)) + msg = await text_channel.send("Trying to join channel {}...".format(channel.name)) except discord.Forbidden: msg = None @@ -283,7 +322,9 @@ class Music: await channel.connect() # If we have connnected, create our voice state - self.voice_states[channel.guild.id] = VoiceState(channel.guild, self.bot) + queue_type = self.bot.db.load('server_settings', key=channel.guild.id, pluck='queue_type') + user_queue = queue_type == "user" + self.voice_states[channel.guild.id] = VoiceState(channel.guild, self.bot, user_queue=user_queue) # If we can send messages, edit it to let the channel know we have succesfully joined if msg: @@ -309,7 +350,7 @@ class Music: await channel.send("Sorry but I couldn't connect...try again?") return False - @commands.command(pass_context=True) + @commands.command() @commands.guild_only() @utils.custom_perms(send_messages=True) async def progress(self, ctx): @@ -356,7 +397,7 @@ class Music: "voice activation`".format(channel.name)) return False - return await self.join_channel(channel) + return await self.join_channel(channel, ctx.channel) @commands.command() @commands.guild_only() @@ -372,6 +413,10 @@ class Music: if ctx.message.guild.id not in self.voice_states: if not await ctx.invoke(self.join): return + if self.voice_states.get(ctx.message.guild.id).user_queue: + await ctx.send("The current queue type is the DJ queue. " + "Use the command {}dj to join this queue".format(ctx.prefix)) + return song = re.sub('[<>\[\]]', '', song) if len(song) == 11: @@ -382,8 +427,6 @@ class Music: try: entry = await self.add_entry(song, ctx) - except asyncio.TimeoutError: - await ctx.send("You took too long!") except LiveStreamError as e: await ctx.send(str(e)) except WrongEntryTypeError: @@ -396,7 +439,7 @@ class Music: # We want youtube_dl's error message, but just the first part, the actual "error" error = error[2] # This is colour formatting for the console...it's just going to show up as text on discord - error = error.strip("ERROR: ") + error = error.replace("ERROR: ", "") else: # This happens when the download just returns `None` error = error[0] @@ -409,10 +452,10 @@ class Music: embed = entry.to_embed() embed.title = "Enqueued song!" await ctx.send(embed=embed) - except (discord.Forbidden, discord.HTTPException): + except discord.Forbidden: pass - @commands.command(pass_context=True) + @commands.command() @commands.guild_only() @utils.custom_perms(kick_members=True) async def volume(self, ctx, value: int = None): @@ -431,7 +474,7 @@ class Music: state.volume = value await ctx.send('Set the volume to {:.0%}'.format(state.volume)) - @commands.command(pass_context=True) + @commands.command() @commands.guild_only() @utils.custom_perms(kick_members=True) async def pause(self, ctx): @@ -440,7 +483,7 @@ class Music: if state and state.voice and state.voice.is_connected(): state.voice.pause() - @commands.command(pass_context=True) + @commands.command() @commands.guild_only() @utils.custom_perms(kick_members=True) async def resume(self, ctx): @@ -449,7 +492,7 @@ class Music: if state and state.voice and state.voice.is_connected(): state.voice.resume() - @commands.command(pass_context=True) + @commands.command() @commands.guild_only() @utils.custom_perms(kick_members=True) async def stop(self, ctx): @@ -457,23 +500,22 @@ class Music: This also clears the queue. """ state = self.voice_states.get(ctx.message.guild.id) - voice = ctx.message.guild.voice_client - if voice: - voice.stop() - await voice.disconnect(force=True) - if state: + # Stop playing whatever song is playing. + if state and state.voice: + state.voice.stop() + state.songs.clear() # This will cancel the audio event we're using to loop through the queue # Then erase the voice_state entirely, and disconnect from the channel - state.audio_player.cancel() + await state.voice.disconnect() try: del self.voice_states[ctx.message.guild.id] except KeyError: pass - @commands.command(pass_context=True) + @commands.command() @commands.guild_only() @utils.custom_perms(send_messages=True) async def eta(self, ctx): @@ -485,7 +527,10 @@ class Music: await ctx.send('Not playing any music right now...') return - queue = state.songs.entries + if state.user_queue: + queue = [x.peek() for x in state.djs if x.peek()] + else: + queue = state.songs.entries if len(queue) == 0: await ctx.send("Nothing currently in the queue") return @@ -506,7 +551,7 @@ class Music: return await ctx.send("ETA till your next play is: {0[0]}m {0[1]}s".format(divmod(round(count, 0), 60))) - @commands.command(pass_context=True) + @commands.command() @commands.guild_only() @utils.custom_perms(send_messages=True) async def queue(self, ctx): @@ -515,25 +560,35 @@ class Music: if state is None: await ctx.send("Nothing currently in the queue") return - # Asyncio provides no non-private way to access the queue, so we have to use _queue - _queue = state.songs.entries + + if state.user_queue: + _queue = [x.peek() for x in state.djs if x.peek()] + else: + _queue = state.songs.entries if len(_queue) == 0: await ctx.send("Nothing currently in the queue") else: self.bot.loop.create_task(self.queue_embed_task(state, ctx.message.channel, ctx.message.author)) - @commands.command(pass_context=True) + @commands.command() @commands.guild_only() @utils.custom_perms(send_messages=True) async def queuelength(self, ctx): """Prints the length of the queue""" state = self.voice_states.get(ctx.message.guild.id) - if state: - await ctx.send("There are a total of {} songs in the queue".format(len(state.songs.entries))) - else: - await ctx.send("There are no songs in the queue") + if state is None: + await ctx.send("Nothing currently in the queue") + return - @commands.command(pass_context=True) + if state.user_queue: + _queue = [x.peek() for x in state.djs if x.peek()] + else: + _queue = state.songs.entries + if len(_queue) == 0: + await ctx.send("Nothing currently in the queue") + await ctx.send("There are a total of {} songs in the queue".format(len(_queue))) + + @commands.command() @commands.guild_only() @utils.custom_perms(send_messages=True) async def skip(self, ctx): @@ -566,7 +621,7 @@ class Music: else: await ctx.send('You have already voted to skip this song.') - @commands.command(pass_context=True) + @commands.command() @commands.guild_only() @utils.custom_perms(kick_members=True) async def modskip(self, ctx): @@ -579,7 +634,7 @@ class Music: state.skip() await ctx.send('Song has just been skipped.') - @commands.command(pass_context=True) + @commands.command() @commands.guild_only() @utils.custom_perms(send_messages=True) async def playing(self, ctx): @@ -608,6 +663,38 @@ class Music: # And send the embed await ctx.send(embed=embed) + @commands.command() + @commands.guild_only() + @utils.custom_perms(send_messages=True) + async def dj(self, ctx): + """Attempts to join the current DJ queue + + EXAMPLE: !dj + RESULT: You are 7th on the waitlist for the queue""" + if ctx.message.guild.id not in self.voice_states: + if not await ctx.invoke(self.join): + return + + state = self.voice_states.get(ctx.message.guild.id) + if not state.user_queue: + await ctx.send("The current queue type is the song queue. " + "Use the command {}play to add a song to the queue".format(ctx.prefix)) + return + + if state.get_dj(ctx.message.author): + await ctx.send("You are already in the DJ queue!") + else: + new_dj = self.bot.get_cog('DJEvents').djs[ctx.message.author.id] + state.djs.append(new_dj) + try: + await ctx.send("You have joined the DJ queue; there are currently {} people ahead of you".format( + state.djs.index(new_dj))) + except discord.Forbidden: + pass + + if not state.playing: + await state.play_next_song() + def setup(bot): bot.add_cog(Music(bot)) diff --git a/cogs/osu.py b/cogs/osu.py index 4d5962c..232c256 100644 --- a/cogs/osu.py +++ b/cogs/osu.py @@ -44,7 +44,7 @@ class Osu: async def get_users(self): """A task used to 'cache' all member's and their Osu profile's""" - data = await utils.get_content('osu') + data = await self.bot.db.actual_load('osu') if data is None: return @@ -56,11 +56,12 @@ class Osu: @commands.group(invoke_without_command=True) @utils.custom_perms(send_messages=True) - async def osu(self, ctx, member: discord.Member=None): + async def osu(self, ctx, member: discord.Member = None): """Provides basic information about a specific user EXAMPLE: !osu @Person RESULT: Informationa bout that person's osu account""" + await ctx.message.channel.trigger_typing() if member is None: member = ctx.message.author @@ -92,6 +93,7 @@ class Osu: EXAMPLE: !osu add username RESULT: Links your username to your account, and allows stats to be pulled from it""" + await ctx.message.channel.trigger_typing() author = ctx.message.author user = await self.get_user(author, username) if user is None: @@ -103,8 +105,7 @@ class Osu: 'osu_username': user.username } - if not await utils.add_content('osu', entry): - await utils.update_content('osu', entry, str(author.id)) + self.bot.db.save('osu', entry) await ctx.send("I have just saved your Osu user {}".format(author.display_name)) @@ -116,7 +117,7 @@ class Osu: EXAMPLE: !osu scores @Person 5 RESULT: The top 5 maps for the user @Person""" - + await ctx.message.channel.trigger_typing() # Set the defaults before we go through our passed data to figure out what we want limit = 5 member = ctx.message.author @@ -166,7 +167,9 @@ class Osu: {'name': 'Length', 'value': m.total_length}, {'name': 'Score', 'value': i.score}, {'name': 'Max Combo', 'value': i.maxcombo}, - {'name': 'Hits', 'value': "{}/{}/{}/{} (300/100/50/misses)".format(i.count300, i.count100, i.count50, i.countmiss), "inline": False}, + {'name': 'Hits', + 'value': "{}/{}/{}/{} (300/100/50/misses)".format(i.count300, i.count100, i.count50, i.countmiss), + "inline": False}, {'name': 'Perfect', 'value': "Yes" if i.perfect else "No"}, {'name': 'Rank', 'value': i.rank}, {'name': 'PP', 'value': i.pp}, @@ -182,5 +185,6 @@ class Osu: except utils.CannotPaginate as e: await ctx.send(str(e)) + def setup(bot): bot.add_cog(Osu(bot)) diff --git a/cogs/overwatch.py b/cogs/overwatch.py index 92ca5e3..9b08135 100644 --- a/cogs/overwatch.py +++ b/cogs/overwatch.py @@ -38,14 +38,11 @@ class Overwatch: await ctx.message.channel.trigger_typing() user = user or ctx.message.author - ow_stats = await utils.get_content('overwatch', str(user.id)) + bt = self.bot.db.load('overwatch', key=str(user.id), pluck='battletag') - if ow_stats is None: + if bt is None: await ctx.send("I do not have this user's battletag saved!") return - # This API sometimes takes a while to look up information, so send a message saying we're processing - - bt = ow_stats['battletag'] if hero == "": # If no hero was provided, we just want the base stats for a player @@ -114,11 +111,12 @@ class Overwatch: return # Now just save the battletag - entry = {'member_id': key, 'battletag': bt} - update = {'battletag': bt} - # Try adding this first, if that fails, update the saved entry - if not await utils.add_content('overwatch', entry): - await utils.update_content('overwatch', update, key) + entry = { + 'member_id': key, + 'battletag': bt + } + + self.bot.db.save('overwatch', entry) await ctx.send("I have just saved your battletag {}".format(ctx.message.author.mention)) @ow.command(pass_context=True, name="delete", aliases=['remove']) @@ -128,10 +126,12 @@ class Overwatch: EXAMPLE: !ow delete RESULT: Your battletag is no longer saved""" - if await utils.remove_content('overwatch', str(ctx.message.author.id)): - await ctx.send("I no longer have your battletag saved {}".format(ctx.message.author.mention)) - else: - await ctx.send("I don't even have your battletag saved {}".format(ctx.message.author.mention)) + entry = { + 'member_id': str(ctx.message.author.id), + 'battletag': None + } + self.bot.db.save('overwatch', entry) + await ctx.send("I no longer have your battletag saved {}".format(ctx.message.author.mention)) def setup(bot): diff --git a/cogs/owner.py b/cogs/owner.py index 04fa4ba..27548b2 100644 --- a/cogs/owner.py +++ b/cogs/owner.py @@ -180,19 +180,6 @@ class Owner: except discord.HTTPException as e: await ctx.send('Unexpected error: `{}`'.format(e)) - @commands.command() - @commands.check(utils.is_owner) - async def motd_push(self, ctx, *, message): - """Used to push a new message to the message of the day""" - date = pendulum.utcnow().to_date_string() - key = date - entry = {'motd': message, 'date': date} - # Try to add this, if there's an entry for that date, lets update it to make sure only one motd is sent a day - # I should be managing this myself, more than one should not be sent in a day - if await utils.add_content('motd', entry): - await utils.update_content('motd', entry, key) - await ctx.send("New motd update for {}!".format(date)) - @commands.command() @commands.check(utils.is_owner) async def sendtochannel(self, ctx, cid: int, *, message): diff --git a/cogs/picarto.py b/cogs/picarto.py index d8b2556..ab3aa86 100644 --- a/cogs/picarto.py +++ b/cogs/picarto.py @@ -17,10 +17,6 @@ BASE_URL = 'https://ptvappapi.picarto.tv' api_key = '03e26294-b793-11e5-9a41-005056984bd4' - - - - class Picarto: def __init__(self, bot): self.bot = bot @@ -47,7 +43,7 @@ class Picarto: try: while not self.bot.is_closed(): await self.get_online_users() - picarto = await utils.filter_content('picarto', {'notifications_on': 1}) + picarto = self.bot.db.load('picarto', table_filter={'notifications_on': 1}) for data in picarto: m_id = int(data['member_id']) url = data['picarto_url'] @@ -62,17 +58,16 @@ class Picarto: member = server.get_member(m_id) if member is None: continue - server_settings = await utils.get_content('server_settings', s_id) - if server_settings is not None: - channel_id = int(server_settings.get('notification_channel', s_id)) - else: - channel_id = int(s_id) + channel_id = self.bot.db.load('server_settings', key=s_id, + pluck='notifications_channel') or int(s_id) channel = server.get_channel(channel_id) try: - await channel.send("{} has just gone live! View their stream at <{}>".format(member.display_name, data['picarto_url'])) + await channel.send( + "{} has just gone live! View their stream at <{}>".format(member.display_name, + data['picarto_url'])) except discord.Forbidden: pass - self.bot.loop.create_task(utils.update_content('picarto', {'live': 1}, str(m_id))) + self.bot.db.save('picarto', {'live': 1, 'member_id': str(m_id)}) elif not online and data['live'] == 1: for s_id in data['servers']: server = self.bot.get_guild(int(s_id)) @@ -81,17 +76,16 @@ class Picarto: member = server.get_member(m_id) if member is None: continue - server_settings = await utils.get_content('server_settings', s_id) - if server_settings is not None: - channel_id = int(server_settings.get('notification_channel', s_id)) - else: - channel_id = int(s_id) + channel_id = self.bot.db.load('server_settings', key=s_id, + pluck='notifications_channel') or int(s_id) channel = server.get_channel(channel_id) try: - await channel.send("{} has just gone offline! View their stream next time at <{}>".format(member.display_name, data['picarto_url'])) + await channel.send( + "{} has just gone offline! View their stream next time at <{}>".format( + member.display_name, data['picarto_url'])) except discord.Forbidden: pass - self.bot.loop.create_task(utils.update_content('picarto', {'live': 0}, str(m_id))) + self.bot.db.save('picarto', {'live': 0, 'member_id': str(m_id)}) await asyncio.sleep(30) except Exception as e: tb = traceback.format_exc() @@ -109,13 +103,11 @@ class Picarto: # If member is not given, base information on the author member = member or ctx.message.author - picarto_entry = await utils.get_content('picarto', str(member.id)) - if picarto_entry is None: + member_url = self.bot.db.load('picarto', key=member.id, pluck='picarto_url') + if member_url is None: await ctx.send("That user does not have a picarto url setup!") return - member_url = picarto_entry['picarto_url'] - # Use regex to get the actual username so that we can make a request to the API stream = re.search("(?<=picarto.tv/)(.*)", member_url).group(1) url = BASE_URL + '/channel/{}'.format(stream) @@ -141,9 +133,9 @@ class Picarto: # Social URL's can be given if a user wants them to show # Print them if they exist, otherwise don't try to include them - social_links = data.get('social_urls') + social_links = data.get('social_urls', {}) - for i, result in data['social_urls'].items(): + for i, result in social_links.items(): embed.add_field(name=i.title(), value=result) await ctx.send(embed=embed) @@ -183,29 +175,38 @@ class Picarto: return key = str(ctx.message.author.id) - entry = {'picarto_url': url, - 'servers': [str(ctx.message.guild.id)], - 'notifications_on': 1, - 'live': 0, - 'member_id': key} - if await utils.add_content('picarto', entry): - await ctx.send( + + # Check if it exists first, if it does we don't want to override some of the settings + result = self.bot.db.load('picarto', key=key) + if result: + entry = { + 'picarto_url': url, + 'member_id': key + } + else: + entry = { + 'picarto_url': url, + 'servers': [str(ctx.message.guild.id)], + 'notifications_on': 1, + 'live': 0, + 'member_id': key + } + self.bot.db.save('picarto', entry) + await ctx.send( "I have just saved your Picarto URL {}, this guild will now be notified when you go live".format( ctx.message.author.mention)) - else: - await utils.update_content('picarto', {'picarto_url': url}, key) - await ctx.send("I have just updated your Picarto URL") @picarto.command(name='remove', aliases=['delete']) @utils.custom_perms(send_messages=True) async def remove_picarto_url(self, ctx): """Removes your picarto URL""" - if await utils.remove_content('picarto', str(ctx.message.author.id)): - await ctx.send("I am no longer saving your picarto URL {}".format(ctx.message.author.mention)) - else: - await ctx.send( - "I do not have your picarto URL added {}. You can save your picarto url with {}picarto add".format( - ctx.message.author.mention, ctx.prefix)) + entry = { + 'picarto_url': None, + 'member_id': str(ctx.message.author.id) + } + + self.bot.db.save('picarto', entry) + await ctx.send("I am no longer saving your picarto URL {}".format(ctx.message.author.mention)) @picarto.group(invoke_without_command=True) @commands.guild_only() @@ -217,39 +218,61 @@ class Picarto: EXAMPLE: !picarto notify RESULT: This guild will now be notified of you going live""" key = str(ctx.message.author.id) - result = await utils.get_content('picarto', key) + servers = self.bot.db.load('picarto', key=key, pluck='servers') # Check if this user is saved at all - if result is None: + if servers is None: await ctx.send( "I do not have your Picarto URL added {}. You can save your Picarto url with !picarto add".format( ctx.message.author.mention)) # Then check if this guild is already added as one to notify in - elif ctx.message.guild.id in result['servers']: + elif str(ctx.message.guild.id) in servers: await ctx.send("I am already set to notify in this guild...") else: - await utils.update_content('picarto', {'servers': r.row['servers'].append(str(ctx.message.guild.id))}, key) + servers.append(str(ctx.message.guild.id)) + entry = { + 'member_id': key, + 'servers': servers + } + self.bot.db.save('picarto', entry) + await ctx.send("This server will now be notified if you go live") @notify.command(name='on', aliases=['start,yes']) + @commands.guild_only() @utils.custom_perms(send_messages=True) async def notify_on(self, ctx): """Turns picarto notifications on EXAMPLE: !picarto notify on RESULT: Notifications are sent when you go live""" - if await utils.update_content('picarto', {'notifications_on': 1}, str(ctx.message.author.id)): + key = str(ctx.message.author.id) + result = self.bot.db.load('picarto', key=key) + if result: + entry = { + 'member_id': key, + 'notifications_on': 1 + } + self.bot.db.save('picarto', entry) await ctx.send("I will notify if you go live {}, you'll get a bajillion followers I promise c:".format( ctx.message.author.mention)) else: await ctx.send("I can't notify if you go live if I don't know your picarto URL yet!") - @notify.command(name='off', aliases=['stop,no'], pass_context=True) + @notify.command(name='off', aliases=['stop,no']) + @commands.guild_only() @utils.custom_perms(send_messages=True) async def notify_off(self, ctx): """Turns picarto notifications off EXAMPLE: !picarto notify off RESULT: No more notifications sent when you go live""" - if await utils.update_content('picarto', {'notifications_on': 0}, str(ctx.message.author.id)): + key = str(ctx.message.author.id) + result = self.bot.db.load('picarto', key=key) + if result: + entry = { + 'member_id': key, + 'notifications_on': 0 + } + self.bot.db.save('picarto', entry) await ctx.send( "I will not notify if you go live anymore {}, " "are you going to stream some lewd stuff you don't want people to see?~".format( diff --git a/cogs/playlist.py b/cogs/playlist.py new file mode 100644 index 0000000..a5360d0 --- /dev/null +++ b/cogs/playlist.py @@ -0,0 +1,337 @@ +import discord +import asyncio +from discord.ext import commands + +from . import utils + + +class Playlist: + """Used to manage user playlists""" + + def __init__(self, bot): + self.bot = bot + + async def get_response(self, ctx, question): + # Save our simple variables + channel = ctx.message.channel + author = ctx.message.author + + # Create our check function used to ensure the author and channel are the only possible message we get + check = lambda m: m.author == author and m.channel == channel + + try: + # Ask our question, wait 60 seconds for a response + my_msg = await ctx.send(question) + response = await self.bot.wait_for('message', check=check, timeout=60) + except asyncio.TimeoutError: + # If we timeout, let them know and return None + await ctx.send("You took too long. I'm impatient, don't make me wait") + return None + else: + # If succesful try to delete the message we sent, and the response + try: + await my_msg.delete() + await response.delete() + except discord.Forbidden: + pass + + # For our case here, everything needs to be lowered and stripped, so just do this now + return response.content.lower().strip() + + async def get_info(self, song_url): + try: + # Just download the information + info = await self.bot.downloader.extract_info(self.bot.loop, song_url, download=False) + except Exception as e: + # If we fail, it's possibly due to an incorrect detection as a URL instead of a search + if "gaierror" in str(e) or "unknown url type" in str(e): + # So just force a search + song_url = "ytsearch:" + song_url + info = await self.bot.downloader.extract_info(self.bot.loop, song_url, download=False) + else: + # Otherwise if we fail, we just want to return None + return None + + # If we detected a search, get the first entry in the results + if info.get('_type', None) == 'playlist': + if info.get('extractor') == 'youtube:search': + if len(info['entries']) == 0: + return None + else: + info = info['entries'][0] + song_url = info['webpage_url'] + + # If we are successful, create the entry we'll need to add to the playlist database, and return it + if info: + return { + 'title': info.get('title', 'Untitled'), + 'url': song_url + } + else: + return None + + async def add_to_playlist(self, author, playlist, url): + # Simply get the database entry for this user's playlist + key = str(author.id) + playlist = playlist.lower().strip() + playlists = self.bot.db.load('user_playlists', key=key, pluck='playlists') or [] + + entry = await self.get_info(url) + + # Search through, find the name that matches the playlist + if entry: + for pl in playlists: + if pl['name'] == playlist: + # If we find it, add the song entry to the songs + pl['songs'].append(entry) + # Create the json needed to save to the database, and save + update = { + 'member_id': key, + 'playlists': playlists + } + self.bot.db.save('user_playlists', update) + return True + + async def rename_playlist(self, author, old_name, new_name): + # Simply get the database entry for this user's playlist + key = str(author.id) + old_name = old_name.lower().strip() + new_name = new_name.lower().strip() + playlists = self.bot.db.load('user_playlists', key=key, pluck='playlists') or [] + + # Find the playlist that matches the old name + for pl in playlists: + if pl['name'] == old_name: + # Once found, change the name, update the json, save + pl['name'] = new_name + update = { + 'member_id': key, + 'playlists': playlists + } + self.bot.db.save('user_playlists', update) + return True + + async def remove_from_playlist(self, author, playlist, index): + # Simply get the database entry for this user's playlist + key = str(author.id) + playlist = playlist.lower().strip() + playlists = self.bot.db.load('user_playlists', key=key, pluck='playlists') or [] + + # Loop through till we find the playlist that matches + for pl in playlists: + if pl['name'] == playlist: + song = pl['songs'][index] + # Once found, remove the matching song, update json, save + pl['songs'].remove(song) + update = { + 'member_id': key, + 'playlists': playlists + } + self.bot.db.save('user_playlists', update) + return song + + async def update_dj_for_member(self, member): + music = self.bot.get_cog('Music') + if music: + for state in music.voice_states.values(): + dj = state.get_dj(member) + if dj: + # We want to add a slight delay to this, because our database method launches a task to update + # Before we update what is live, we need the information saved in (at least the cache) the database + await asyncio.sleep(2) + self.bot.loop.create_task(dj.resolve_playlist()) + + @commands.command() + @utils.custom_perms(send_messages=True) + async def playlists(self, ctx): + """Displays the playlists you have + + EXAMPLE: !playlists + RESULT: All your playlists""" + # Get all the author's playlists + playlists = self.bot.db.load('user_playlists', key=ctx.message.author.id, pluck='playlists') + if playlists: + # Create the entries for our paginator detailing the name of the playlist, and the number of songs in it + entries = [ + "{} ({} songs)".format(x['name'], len(x['songs'])) if not x.get('active') + else "{} ({} songs) - Active playlist".format(x['name'], len(x['songs'])) + for x in playlists + ] + + try: + # And paginate + pages = utils.Pages(self.bot, message=ctx.message, entries=entries) + await pages.paginate() + except utils.CannotPaginate as e: + await ctx.send(str(e)) + else: + await ctx.send("You do not have any playlists") + + @commands.group(invoke_without_command=True) + @utils.custom_perms(send_messages=True) + async def playlist(self, ctx, *, playlist_name): + """Used to view your playlists + + EXAMPLE: !playlist Playlist 2 + RESULT: Displays the songs in your playlist called "Playlist 2" """ + playlist_name = playlist_name.lower().strip() + + playlists = self.bot.db.load('user_playlists', key=ctx.message.author.id, pluck='playlists') + try: + # Get the playlist if the name matches + playlist = [x for x in playlists if playlist_name == x['name']][0] + # Create the entries for our paginator just based on the title of the songs in the playlist + entries = ["{}".format(x['title']) for x in playlist['songs']] + # Paginate + pages = utils.Pages(self.bot, message=ctx.message, entries=entries) + await pages.paginate() + except (IndexError, TypeError, KeyError): + await ctx.send("You do not have a playlist named {}!".format(playlist_name)) + except utils.CannotPaginate as e: + await ctx.send(str(e)) + + @playlist.command(name='create') + @utils.custom_perms(send_messages=True) + async def _pl_create(self, ctx, *, name): + """Used to create a new playlist + + EXAMPLE: !playlist create Playlist + RESULT: A new playlist called Playlist""" + key = str(ctx.message.author.id) + playlists = self.bot.db.load('user_playlists', key=key, pluck='playlists') or [] + + # Create the new playlist entry + entry = { + 'name': name.lower().strip(), + 'songs': [] + } + + # Check to make sure that there isn't a playlist with the same name + names = [x['name'] for x in playlists] + if name in names: + await ctx.send('You already have a playlist called {}'.format(name)) + # Otherwise add this new playlist, and save + else: + # This is here to set the first playlist we create as the active one. + # If someone has a playlist already, we don't want to change which is the active one + # If they don't have any, then we want to set our first one as the active one + entry['active'] = len(playlists) == 0 + playlists.append(entry) + update = { + 'member_id': key, + 'playlists': playlists + } + self.bot.db.save('user_playlists', update) + await ctx.send("You have just created a new playlist called {}".format(name)) + + @playlist.command(name='edit') + @utils.custom_perms(send_messages=True) + async def _pl_edit(self, ctx): + """A command used to edit a current playlist + The available ways to edit a playlist are to rename, add a song, remove a song, or delete the playlist + + EXAMPLE: !playlist edit + RESULT: A followalong asking for what you need""" + # Load the playlists for the author + key = str(ctx.message.author.id) + playlists = self.bot.db.load('user_playlists', key=key, pluck='playlists') or [] + # Also create a list of the names for easy comparision + names = [x['name'] for x in playlists] + + if not playlists: + await ctx.send("You have no playlists to edit!") + return + + # Show the playlists we have, and ask which to choose from + await ctx.invoke(self.playlists) + question = "Please provide what playlist you would like to edit, the playlists you have available are above." + playlist = await self.get_response(ctx, question) + if not playlist: + return + if playlist not in names: + await ctx.send("You do not have a playlist named {}!".format(playlist)) + return + + q1 = "How would you like to edit {}? Choices are `add`, `remove`, `rename`, `delete`, or `activate`.\n" \ + "**add** - Adds a song to this playlist\n" \ + "**remove** - Removes a song from this playlist\n" \ + "**rename** - Changes the name of this playlist\n" \ + "**delete** - Deletes this playlist\n" \ + "**activate** - Sets this as the active playlist\n\n" \ + "Type **quit** to stop editing this playlist".format(playlist) + + # Lets create a list of the messages we'll delete after + delete_msgs = [] + + # We want to loop this in order to continue editing, till the user is done + while True: + response = await self.get_response(ctx, q1) + + if 'add' in response: + # Ask the user what song to add, get the response, add it + question = "What is the song you would like to add to {}?".format(playlist) + response = await self.get_response(ctx, question) + # If we didn't get a response, just continue with the loop, we have no need to say anything + # The "error" message is sent with our `get_response` helper method + if response: + await ctx.message.channel.trigger_typing() + await self.add_to_playlist(ctx.message.author, playlist, response) + delete_msgs.append(await ctx.send("Successfully added song {} to playlist {}".format(response, + playlist))) + elif 'remove' in response: + await ctx.invoke(self.playlist, playlist_name=playlist) + question = "Please provide just the number of the song you want to delete" + try: + response = await self.get_response(ctx, question) + if response: + num = int(response) - 1 + song = await self.remove_from_playlist(ctx.author, playlist, num) + await ctx.send("Successfully removed {} from {}".format(song['title'], playlist)) + except (ValueError, IndexError): + delete_msgs.append(await ctx.send("Please provide just the number of the song you want to delete " + "next time!")) + elif 'delete' in response: + playlists = [x for x in playlists if x['name'] != playlist] + entry = { + 'member_id': str(key), + 'playlists': playlists + } + self.bot.db.save('user_playlists', entry) + delete_msgs.append(await ctx.send("Successfully deleted playlist {}".format(playlist))) + await ctx.send("Finished editing {}".format(playlist)) + break + elif 'rename' in response: + question = "What would you like to rename the playlist {} to?".format(playlist) + new_name = await self.get_response(ctx, question) + if new_name: + await self.rename_playlist(ctx.message.author, playlist, new_name) + playlist = new_name + delete_msgs.append(await ctx.send("Successfully renamed {} to {}!".format(playlist, new_name))) + elif 'activate' in response: + for x in playlists: + x['active'] = x['name'] == playlist + + entry = { + 'member_id': str(key), + 'playlists': playlists + } + self.bot.db.save('user_playlists', entry) + # Now we have edited the user's actual playlist...but we need to + delete_msgs.append(await ctx.send("{} is now your active playlist".format(playlist))) + elif 'quit' in response: + await ctx.send("Finished editing {}".format(playlist)) + break + else: + delete_msgs.append(await ctx.send("That is not a valid option!")) + + # After whatever has been edited, we need to update the live DJ's + await self.update_dj_for_member(ctx.message.author) + + if len(delete_msgs) == 1: + await delete_msgs[0].delete() + elif len(delete_msgs) > 1: + await ctx.message.channel.delete_messages(delete_msgs) + + +def setup(bot): + bot.add_cog(Playlist(bot)) diff --git a/cogs/polls.py b/cogs/polls.py index efebe03..55d16d2 100644 --- a/cogs/polls.py +++ b/cogs/polls.py @@ -1,9 +1,11 @@ from discord.ext import commands from . import utils + def to_keycap(c): return '\N{KEYCAP TEN}' if c == 10 else str(c) + '\u20e3' + class Poll: def __init__(self, message): self.message = message @@ -68,7 +70,6 @@ class Polls: if poll: await poll.remove_other_reaction(reaction, user) - @commands.command(pass_context=True) @commands.guild_only() @utils.custom_perms(send_messages=True) diff --git a/cogs/raffle.py b/cogs/raffle.py index c2674bb..10c0d55 100644 --- a/cogs/raffle.py +++ b/cogs/raffle.py @@ -8,6 +8,7 @@ import pendulum import re import asyncio import traceback +import rethinkdb as r class Raffle: @@ -28,16 +29,20 @@ class Raffle: async def check_raffles(self): # This is used to periodically check the current raffles, and see if they have ended yet # If the raffle has ended, we'll pick a winner from the entrants - raffles = await utils.get_content('raffles') + raffles = self.bot.db.load('raffles') if raffles is None: return for raffle in raffles: server = self.bot.get_guild(int(raffle['server_id'])) + title = raffle['title'] + entrants = raffle['entrants'] + raffle_id = raffle['id'] # Check to see if this cog can find the server in question if server is None: + await self.bot.db.query(r.table('raffles').get(raffle_id).delete()) continue now = pendulum.utcnow() @@ -47,10 +52,6 @@ class Raffle: if expires > now: continue - title = raffle['title'] - entrants = raffle['entrants'] - raffle_id = raffle['id'] - # Make sure there are actually entrants if len(entrants) == 0: fmt = 'Sorry, but there were no entrants for the raffle `{}`!'.format(title) @@ -72,21 +73,17 @@ class Raffle: else: fmt = 'The raffle `{}` has just ended! The winner is {}!'.format(title, winner.display_name) - # No matter which one of these matches were met, the raffle has ended and we want to remove it - # We don't have to wait for it however, so create a task for it - self.bot.loop.create_task(utils.remove_content('raffles', raffle_id)) - - server_settings = await utils.get_content('server_settings', str(server.id)) - if server_settings is None: - channel = self.bot.get_channel(server.id) - else: - channel_id = server_settings.get('notification_channel', server.id) - channel = self.bot.get_channel(channel_id) + channel_id = self.bot.db.load('server_settings', key=server.id, + pluck='notifications_channel') or server.id + channel = self.bot.get_channel(channel_id) try: await channel.send(fmt) except (discord.Forbidden, AttributeError): pass + # No matter which one of these matches were met, the raffle has ended and we want to remove it + await self.bot.db.query(r.table('raffles').get(raffle_id).delete()) + @commands.command() @commands.guild_only() @utils.custom_perms(send_messages=True) @@ -96,11 +93,16 @@ class Raffle: EXAMPLE: !raffles RESULT: A list of the raffles setup on this server""" r_filter = {'server_id': str(ctx.message.guild.id)} - raffles = await utils.filter_content('raffles', r_filter) + raffles = self.bot.db.load('raffles', table_filter=r_filter) if raffles is None: await ctx.send("There are currently no raffles setup on this server!") return + # For EVERY OTHER COG, when we get one result, it is nice to have it return that exact object + # This is the only cog where that is different, so just to make this easier lets throw it + # back in a one-indexed list, for easier parsing + if isinstance(raffles, dict): + raffles = [raffles] fmt = "\n\n".join("**Raffle:** {}\n**Title:** {}\n**Total Entrants:** {}\n**Ends:** {} UTC".format( num + 1, raffle['title'], @@ -122,12 +124,16 @@ class Raffle: r_filter = {'server_id': str(ctx.message.guild.id)} author = ctx.message.author - raffles = await utils.filter_content('raffles', r_filter) + raffles = self.bot.db.load('raffles', table_filter=r_filter) if raffles is None: await ctx.send("There are currently no raffles setup on this server!") return - raffle_count = len(raffles) + if isinstance(raffles, list): + raffle_count = len(raffles) + elif isinstance(raffles, dict): + raffles = [raffles] + raffle_count = 1 # There is only one raffle, so use the first's info if raffle_count == 1: @@ -138,8 +144,11 @@ class Raffle: return entrants.append(str(author.id)) - update = {'entrants': entrants} - await utils.update_content('raffles', update, raffles[0]['id']) + update = { + 'entrants': entrants, + 'id': raffles[0]['id'] + } + self.bot.db.save('raffles', update) await ctx.send("{} you have just entered the raffle!".format(author.mention)) # Otherwise, make sure the author gave a valid raffle_id elif raffle_id in range(raffle_count - 1): @@ -152,13 +161,17 @@ class Raffle: entrants.append(str(author.id)) # Since we have no good thing to filter things off of, lets use the internal rethinkdb id - update = {'entrants': entrants} - await utils.update_content('raffles', update, raffles[raffle_id]['id']) + + update = { + 'entrants': entrants, + 'id': raffles[0]['id'] + } + self.bot.db.save('raffles', update) await ctx.send("{} you have just entered the raffle!".format(author.mention)) else: fmt = "Please provide a valid raffle ID, as there are more than one setup on the server! " \ "There are currently `{}` raffles running, use {}raffles to view the current running raffles".format( - raffle_count, ctx.prefix) + raffle_count, ctx.prefix) await ctx.send(fmt) @raffle.command(pass_context=True, name='create', aliases=['start', 'begin', 'add']) @@ -198,6 +211,7 @@ class Raffle: return re.search("\d+ (minutes?|hours?|days?|weeks?|months?)", m.content.lower()) is not None else: return False + try: msg = await self.bot.wait_for('message', timeout=120, check=check) except asyncio.TimeoutError: @@ -238,14 +252,16 @@ class Raffle: expires = now.add(**payload) # Now we're ready to add this as a new raffle - entry = {'title': title, - 'expires': expires.to_datetime_string(), - 'entrants': [], - 'author': str(author.id), - 'server_id': str(server.id)} + entry = { + 'title': title, + 'expires': expires.to_datetime_string(), + 'entrants': [], + 'author': str(author.id), + 'server_id': str(server.id) + } # We don't want to pass a filter to this, because we can have multiple raffles per server - await utils.add_content('raffles', entry) + self.bot.db.save('raffles', entry) await ctx.send("I have just saved your new raffle!") diff --git a/cogs/roles.py b/cogs/roles.py index 02ac908..038b44a 100644 --- a/cogs/roles.py +++ b/cogs/roles.py @@ -324,12 +324,8 @@ class Roles: author = ctx.message.author key = str(ctx.message.guild.id) - server_settings = await utils.get_content('server_settings', key) + self_assignable_roles = self.bot.db.load('server_settings', key=key, pluck='self_assignable_roles') or [] - if server_settings is None: - await ctx.send("There are no self-assignable roles on this server") - return - self_assignable_roles = server_settings.get('self_assignable_roles', []) if len(self_assignable_roles) == 0: await ctx.send("There are no self-assignable roles on this server") return @@ -361,12 +357,8 @@ class Roles: author = ctx.message.author key = str(ctx.message.guild.id) - server_settings = await utils.get_content('server_settings', key) + self_assignable_roles = self.bot.db.load('server_settings', key=key, pluck='self_assignable_roles') or [] - if server_settings is None: - await ctx.send("There are no self-assignable roles on this server") - return - self_assignable_roles = server_settings.get('self_assignable_roles', []) if len(self_assignable_roles) == 0: await ctx.send("There are no self-assignable roles on this server") return @@ -394,17 +386,16 @@ class Roles: RESULT: Allows users to self-assign the roles Member, and NSFW""" roles = [str(r.id) for r in role] key = str(ctx.message.guild.id) - server_settings = await utils.get_content('server_settings', key) - if server_settings is None: - entry = {'server_id': key, 'self_assignable_roles': roles} - await utils.add_content('server_settings', entry) - else: - self_assignable_roles = server_settings.get('self_assignable_roles', []) - self_assignable_roles.extend(roles) - self_assignable_roles = list(set(self_assignable_roles)) - update = {'self_assignable_roles': self_assignable_roles} - await utils.update_content('server_settings', update, key) + self_assignable_roles = self.bot.db.load('server_settings', key=key, pluck='self_assignable_roles') or [] + self_assignable_roles.extend(roles) + self_assignable_roles = list(set(self_assignable_roles)) + entry = { + 'server_id': key, + 'self_assignable_roles': self_assignable_roles + } + + self.bot.db.save('server_settings', entry) if len(roles) == 1: fmt = "Successfully added {} as a self-assignable role".format(role[0].name) @@ -423,12 +414,7 @@ class Roles: EXAMPLE: !assigns list RESUL: A list of all the self-assignable roles""" key = str(ctx.message.guild.id) - server_settings = await utils.get_content('server_settings', key) - - if server_settings is None: - await ctx.send("There are no self-assignable roles on this server") - return - self_assignable_roles = server_settings.get('self_assignable_roles', []) + self_assignable_roles = self.bot.db.load('server_settings', key=key, pluck='self_assignable_roles') or [] if len(self_assignable_roles) == 0: await ctx.send("There are no self-assignable roles on this server") return @@ -458,12 +444,7 @@ class Roles: EXAMPLE: !assigns remove Member NSFW RESULT: Removes the ability for users to self-assign the roles Member, and NSFW""" key = str(ctx.message.guild.id) - server_settings = await utils.get_content('server_settings', key) - - if server_settings is None: - await ctx.send("There are no self-assignable roles on this server") - return - self_assignable_roles = server_settings.get('self_assignable_roles', []) + self_assignable_roles = self.bot.db.load('server_settings', key=key, pluck='self_assignable_roles') or [] if len(self_assignable_roles) == 0: await ctx.send("There are no self-assignable roles on this server") return @@ -478,8 +459,11 @@ class Roles: else: fmt += "\n{} is no longer a self-assignable role".format(r.name) - update = {'self_assignable_roles': self_assignable_roles} - await utils.update_content('server_settings', update, key) + update = { + 'self_assignable_roles': self_assignable_roles, + 'server_id': key + } + self.bot.db.save('server_settings', update) await ctx.send(fmt) diff --git a/cogs/stats.py b/cogs/stats.py index aab8d69..58d8a41 100644 --- a/cogs/stats.py +++ b/cogs/stats.py @@ -65,7 +65,7 @@ class Stats: await ctx.send("`{}` is not a valid command".format(command)) return - command_stats = await utils.get_content('command_usage', cmd.qualified_name) + command_stats = self.bot.db.load('command_usage', key=cmd.qualified_name) if command_stats is None: await ctx.send("That command has never been used! You know I worked hard on that! :c") return @@ -103,7 +103,7 @@ class Stats: if re.search('(author|me)', option): author = ctx.message.author # First lets get all the command usage - command_stats = await utils.get_content('command_usage') + command_stats = self.bot.db.load('command_usage') # Now use a dictionary comprehension to get just the command name, and usage # Based on the author's usage of the command stats = {data['command']: data['member_usage'].get(str(author.id)) for data in command_stats @@ -125,7 +125,7 @@ class Stats: elif re.search('server', option): # This is exactly the same as above, except server usage instead of member usage server = ctx.message.guild - command_stats = await utils.get_content('command_usage') + command_stats = self.bot.db.load('command_usage') stats = {data['command']: data['server_usage'].get(str(server.id)) for data in command_stats if data['server_usage'].get(str(server.id), 0) > 0} sorted_stats = sorted(stats.items(), key=lambda x: x[1], reverse=True) @@ -148,7 +148,7 @@ class Stats: EXAMPLE: !mostboops RESULT: You've booped @OtherPerson 351253897120935712093572193057310298 times!""" - boops = await utils.get_content('boops', str(ctx.message.author.id)) + boops = self.bot.db.load('boops', key=ctx.message.author.id) if boops is None: await ctx.send("You have not booped anyone {} Why the heck not...?".format(ctx.message.author.mention)) return @@ -182,7 +182,7 @@ class Stats: RESULT: The list of your booped members!""" await ctx.message.channel.trigger_typing() - boops = await utils.get_content('boops', str(ctx.message.author.id)) + boops = self.bot.db.load('boops', key=ctx.message.author.id) if boops is None: await ctx.send("You have not booped anyone {} Why the heck not...?".format(ctx.message.author.mention)) return @@ -220,7 +220,7 @@ class Stats: # Create a list of the ID's of all members in this server, for comparison to the records saved server_member_ids = [member.id for member in ctx.message.guild.members] - battles = await utils.get_content('battle_records') + battles = self.bot.db.load('battle_records') if battles is None or len(battles) == 0: await ctx.send("No one has battled on this server!") @@ -256,9 +256,10 @@ class Stats: # For this one, we don't want to pass a filter, as we do need all battle records # We need this because we want to make a comparison for overall rank - all_members = await utils.get_content('battle_records') + all_members = self.bot.db.load('battle_records') if all_members is None or len(all_members) == 0: - await ctx.send("You have not battled anyone!") + await ctx.send("That user has not battled yet!") + return # Make a list comprehension to just check if the user has battled if len([entry for entry in all_members if entry['member_id'] == str(member.id)]) == 0: diff --git a/cogs/tags.py b/cogs/tags.py index 96e3c84..896e88f 100644 --- a/cogs/tags.py +++ b/cogs/tags.py @@ -4,7 +4,7 @@ import discord from . import utils import asyncio -import rethinkdb as r + class Tags: """This class contains all the commands for custom tags""" @@ -20,9 +20,9 @@ class Tags: EXAMPLE: !tags RESULT: All tags setup on this server""" - tags = await utils.get_content('tags', str(ctx.message.guild.id)) - if tags and len(tags['tags']) > 0: - entries = [t['trigger'] for t in tags['tags']] + tags = self.bot.db.load('tags', key=ctx.message.guild.id, pluck='tags') + if tags: + entries = [t['trigger'] for t in tags] pages = utils.Pages(self.bot, message=ctx.message, entries=entries) await pages.paginate() else: @@ -36,16 +36,16 @@ class Tags: EXAMPLE: !mytags RESULT: All your tags setup on this server""" - tags = await utils.get_content('tags', str(ctx.message.guild.id)) - if not tags: - await ctx.send("There are no tags setup on this server!") - else: - entries = [t['trigger'] for t in tags['tags'] if t['author'] == str(ctx.message.author.id)] + tags = self.bot.db.load('tags', key=ctx.message.guild.id, pluck='tags') + if tags: + entries = [t['trigger'] for t in tags if t['author'] == str(ctx.message.author.id)] if len(entries) == 0: await ctx.send("You have no tags setup on this server!") else: pages = utils.Pages(self.bot, message=ctx.message, entries=entries) await pages.paginate() + else: + await ctx.send("There are no tags setup on this server!") @commands.group(invoke_without_command=True) @commands.guild_only() @@ -57,9 +57,9 @@ class Tags: EXAMPLE: !tag butts RESULT: Whatever you setup for the butts tag!!""" tag = tag.lower().strip() - tags = await utils.get_content('tags', str(ctx.message.guild.id)) - if tags and len(tags['tags']) > 0: - for t in tags['tags']: + tags = self.bot.db.load('tags', key=ctx.message.guild.id, pluck='tags') + if tags: + for t in tags: if t['trigger'].lower().strip() == tag: await ctx.send("\u200B{}".format(t['result'])) return @@ -67,7 +67,6 @@ class Tags: else: await ctx.send("There are no tags setup on this server!") - @tag.command(name='add', aliases=['create', 'setup']) @commands.guild_only() @utils.custom_perms(send_messages=True) @@ -76,8 +75,10 @@ class Tags: EXAMPLE: !tag add RESULT: A follow-along in order to create a new tag""" + def check(m): return m.channel == ctx.message.channel and m.author == ctx.message.author and len(m.content) > 0 + my_msg = await ctx.send("Ready to setup a new tag! What do you want the trigger for the tag to be?") try: @@ -92,12 +93,14 @@ class Tags: await ctx.send("Please keep tag triggers under 100 characters") return elif trigger in forbidden_tags: - await ctx.send("Sorry, but your tag trigger was detected to be forbidden. Current forbidden tag triggers are: \n{}".format("\n".join(forbidden_tags))) + await ctx.send( + "Sorry, but your tag trigger was detected to be forbidden. " + "Current forbidden tag triggers are: \n{}".format("\n".join(forbidden_tags))) return - tags = await utils.get_content('tags', str(ctx.message.guild.id)) - if tags and len(tags['tags']) > 0: - for t in tags['tags']: + tags = self.bot.db.load('tags', key=ctx.message.guild.id, pluck='tags') or [] + if tags: + for t in tags: if t['trigger'].lower().strip() == trigger: await ctx.send("There is already a tag setup called {}!".format(trigger)) return @@ -112,7 +115,9 @@ class Tags: await ctx.send("You can't create a tag with {}!".format(trigger)) return - my_msg = await ctx.send("Alright, your new tag can be called with {}!\n\nWhat do you want to be displayed with this tag?".format(trigger)) + my_msg = await ctx.send( + "Alright, your new tag can be called with {}!\n\nWhat do you want to be displayed with this tag?".format( + trigger)) try: msg = await self.bot.wait_for("message", check=check, timeout=60) @@ -133,13 +138,12 @@ class Tags: 'trigger': trigger, 'result': result } + tags.append(tag) entry = { 'server_id': str(ctx.message.guild.id), - 'tags': [tag] + 'tags': tags } - key = str(ctx.message.guild.id) - if not await utils.add_content('tags', entry): - await utils.update_content('tags', {'tags': r.row['tags'].append(tag)}, key) + self.bot.db.save('tags', entry) await ctx.send("I have just setup a new tag for this server! You can call your tag with {}".format(trigger)) @tag.command(name='edit') @@ -149,16 +153,17 @@ class Tags: """This will allow you to edit a tag that you have created EXAMPLE: !tag edit this tag RESULT: I'll ask what you want the new result to be""" - key = str(ctx.message.guild.id) - tags = await utils.get_content('tags', key) + tags = self.bot.db.load('tags', key=ctx.message.guild.id, pluck='tags') + def check(m): return m.channel == ctx.message.channel and m.author == ctx.message.author and len(m.content) > 0 - if tags and len(tags['tags']) > 0: - for i, t in enumerate(tags['tags']): + if tags: + for i, t in enumerate(tags): if t['trigger'] == tag: if t['author'] == str(ctx.message.author.id): - my_msg = await ctx.send("Alright, what do you want the new result for the tag {} to be".format(tag)) + my_msg = await ctx.send( + "Alright, what do you want the new result for the tag {} to be".format(tag)) try: msg = await self.bot.wait_for("message", check=check, timeout=60) except asyncio.TimeoutError: @@ -166,13 +171,17 @@ class Tags: return new_tag = t.copy() new_tag['result'] = msg.content - tags['tags'][i] = new_tag + tags[i] = new_tag try: await my_msg.delete() await msg.delete() except discord.Forbidden: pass - await utils.update_content('tags', tags, key) + entry = { + 'server_id': str(ctx.message.guild.id), + 'tags': tags + } + self.bot.db.save('tags', entry) await ctx.send("Alright, the tag {} has been updated".format(tag)) return else: @@ -182,7 +191,6 @@ class Tags: else: await ctx.send("There are no tags setup on this server!") - @tag.command(name='delete', aliases=['remove', 'stop']) @commands.guild_only() @utils.custom_perms(send_messages=True) @@ -192,14 +200,18 @@ class Tags: EXAMPLE: !tag delete stupid_tag RESULT: Deletes that stupid tag""" - key = str(ctx.message.guild.id) - tags = await utils.get_content('tags', key) - if tags and len(tags['tags']) > 0: - for t in tags['tags']: - if t['trigger'] == tag: - if ctx.message.author.permissions_in(ctx.message.channel).manage_guild or str(ctx.message.author.id) == t['author']: - tags['tags'].remove(t) - await utils.update_content('tags', tags, key) + tags = self.bot.db.load('tags', key=ctx.message.guild.id, pluck='tags') + if tags: + for t in tags: + if t['trigger'].lower().strip() == tag: + if ctx.message.author.permissions_in(ctx.message.channel).manage_guild or str( + ctx.message.author.id) == t['author']: + tags.remove(t) + entry = { + 'server_id': str(ctx.message.guild.id), + 'tags': tags + } + self.bot.db.save('tags', entry) await ctx.send("I have just removed the tag {}".format(tag)) else: await ctx.send("You don't own that tag! You can't remove it!") diff --git a/cogs/tictactoe.py b/cogs/tictactoe.py index cbf352d..edaf18c 100644 --- a/cogs/tictactoe.py +++ b/cogs/tictactoe.py @@ -191,7 +191,7 @@ class TicTacToe: await ctx.send("{} has won this game of TicTacToe, better luck next time {}".format(winner.display_name, loser.display_name)) # Handle updating ratings based on the winner and loser - await utils.update_records('tictactoe', winner, loser) + await utils.update_records('tictactoe', self.bot.db, winner, loser) # This game has ended, delete it so another one can be made del self.boards[ctx.message.guild.id] else: diff --git a/cogs/twitch.py b/cogs/twitch.py index d33ec7e..2418838 100644 --- a/cogs/twitch.py +++ b/cogs/twitch.py @@ -44,7 +44,7 @@ class Twitch: # Loop through as long as the bot is connected try: while not self.bot.is_closed(): - twitch = await utils.filter_content('twitch', {'notifications_on': 1}) + twitch = self.bot.db.load('twitch', table_filter={'notifications_on': 1}) for data in twitch: m_id = int(data['member_id']) url = data['twitch_url'] @@ -59,11 +59,8 @@ class Twitch: member = server.get_member(m_id) if member is None: continue - server_settings = await utils.get_content('server_settings', s_id) - if server_settings is not None: - channel_id = int(server_settings.get('notification_channel', s_id)) - else: - channel_id = int(s_id) + channel_id = self.bot.db.load('server_settings', key=s_id, + pluck='notifications_channel') or int(s_id) channel = server.get_channel(channel_id) if channel is None: continue @@ -71,7 +68,7 @@ class Twitch: await channel.send("{} has just gone live! View their stream at <{}>".format(member.display_name, data['twitch_url'])) except discord.Forbidden: pass - self.bot.loop.create_task(utils.update_content('twitch', {'live': 1}, str(m_id))) + self.bot.db.save('twitch', {'live': 1, 'member_id': str(m_id)}) elif not online and data['live'] == 1: for s_id in data['servers']: server = self.bot.get_guild(int(s_id)) @@ -80,17 +77,14 @@ class Twitch: member = server.get_member(m_id) if member is None: continue - server_settings = await utils.get_content('server_settings', s_id) - if server_settings is not None: - channel_id = int(server_settings.get('notification_channel', s_id)) - else: - channel_id = int(s_id) + channel_id = self.bot.db.load('server_settings', key=s_id, + pluck='notifications_channel') or int(s_id) channel = server.get_channel(channel_id) try: await channel.send("{} has just gone offline! View their stream next time at <{}>".format(member.display_name, data['twitch_url'])) except discord.Forbidden: pass - self.bot.loop.create_task(utils.update_content('twitch', {'live': 0}, str(m_id))) + self.bot.db.save('twitch', {'live': 0, 'member_id': str(m_id)}) await asyncio.sleep(30) except Exception as e: tb = traceback.format_exc() @@ -110,12 +104,11 @@ class Twitch: if member is None: member = ctx.message.author - result = await utils.get_content('twitch', str(member.id)) - if result is None: + url = self.bot.db.load('twitch', key=member.id, pluck='twitch_url') + if url is None: await ctx.send("{} has not saved their twitch URL yet!".format(member.name)) return - url = result['twitch_url'] user = re.search("(?<=twitch.tv/)(.*)", url).group(1) twitch_url = "https://api.twitch.tv/kraken/channels/{}".format(user) payload = {'client_id': self.key} @@ -170,18 +163,23 @@ class Twitch: return key = str(ctx.message.author.id) - entry = {'twitch_url': url, - 'servers': [str(ctx.message.guild.id)], - 'notifications_on': 1, - 'live': 0, - 'member_id': key} - update = {'twitch_url': url} - # Check to see if this user has already saved a twitch URL - # If they have, update the URL, otherwise create a new entry - # Assuming they're not live, and notifications should be on - if not await utils.add_content('twitch', entry): - await utils.update_content('twitch', update, key) + # Check if it exists first, if it does we don't want to override some of the settings + result = self.bot.db.load('twitch', key=key) + if result: + entry = { + 'twitch_url': url, + 'member_id': key + } + else: + entry = { + 'twitch_url': url, + 'servers': [str(ctx.message.guild.id)], + 'notifications_on': 1, + 'live': 0, + 'member_id': key + } + self.bot.db.save('twitch', entry) await ctx.send("I have just saved your twitch url {}".format(ctx.message.author.mention)) @twitch.command(name='remove', aliases=['delete']) @@ -192,8 +190,12 @@ class Twitch: EXAMPLE: !twitch remove RESULT: I stop saving your twitch URL""" - # Just try to remove it, if it doesn't exist, nothing is going to happen - await utils.remove_content('twitch', str(ctx.message.author.id)) + entry = { + 'twitch_url': None, + 'member_id': str(ctx.message.author.id) + } + + self.bot.db.save('twitch', entry) await ctx.send("I am no longer saving your twitch URL {}".format(ctx.message.author.mention)) @twitch.group(invoke_without_command=True) @@ -206,17 +208,22 @@ class Twitch: EXAMPLE: !twitch notify RESULT: This server will now be notified when you go live""" key = str(ctx.message.author.id) - result = await utils.get_content('twitch', key) + servers = self.bot.db.load('twitch', key=key, pluck='servers') # Check if this user is saved at all - if result is None: + if servers is None: await ctx.send( "I do not have your twitch URL added {}. You can save your twitch url with !twitch add".format( ctx.message.author.mention)) # Then check if this server is already added as one to notify in - elif str(ctx.message.guild.id) in result['servers']: + elif str(ctx.message.guild.id) in servers: await ctx.send("I am already set to notify in this server...") else: - await utils.update_content('twitch', {'servers': r.row['servers'].append(str(ctx.message.guild.id))}, key) + servers.append(str(ctx.message.guild.id)) + entry = { + 'member_id': key, + 'servers': servers + } + self.bot.db.save('twitch', entry) await ctx.send("This server will now be notified if you go live") @notify.command(name='on', aliases=['start,yes']) @@ -227,7 +234,14 @@ class Twitch: EXAMPLE: !twitch notify on RESULT: Notifications will be sent when you go live""" - if await utils.update_content('twitch', {"notifications_on": 1}, str(ctx.message.author.id)): + key = str(ctx.message.author.id) + result = self.bot.db.load('twitch', key=key) + if result: + entry = { + 'member_id': key, + 'notifications_on': 1 + } + self.bot.db.save('twitch', entry) await ctx.send("I will notify if you go live {}, you'll get a bajillion followers I promise c:".format( ctx.message.author.mention)) else: @@ -241,7 +255,14 @@ class Twitch: EXAMPLE: !twitch notify off RESULT: Notifications will not be sent when you go live""" - if await utils.update_content('twitch', {"notifications_on": 0}, str(ctx.message.author.id)): + key = str(ctx.message.author.id) + result = self.bot.db.load('twitch', key=key) + if result: + entry = { + 'member_id': key, + 'notifications_on': 0 + } + self.bot.db.save('twitch', entry) await ctx.send( "I will not notify if you go live anymore {}, " "are you going to stream some lewd stuff you don't want people to see?~".format( diff --git a/cogs/utils/__init__.py b/cogs/utils/__init__.py index 7a3867b..6e9baee 100644 --- a/cogs/utils/__init__.py +++ b/cogs/utils/__init__.py @@ -4,3 +4,4 @@ from .config import * from .utilities import * from .images import create_banner from .paginator import Pages, CannotPaginate, DetailedPages +from .database import DB diff --git a/cogs/utils/checks.py b/cogs/utils/checks.py index bc3f47c..58dc181 100644 --- a/cogs/utils/checks.py +++ b/cogs/utils/checks.py @@ -12,7 +12,6 @@ required_tables = { 'battle_records': 'member_id', 'boops': 'member_id', 'command_usage': 'command', - 'motd': 'date', 'overwatch': 'member_id', 'picarto': 'member_id', 'server_settings': 'server_id', @@ -21,7 +20,8 @@ required_tables = { 'osu': 'member_id', 'tags': 'server_id', 'tictactoe': 'member_id', - 'twitch': 'member_id' + 'twitch': 'member_id', + 'user_playlists': 'member_id' } @@ -68,16 +68,13 @@ def is_owner(ctx): return ctx.bot.owner.id == ctx.message.author.id -def should_ignore(message): +def should_ignore(bot, message): if message.guild is None: return False - try: - server_settings = config.cache.get('server_settings').values - ignored = [x for x in server_settings if x['server_id'] == str( - message.guild.id)][0]['ignored'] - return str(message.author.id) in ignored['members'] or str(message.channel.id) in ignored['channels'] - except (TypeError, IndexError, KeyError): + ignored = bot.db.load('server_settings', key=message.guild.id, pluck='ignored') + if not ignored: return False + return str(message.author.id) in ignored['members'] or str(message.channel.id) in ignored['channels'] def custom_perms(**perms): @@ -94,13 +91,10 @@ def custom_perms(**perms): for perm, setting in perms.items(): setattr(required_perm, perm, setting) - try: - server_settings = config.cache.get('server_settings').values - required_perm_value = [x for x in server_settings if x['server_id'] == str( - ctx.message.guild.id)][0]['permissions'][ctx.command.qualified_name] + required_perm_value = ctx.bot.db.load('server_settings', key=ctx.message.guild.id, pluck='permissions') or {} + required_perm_value = required_perm_value.get(ctx.command.qualified_name) + if required_perm_value: required_perm = discord.Permissions(required_perm_value) - except (TypeError, IndexError, KeyError): - pass # Now just check if the person running the command has these permissions return member_perms >= required_perm diff --git a/cogs/utils/config.py b/cogs/utils/config.py index dd52d5c..731bbad 100644 --- a/cogs/utils/config.py +++ b/cogs/utils/config.py @@ -23,20 +23,6 @@ except KeyError: quit() -# This is a simple class for the cache concept, all it holds is it's own key and the values -# With a method that gets content based on it's key -class Cache: - def __init__(self, key): - self.key = key - self.values = {} - self.refreshed = pendulum.utcnow() - loop.create_task(self.update()) - - async def update(self): - self.values = await get_content(self.key) - self.refreshed = pendulum.utcnow() - - # Default bot's description bot_description = global_config.get("description") # Bot's default prefix for commands @@ -80,7 +66,9 @@ extensions = [ 'cogs.tags', 'cogs.roulette', 'cogs.music', - 'cogs.polls' + 'cogs.polls', + 'cogs.playlist', + 'cogs.dj' ] @@ -103,155 +91,8 @@ db_pass = global_config.get('db_pass', '') # {'ca_certs': db_cert}, 'user': db_user, 'password': db_pass} db_opts = {'host': db_host, 'db': db_name, 'port': db_port, 'user': db_user, 'password': db_pass} -possible_keys = ['prefixes', 'battle_records', 'boops', 'server_alerts', 'user_notifications', 'nsfw_channels', - 'custom_permissions', 'rules', 'overwatch', 'picarto', 'twitch', 'strawpolls', 'tags', - 'tictactoe', 'bot_data', 'command_manage'] - -# This will be a dictionary that holds the cache object, based on the key that is saved -cache = {} - -# Populate cache with each object -# With the new saving method, we're not going to be able to cache the way that I was before -# This is on standby until I rethink how to do this, because I do still want to cache data -"""for k in possible_keys: - ca che[k] = Cache(k)""" - -# We still need 'cache' for prefixes and custom permissions however, so for now, just include that -cache['prefixes'] = Cache('prefixes') -cache['server_settings'] = Cache('server_settings') - -async def update_cache(): - for value in cache.values(): - await value.update() - def command_prefix(bot, message): - # We do not want to make a query for every message that is sent - # So assume it's in cache, or it doesn't exist - # If the prefix does exist in the database and isn't in our cache; too bad, something has messed up - # But it is not worth a query for every single message the bot detects, to fix - try: - prefixes = cache['server_settings'].values - prefix = [x for x in prefixes if x['server_id'] == str(message.guild.id)][0]['prefix'] - return prefix or default_prefix - except (KeyError, TypeError, IndexError, AttributeError, KeyError): + if not message.guild: return default_prefix - - -async def add_content(table, content): - r.set_loop_type("asyncio") - conn = await r.connect(**db_opts) - # First we need to make sure that this entry doesn't exist - # For all rethinkDB cares, multiple entries can exist with the same content - # For our purposes however, we do not want this - try: - result = await r.table(table).insert(content).run(conn) - except r.ReqlOpFailedError: - # This means the table does not exist - await r.table_create(table).run(conn) - await r.table(table).insert(content).run(conn) - result = {} - - await conn.close() - if table == 'prefixes' or table == 'server_settings': - loop.create_task(cache[table].update()) - return result.get('inserted', 0) > 0 - - -async def remove_content(table, key): - r.set_loop_type("asyncio") - conn = await r.connect(**db_opts) - try: - result = await r.table(table).get(key).delete().run(conn) - except r.ReqlOpFailedError: - result = {} - pass - - await conn.close() - if table == 'prefixes' or table == 'server_settings': - loop.create_task(cache[table].update()) - return result.get('deleted', 0) > 0 - - -async def update_content(table, content, key): - r.set_loop_type("asyncio") - conn = await r.connect(**db_opts) - # This method is only for updating content, so if we find that it doesn't exist, just return false - try: - # Update based on the content and filter passed to us - # rethinkdb allows you to do many many things inside of update - # This is why we're accepting a variable and using it, whatever it may be, as the query - result = await r.table(table).get(key).update(content).run(conn) - except r.ReqlOpFailedError: - result = {} - - await conn.close() - if table == 'prefixes' or table == 'server_settings': - loop.create_task(cache[table].update()) - return result.get('replaced', 0) > 0 or result.get('unchanged', 0) > 0 - - -async def replace_content(table, content, key): - # This method is here because .replace and .update can have some different functionalities - r.set_loop_type("asyncio") - conn = await r.connect(**db_opts) - try: - result = await r.table(table).get(key).replace(content).run(conn) - except r.ReqlOpFailedError: - result = {} - - await conn.close() - if table == 'prefixes' or table == 'server_settings': - loop.create_task(cache[table].update()) - return result.get('replaced', 0) > 0 or result.get('unchanged', 0) > 0 - - -async def get_content(table, key=None): - r.set_loop_type("asyncio") - conn = await r.connect(**db_opts) - - try: - if key: - cursor = await r.table(table).get(key).run(conn) - else: - cursor = await r.table(table).run(conn) - if cursor is None: - content = None - elif type(cursor) is not dict: - content = await _convert_to_list(cursor) - if len(content) == 0: - content = None - else: - content = cursor - except (IndexError, r.ReqlOpFailedError): - content = None - - await conn.close() - return content - -async def filter_content(table: str, r_filter): - r.set_loop_type("asyncio") - conn = await r.connect(**db_opts) - try: - cursor = await r.table(table).filter(r_filter).run(conn) - content = await _convert_to_list(cursor) - if len(content) == 0: - content = None - except (IndexError, r.ReqlOpFailedError): - content = None - - await conn.close() - return content - - -async def _convert_to_list(cursor): - # This method is here because atm, AsyncioCursor is not iterable - # For our purposes, we want a list, so we need to do this manually - cursor_list = [] - while True: - try: - val = await cursor.next() - cursor_list.append(val) - except r.ReqlCursorEmpty: - break - return cursor_list + return bot.db.load('server_settings', key=message.guild.id, pluck='prefix') or default_prefix diff --git a/cogs/utils/database.py b/cogs/utils/database.py new file mode 100644 index 0000000..cdd853a --- /dev/null +++ b/cogs/utils/database.py @@ -0,0 +1,141 @@ +import asyncio +import rethinkdb as r +from datetime import datetime +from .checks import required_tables + +from . import config + + +async def _convert_to_list(cursor): + # This method is here because atm, AsyncioCursor is not iterable + # For our purposes, we want a list, so we need to do this manually + cursor_list = [] + while True: + try: + val = await cursor.next() + cursor_list.append(val) + except r.ReqlCursorEmpty: + break + return cursor_list + + +class Cache: + """A class to hold the cached database entries""" + + def __init__(self, table, key, db, loop): + self.table = table # The name of the database table + self.key = key # The name of primary key + self.db = db # The database class connections are made through + self.loop = loop + self.values = [] # The values returned from the database + self.refreshed_time = None + self.loop.create_task(self.check_refresh()) + + async def refresh(self): + self.values = await self.db.actual_load(self.table) + self.refreshed_time = datetime.now() + + async def check_refresh(self): + if self.refreshed_time is None: + await self.refresh() + else: + difference = datetime.now() - self.refreshed_time + if difference.total_seconds() > 300: + await self.refresh() + + self.loop.call_later(60, self.check_refresh()) + + def get(self, key=None, table_filter=None, pluck=None): + """This simulates the database call, to make it easier to get the data""" + if key is None and table_filter is None: + return self.values + elif key: + for value in self.values: + if value[self.key] == key: + if pluck: + return value.get(pluck) + else: + return value + elif table_filter: + req_key = list(table_filter.keys())[0] + req_val = list(table_filter.values())[0] + for value in self.values: + if value[req_key] == req_val: + if pluck: + return value.get(pluck) + else: + return value + + +class DB: + def __init__(self): + self.loop = asyncio.get_event_loop() + self.opts = config.db_opts + self.cache = {} + + for table, key in required_tables.items(): + self.cache[table] = Cache(table, key, self, self.loop) + + async def query(self, query): + """Lets you run a manual query""" + r.set_loop_type("asyncio") + conn = await r.connect(**self.opts) + try: + cursor = await query.run(conn) + except (r.ReqlOpFailedError, r.ReqlNonExistenceError): + cursor = None + if isinstance(cursor, r.Cursor): + cursor = await _convert_to_list(cursor) + await conn.close() + return cursor + + def save(self, table, content): + """A synchronous task to throw saving content into a task""" + self.loop.create_task(self._save(table, content)) + + async def _save(self, table, content): + """Saves data in the table""" + + index = await self.query(r.table(table).info()) + index = index.get("primary_key") + key = content.get(index) + if key: + cur_content = await self.query(r.table(table).get(key)) + if cur_content: + # We have content...we either need to update it, or replace + # Update will typically be more common so lets try that first + result = await self.query(r.table(table).get(key).update(content)) + print(result) + if result.get('replaced', 0) == 0 and result.get('unchanged', 0) == 0: + print("Replacing...") + await self.query(r.table(table).get(key).replace(content)) + else: + await self.query(r.table(table).insert(content)) + else: + await self.query(r.table(table).insert(content)) + + await self.cache.get(table).refresh() + + def load(self, table, **kwargs): + if kwargs.get('key'): + kwargs['key'] = str(kwargs.get('key')) + return self.cache.get(table).get(**kwargs) + + async def actual_load(self, table, key=None, table_filter=None, pluck=None): + """Loads the specified content from the specific table""" + query = r.table(table) + + # If a key has been provided, get content with that key + if key: + query = query.get(str(key)) + # A key and a filter shouldn't be combined for any case we'll ever use, so seperate these + elif table_filter: + query = query.filter(table_filter) + + # If we want to pluck something specific, do that + if pluck: + query = query.pluck(pluck).values()[0] + + cursor = await self.query(query) + + return cursor diff --git a/cogs/utils/utilities.py b/cogs/utils/utilities.py index 2e297f8..671c1fd 100644 --- a/cogs/utils/utilities.py +++ b/cogs/utils/utilities.py @@ -33,7 +33,7 @@ def get_all_subcommands(command): yield from get_all_subcommands(subcmd) -async def channel_is_nsfw(channel): +async def channel_is_nsfw(channel, db): if type(channel) is discord.DMChannel: server = 'DMs' elif channel.is_nsfw(): @@ -43,12 +43,8 @@ async def channel_is_nsfw(channel): channel = str(channel.id) - server_settings = await config.get_content('server_settings', server) - - try: - return channel in server_settings['nsfw_channels'] - except (TypeError, IndexError, KeyError): - return False + channels = db.load('server_settings', key=server, pluck='nsfw_channels') or [] + return channel in channels async def download_image(url): @@ -101,11 +97,11 @@ async def request(url, *, headers=None, payload=None, method='GET', attr='json') continue -async def update_records(key, winner, loser): +async def update_records(key, db, winner, loser): # We're using the Harkness scale to rate # http://opnetchessclub.wikidot.com/harkness-rating-system r_filter = lambda row: (row['member_id'] == str(winner.id)) | (row['member_id'] == str(loser.id)) - matches = await config.filter_content(key, r_filter) + matches = db.load(key, table_filter=r_filter) winner_stats = {} loser_stats = {} @@ -150,12 +146,8 @@ async def update_records(key, winner, loser): loser_losses += 1 # Now save the new wins, losses, and ratings - winner_stats = {'wins': winner_wins, 'losses': winner_losses, 'rating': winner_rating} - loser_stats = {'wins': loser_wins, 'losses': loser_losses, 'rating': loser_rating} + winner_stats = {'wins': winner_wins, 'losses': winner_losses, 'rating': winner_rating, 'member_id': str(winner.id)} + loser_stats = {'wins': loser_wins, 'losses': loser_losses, 'rating': loser_rating, 'member_id': str(loser.id)} - if not await config.update_content(key, winner_stats, str(winner.id)): - winner_stats['member_id'] = str(winner.id) - await config.add_content(key, winner_stats) - if not await config.update_content(key, loser_stats, str(loser.id)): - loser_stats['member_id'] = str(loser.id) - await config.add_content(key, loser_stats) + db.save(key, winner_stats) + db.save(key, loser_stats) diff --git a/cogs/voice_utilities/__init__.py b/cogs/voice_utilities/__init__.py index 3161697..c60540c 100644 --- a/cogs/voice_utilities/__init__.py +++ b/cogs/voice_utilities/__init__.py @@ -1,3 +1,4 @@ from .downloader import Downloader -from .playlist import Playlist from .exceptions import * +from .playlist import Playlist + diff --git a/cogs/voice_utilities/entry.py b/cogs/voice_utilities/entry.py index 3803d29..3be0d51 100644 --- a/cogs/voice_utilities/entry.py +++ b/cogs/voice_utilities/entry.py @@ -92,7 +92,7 @@ class BasePlaylistEntry: class URLPlaylistEntry(BasePlaylistEntry): - def __init__(self, playlist, url, title, ctx, thumbnail, duration=0, expected_filename=None, **meta): + def __init__(self, playlist, url, title, thumbnail, duration=0, expected_filename=None, **meta): super().__init__() self.playlist = playlist @@ -102,8 +102,6 @@ class URLPlaylistEntry(BasePlaylistEntry): self.thumbnail = thumbnail self.expected_filename = expected_filename self.meta = meta - self.requester = ctx.message.author - self.channel = ctx.message.channel self.download_folder = self.playlist.downloader.download_folder def __str__(self): @@ -119,7 +117,7 @@ class URLPlaylistEntry(BasePlaylistEntry): @property def progress(self): - if self.start_time: + if hasattr(self, 'start_time') and self.start_time: return round(time.time() - self.start_time) @property @@ -132,7 +130,6 @@ class URLPlaylistEntry(BasePlaylistEntry): @classmethod def from_json(cls, playlist, jsonstring): data = json.loads(jsonstring) - print(data) # TODO: version check url = data['url'] title = data['title'] @@ -187,7 +184,6 @@ class URLPlaylistEntry(BasePlaylistEntry): # the generic extractor requires special handling if extractor == 'generic': - # print("Handling generic") flistdir = [f.rsplit('-', 1)[0] for f in os.listdir(self.download_folder)] expected_fname_noex, fname_ex = os.path.basename(self.expected_filename).rsplit('.', 1) @@ -202,18 +198,14 @@ class URLPlaylistEntry(BasePlaylistEntry): os.listdir(self.download_folder)[flistdir.index(expected_fname_noex)] ) - # print("Resolved %s to %s" % (self.expected_filename, lfile)) lsize = os.path.getsize(lfile) - # print("Remote size: %s Local size: %s" % (rsize, lsize)) if lsize != rsize: await self._really_download(hash=True) else: - # print("[Download] Cached:", self.url) self.filename = lfile else: - # print("File not found in cache (%s)" % expected_fname_noex) await self._really_download(hash=True) else: @@ -227,15 +219,9 @@ class URLPlaylistEntry(BasePlaylistEntry): if expected_fname_base in ldir: self.filename = os.path.join(self.download_folder, expected_fname_base) - print("[Download] Cached:", self.url) elif expected_fname_noex in flistdir: - print("[Download] Cached (different extension):", self.url) self.filename = os.path.join(self.download_folder, ldir[flistdir.index(expected_fname_noex)]) - print("Expected %s, got %s" % ( - self.expected_filename.rsplit('.', 1)[-1], - self.filename.rsplit('.', 1)[-1] - )) else: await self._really_download() @@ -252,15 +238,12 @@ class URLPlaylistEntry(BasePlaylistEntry): # noinspection PyShadowingBuiltins async def _really_download(self, *, hash=False): - print("[Download] Started:", self.url) try: result = await self.playlist.downloader.extract_info(self.playlist.loop, self.url, download=True) except Exception as e: raise ExtractionError(e) - print("[Download] Complete:", self.url) - if result is None: raise ExtractionError("ytdl broke and hell if I know why") # What the fuck do I do now? @@ -293,4 +276,4 @@ class URLPlaylistEntry(BasePlaylistEntry): fmt = "{0[0]}m {0[1]}s".format(length) embed.add_field(name='Duration', value=fmt, inline=False) # And return the embed we created - return embed + return embed \ No newline at end of file diff --git a/cogs/voice_utilities/playlist.py b/cogs/voice_utilities/playlist.py index 7ff4599..ec2ea49 100644 --- a/cogs/voice_utilities/playlist.py +++ b/cogs/voice_utilities/playlist.py @@ -37,7 +37,7 @@ class Playlist(EventEmitter): else: return 0 - async def add_entry(self, song_url, ctx, **meta): + async def add_entry(self, song_url, **meta): """ Validates and adds a song_url to be played. This does not start the download of the song. @@ -76,10 +76,8 @@ class Playlist(EventEmitter): # https://github.com/KeepSafe/aiohttp/issues/758 # https://github.com/KeepSafe/aiohttp/issues/852 content_type = await get_header(self.bot.aiosession, info['url'], 'CONTENT-TYPE') - print("Got content type", content_type) except Exception as e: - print("[Warning] Failed to get content type for url %s (%s)" % (song_url, e)) content_type = None if content_type: @@ -87,9 +85,6 @@ class Playlist(EventEmitter): if '/ogg' not in content_type: # How does a server say `application/ogg` what the actual fuck raise ExtractionError("Invalid content type \"%s\" for url %s" % (content_type, song_url)) - elif not content_type.startswith(('audio/', 'video/')): - print("[Warning] Questionable content type \"%s\" for url %s" % (content_type, song_url)) - if info.get('is_live', False): raise LiveStreamError("Cannot download from a livestream") @@ -97,7 +92,6 @@ class Playlist(EventEmitter): self, song_url, info.get('title', 'Untitled'), - ctx, info.get('thumbnail', None), info.get('duration', 0) or 0, self.downloader.ytdl.prepare_filename(info), @@ -151,14 +145,9 @@ class Playlist(EventEmitter): baditems += 1 # Once I know more about what's happening here I can add a proper message traceback.print_exc() - print(items) - print("Could not add item") else: baditems += 1 - if baditems: - print("Skipped %s bad entries" % baditems) - return entry_list, position async def async_process_youtube_playlist(self, playlist_url, **meta): @@ -191,14 +180,9 @@ class Playlist(EventEmitter): baditems += 1 except Exception as e: baditems += 1 - print("There was an error adding the song {}: {}: {}\n".format( - entry_data['id'], e.__class__.__name__, e)) else: baditems += 1 - if baditems: - print("Skipped %s bad entries" % baditems) - return gooditems async def async_process_sc_bc_playlist(self, playlist_url, **meta): @@ -230,14 +214,9 @@ class Playlist(EventEmitter): baditems += 1 except Exception as e: baditems += 1 - print("There was an error adding the song {}: {}: {}\n".format( - entry_data['id'], e.__class__.__name__, e)) else: baditems += 1 - if baditems: - print("Skipped %s bad entries" % baditems) - return gooditems def _add_entry(self, entry): @@ -286,4 +265,4 @@ class Playlist(EventEmitter): return datetime.timedelta(seconds=estimated_time) def count_for_user(self, user): - return sum(1 for e in self.entries if e.meta.get('author', None) == user) + return sum(1 for e in self.entries if e.meta.get('author', None) == user) \ No newline at end of file diff --git a/requirements.txt b/requirements.txt index fda0300..e9f5e4f 100644 --- a/requirements.txt +++ b/requirements.txt @@ -1,9 +1,9 @@ -beautifulsoup4 +beautifulsoup4==4.6.0 Pillow==3.4.1 -rethinkdb -ruamel.yaml +rethinkdb==2.3.0.post6 +ruamel.yaml==0.14.12 youtube-dl -psutil -pendulum +psutil==5.2.2 +pendulum==1.2.0 -e git+https://github.com/Rapptz/discord.py@rewrite#egg=discord.py[voice] --e git+https://github.com/khazhyk/osuapi.git#egg=osuapi +-e git+https://github.com/khazhyk/osuapi.git#egg=osuapi \ No newline at end of file