1
0
Fork 0
mirror of synced 2024-06-08 13:44:32 +12:00

Update to saving method, using a much more robust method

This commit is contained in:
phxntxm 2016-09-28 18:39:34 -05:00
parent 56be540d55
commit 26f32f9967
14 changed files with 574 additions and 667 deletions

View file

@ -18,10 +18,9 @@ class Core:
def __init__(self, bot):
self.bot = bot
@commands.command()
@checks.custom_perms(send_messages=True)
async def calendar(self, month: str=None, year: int=None):
async def calendar(self, month: str = None, year: int = None):
"""Provides a printout of the current month's calendar
Provide month and year to print the calendar of that year and month"""
@ -66,13 +65,11 @@ class Core:
fmt = {}
bot_data = await config.get_content('bot_data')
total_data = {}
for shard, values in bot_data.items():
for key, value in values.items():
if key in total_data:
total_data[key] += value
else:
total_data[key] = value
total_data = {'member_count': 0,
'server_count': 0}
for entry in bot_data:
total_data['member_count'] += entry['member_count']
total_data['server_count'] += entry['server_count']
# We can pretty safely assume that the author is going to be in at least one channel with the bot
# So find the author based on that list
@ -88,7 +85,7 @@ class Core:
hm_games = len(
[server_id for server_id, game in self.bot.get_cog('Hangman').games.items()])
ttt_games = len([server_id for server_id,
game in self.bot.get_cog('TicTacToe').boards.items()])
game in self.bot.get_cog('TicTacToe').boards.items()])
count_battles = 0
for battles in self.bot.get_cog('Interaction').battles.values():
count_battles += len(battles)
@ -169,7 +166,7 @@ class Core:
@commands.command(pass_context=True)
@checks.custom_perms(send_messages=True)
async def roll(self, ctx, notation: str="d6"):
async def roll(self, ctx, notation: str = "d6"):
"""Rolls a die based on the notation given
Format should be #d#"""
# Use regex to get the notation based on what was provided

View file

@ -6,55 +6,6 @@ import discord
import random
async def update_battle_records(winner, loser):
# We're using the Harkness scale to rate
# http://opnetchessclub.wikidot.com/harkness-rating-system
battles = await config.get_content('battle_records')
# Start ratings at 1000 if they have no rating
winner_stats = battles.get(winner.id) or {}
winner_rating = winner_stats.get('rating') or 1000
loser_stats = battles.get(loser.id) or {}
loser_rating = loser_stats.get('rating') or 1000
# The scale is based off of increments of 25, increasing the change by 1 for each increment
# That is all this loop does, increment the "change" for every increment of 25
# The change caps off at 300 however, so break once we are over that limit
difference = abs(winner_rating - loser_rating)
rating_change = 0
count = 25
while count <= difference:
if count > 300:
break
rating_change += 1
count += 25
# 16 is the base change, increased or decreased based on whoever has the higher current rating
if winner_rating > loser_rating:
winner_rating += 16 - rating_change
loser_rating -= 16 - rating_change
else:
winner_rating += 16 + rating_change
loser_rating -= 16 + rating_change
# Just increase wins/losses for each person, making sure it's at least 0
winner_wins = winner_stats.get('wins') or 0
winner_losses = winner_stats.get('losses') or 0
loser_wins = loser_stats.get('wins') or 0
loser_losses = loser_stats.get('losses') or 0
winner_wins += 1
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}
battles[winner.id] = winner_stats
battles[loser.id] = loser_stats
return await config.save_content('battle_records', battles)
class Interaction:
"""Commands that interact with another user"""
@ -139,10 +90,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 self.bot.say(fmt.format(battleP1.mention, battleP2.mention))
await update_battle_records(battleP1, battleP2)
await config.update_records('battle_records', battleP1, battleP2)
else:
await self.bot.say(fmt.format(battleP2.mention, battleP1.mention))
await update_battle_records(battleP2, battleP1)
await config.update_records('battle_records', battleP2, battleP1)
@commands.command(pass_context=True, no_pm=True)
@checks.custom_perms(send_messages=True)
@ -182,17 +133,22 @@ class Interaction:
await self.bot.say("Why the heck are you booping me? Get away from me >:c")
return
boops = await config.get_content('boops')
r_filter = {'member_id': booper.id}
boops = await config.get_content('boops', r_filter)
if boops is not None:
boops = boops[0]['boops']
# If the booper has never booped the member provided, assure it's 0
amount = boops.get(boopee.id, 0) + 1
boops[boopee.id] = amount
# Get all the booped stats for the author
# Set to default as having just booped boopee 0 times, so that we can increment that
booper_boops = boops.get(ctx.message.author.id, {boopee.id: 0})
# If the booper has never booped the member provided, assume 0 like above so we can increment like normal
amount = booper_boops.get(boopee.id, 0) + 1
booper_boops[boopee.id] = amount
boops[ctx.message.author.id] = booper_boops
await config.update_content('boops', {'boops': boops}, r_filter)
else:
entry = {'member_id': booper.id,
'boops': {boopee.id: 1}}
await config.add_content('boops', entry, r_filter)
amount = 1
await config.save_content('boops', boops)
fmt = "{0.mention} has just booped you {1.mention}! That's {2} times now!"
await self.bot.say(fmt.format(booper, boopee, amount))

View file

@ -69,12 +69,13 @@ class Links:
if len(search) > 0:
# This sets the url as url?q=search+terms
url = 'https://derpibooru.org/search.json?q={}'.format('+'.join(search))
nsfw_channels = await config.get_content("nsfw_channels")
nsfw_channels = nsfw_channels.get('registered') or []
r_filter = {'channel_id': ctx.message.channel.id}
nsfw_channels = await config.get_content("nsfw_channels", r_filter)
# 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
if ctx.message.channel.id in nsfw_channels:
if nsfw_channels is not None:
url += ",+%28explicit+OR+suggestive%29&filter_id=95938"
else:
url += ",+safe"
@ -126,11 +127,11 @@ class Links:
# Due to this, send a message saying we're looking up the information first
await self.bot.say("Looking up an image with those tags....")
nsfw_channels = await config.get_content("nsfw_channels")
nsfw_channels = nsfw_channels.get('registered') or []
r_filter = {'channel_id': ctx.message.channel.id}
nsfw_channels = await config.get_content("nsfw_channels", r_filter)
# e621 by default does not filter explicit content, so tack on
# safe/explicit based on if this channel is nsfw or not
if ctx.message.channel.id in nsfw_channels:
if nsfw_channels is not None:
url += "%20rating:explicit"
else:
url += "%20rating:safe"

View file

@ -1,9 +1,11 @@
from discord.ext import commands
from .utils import checks
from .utils import config
import discord
import re
import asyncio
import rethinkdb as r
valid_perms = [p for p in dir(discord.Permissions) if isinstance(getattr(discord.Permissions, p), property)]
@ -14,15 +16,39 @@ class Mod:
def __init__(self, bot):
self.bot = bot
def find_command(self, command):
# This method ensures the command given is valid. We need to loop through commands
# As self.bot.commands only includes parent commands
# So we are splitting the command in parts, looping through the commands
# And getting the subcommand based on the next part
# If we try to access commands of a command that isn't a group
# We'll hit an AttributeError, meaning an invalid command was given
# If we loop through and don't find anything, cmd will still be None
# And we'll report an invalid was given as well
cmd = None
for part in command.split():
try:
if cmd is None:
cmd = self.bot.commands.get(part)
else:
cmd = cmd.commands.get(part)
except AttributeError:
cmd = None
break
return cmd
@commands.command(pass_context=True, no_pm=True)
@checks.custom_perms(kick_members=True)
async def alerts(self, ctx, channel: discord.Channel):
"""This command is used to set a channel as the server's 'notifications' channel
Any notifications (like someone going live on Twitch, or Picarto) will go to that channel"""
server_alerts = await config.get_content('server_alerts')
# This will update/add the channel if an entry for this server exists or not
server_alerts[ctx.message.server.id] = channel.id
await config.save_content('server_alerts', server_alerts)
r_filter = {'server_id': ctx.message.server.id}
entry = {'server_id': ctx.message.server.id,
'channel_id': ctx.message.channel.id}
if not await config.add_content('server_alerts', entry, r_filter):
await config.update_content('server_alerts', entry, r_filter)
await self.bot.say("I have just changed this server's 'notifications' channel"
"\nAll notifications will now go to `{}`".format(channel))
@ -36,9 +62,11 @@ class Mod:
# 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 = ctx.message.channel.id if re.search("(on|yes|true)", on_off.lower()) else None
notifications = await config.get_content('user_notifications')
notifications[ctx.message.server.id] = on_off
await config.save_content('user_notifications', notifications)
r_filter = {'server_id': ctx.message.server.id}
entry = {'server_id': ctx.message.server.id,
'channel_id': on_off}
if not await config.add_content('user_notifications', entry, r_filter):
await config.update_content('user_notifications', entry, r_filter)
fmt = "notify" if on_off else "not notify"
await self.bot.say("This server will now {} if someone has joined or left".format(fmt))
@ -53,29 +81,21 @@ class Mod:
@checks.custom_perms(kick_members=True)
async def nsfw_add(self, ctx):
"""Registers this channel as a 'nsfw' channel"""
nsfw_channels = await config.get_content('nsfw_channels')
# rethinkdb cannot save a list as a field, so we need a dict with one elemtn to store our list
nsfw_channels = nsfw_channels.get('registered') or []
if ctx.message.channel.id in nsfw_channels:
await self.bot.say("This channel is already registered as 'nsfw'!")
else:
# Append instead of setting to a certain channel, so that multiple channels can be nsfw
nsfw_channels.append(ctx.message.channel.id)
await config.save_content('nsfw_channels', {'registered': nsfw_channels})
r_filter = {'channel_id': ctx.message.channel.id}
if await config.add_content('nsfw_channels', r_filter, r_filter):
await self.bot.say("This channel has just been registered as 'nsfw'! Have fun you naughties ;)")
else:
await self.bot.say("This channel is already registered as 'nsfw'!")
@nsfw.command(name="remove", aliases=["delete"], pass_context=True, no_pm=True)
@checks.custom_perms(kick_members=True)
async def nsfw_remove(self, ctx):
"""Removes this channel as a 'nsfw' channel"""
nsfw_channels = await config.get_content('nsfw_channels')
nsfw_channels = nsfw_channels.get('registered') or []
if ctx.message.channel.id not in nsfw_channels:
await self.bot.say("This channel is not registered as a ''nsfw' channel!")
else:
nsfw_channels.remove(ctx.message.channel.id)
await config.save_content('nsfw_channels', {'registered': nsfw_channels})
r_filter = {'channel_id': ctx.message.channel.id}
if await config.remove_content('nsfw_channels', r_filter):
await self.bot.say("This channel has just been unregistered as a nsfw channel")
else:
await self.bot.say("This channel is not registered as a ''nsfw' channel!")
@commands.command(pass_context=True, no_pm=True)
@checks.custom_perms(kick_members=True)
@ -98,21 +118,11 @@ class Mod:
"Valid permissions are: ```\n{}```".format("\n".join("{}".format(i) for i in valid_perms)))
return
custom_perms = await config.get_content('custom_permissions')
server_perms = custom_perms.get(ctx.message.server.id) or {}
r_filter = {'server_id': ctx.message.server.id}
server_perms = await config.get_content('custom_permissions', r_filter)
server_perms = server_perms[0]
cmd = None
# This is the same loop as the add command, we need this to get the
# command object so we can get the qualified_name
for part in command.split():
try:
if cmd is None:
cmd = self.bot.commands.get(part)
else:
cmd = cmd.commands.get(part)
except AttributeError:
cmd = None
break
cmd = self.find_command(command)
if cmd is None:
await self.bot.say("That is not a valid command!")
@ -126,8 +136,7 @@ class Mod:
custom_perms = [func for func in cmd.checks if "custom_perms" in func.__qualname__][0]
except IndexError:
# Loop through and check if there is a check called is_owner
# Ff we loop through and don't find one
# This means that the only other choice is to be
# If we loop through and don't find one, this means that the only other choice is to be
# Able to manage the server (for the checks on perm commands)
for func in cmd.checks:
if "is_owner" in func.__qualname__:
@ -137,7 +146,7 @@ class Mod:
"You are required to have `manage_server` permissions to run `{}`".format(cmd.qualified_name))
return
# Perms will be an attribute if custom_perms is found no matter what, so need to check this
# Perms will be an attribute if custom_perms is found no matter what, so no need to check this
perms = "\n".join(attribute for attribute, setting in custom_perms.perms.items() if setting)
await self.bot.say(
"You are required to have `{}` permissions to run `{}`".format(perms, cmd.qualified_name))
@ -165,34 +174,18 @@ class Mod:
if permissions.lower() == "none":
permissions = "send_messages"
# Check if the permission that was requested is valid
if getattr(discord.Permissions, permissions, None) is None:
# Convert the string to an int value of the permissions object, based on the required permission
# If we hit an attribute error, that means the permission given was not correct
perm_obj = discord.Permissions.none()
try:
setattr(perm_obj, permissions, True)
except AttributeError:
await self.bot.say("{} does not appear to be a valid permission! Valid permissions are: ```\n{}```"
.format(permissions, "\n".join(valid_perms)))
return
# Convert the string to an int value of the permissions object, based on the required permission
perm_obj = discord.Permissions.none()
setattr(perm_obj, permissions, True)
perm_value = perm_obj.value
# This next loop ensures the command given is valid. We need to loop through commands
# As self.bot.commands only includes parent commands
# So we are splitting the command in parts, looping through the commands
# And getting the subcommand based on the next part
# If we try to access commands of a command that isn't a group
# We'll hit an AttributeError, meaning an invalid command was given
# If we loop through and don't find anything, cmd will still be None
# And we'll report an invalid was given as well
cmd = None
for part in msg[0:len(msg) - 1]:
try:
if cmd is None:
cmd = self.bot.commands.get(part)
else:
cmd = cmd.commands.get(part)
except AttributeError:
cmd = None
break
cmd = self.find_command(command)
if cmd is None:
await self.bot.say(
@ -208,13 +201,19 @@ class Mod:
await self.bot.say("This command cannot have custom permissions setup!")
return
custom_perms = await config.get_content('custom_permissions')
server_perms = custom_perms.get(ctx.message.server.id) or {}
# Save the qualified name, so that we don't get screwed up by aliases
server_perms[cmd.qualified_name] = perm_value
custom_perms[ctx.message.server.id] = server_perms
r_filter = {'server_id': ctx.message.server.id}
entry = {'server_id': ctx.message.server.id,
cmd.qualified_name: perm_value}
await config.save_content('custom_permissions', custom_perms)
# In all other cases, I've used add_content before update_content
# In this case, I'm going the other way around, to make the least queries
# As custom permissions are probably going to be ran multiple times per server
# Whereas in most other cases, the command is probably going to be ran once/few times per server
if not await config.update_content('custom_permissions', entry, r_filter):
await config.add_content('custom_permissions', entry, r_filter)
# Same case as prefixes, for now, trigger a manual update
self.bot.loop.create_task(config.cache['custom_permissions'].update())
await self.bot.say("I have just added your custom permissions; "
"you now need to have `{}` permissions to use the command `{}`".format(permissions, command))
@ -222,48 +221,42 @@ class Mod:
@commands.has_permissions(manage_server=True)
async def remove_perms(self, ctx, *command: str):
"""Removes the custom permissions setup on the command specified"""
custom_perms = await config.get_content('custom_permissions')
server_perms = custom_perms.get(ctx.message.server.id) or {}
if server_perms is None:
await self.bot.say("There are no custom permissions setup on this server yet!")
return
cmd = None
# This is the same loop as the add command, we need this to get the
# command object so we can get the qualified_name
for part in command:
try:
if cmd is None:
cmd = self.bot.commands.get(part)
else:
cmd = cmd.commands.get(part)
except AttributeError:
cmd = None
break
cmd = self.find_command(command)
if cmd is None:
await self.bot.say(
"That command does not exist! You can't have custom permissions on a non-existant command....")
return
command_perms = server_perms.get(cmd.qualified_name)
if command_perms is None:
await self.bot.say("You do not have custom permissions setup on this command!")
return
del custom_perms[ctx.message.server.id][cmd.qualified_name]
await config.save_content('custom_permissions', custom_perms)
r_filter = {'server_id': ctx.message.server.id}
await config.replace_content('custom_permissions', r.row.without(cmd.qualified_name), r_filter)
await self.bot.say("I have just removed the custom permissions for {}!".format(cmd))
# Same case as prefixes, for now, trigger a manual update
self.bot.loop.create_task(config.cache['custom_permissions'].update())
@commands.command(pass_context=True, no_pm=True)
@checks.custom_perms(manage_server=True)
async def prefix(self, ctx, *, prefix: str):
"""This command can be used to set a custom prefix per server"""
prefixes = await config.get_content('prefixes')
prefixes[ctx.message.server.id] = prefix
await config.save_content('prefixes', prefixes)
r_filter = {'server_id': ctx.message.server.id}
if prefix.lower == "none":
prefix = None
entry = {'server_id': ctx.message.server.id,
'prefix': prefix}
if not await config.add_content('prefixes', entry, r_filter):
await config.update_content('prefixes', entry, r_filter)
# For now, cache is not fully implemented, however is needed for prefixes
# So we're going to manually trigger an update when this is ran
self.bot.loop.create_task(config.cache['prefixes'].update())
await self.bot.say(
"I have just updated the prefix for this server; you now need to call commands with `{}`".format(prefix))
"I have just updated the prefix for this server; you now need to call commands with `{0}`."
"For example, you can call this command again with {0}prefix".format(
prefix))
@commands.command(pass_context=True, no_pm=True)
@checks.custom_perms(manage_messages=True)
@ -315,71 +308,47 @@ class Mod:
@commands.group(aliases=['rule'], pass_context=True, no_pm=True, invoke_without_command=True)
@checks.custom_perms(send_messages=True)
async def rules(self, ctx):
async def rules(self, ctx, rule: int = None):
"""This command can be used to view the current rules on the server"""
rules = await config.get_content('rules')
server_rules = rules.get(ctx.message.server.id)
if server_rules is None or len(server_rules) == 0:
r_filter = {'server_id': ctx.message.server.id}
rules = await config.get_content('rules', r_filter)
rules = rules[0]['rules']
if rules is None:
await self.bot.say("This server currently has no rules on it! I see you like to live dangerously...")
return
# Enumerate the list, so that we can print the number and the rule for each rule
fmt = "\n".join("{}) {}".format(num + 1, rule) for num, rule in enumerate(server_rules))
await self.bot.say('```\n{}```'.format(fmt))
if rule is None:
# Enumerate the list, so that we can print the number and the rule for each rule
fmt = "\n".join("{}) {}".format(num + 1, rule) for num, rule in enumerate(rules))
await self.bot.say('```\n{}```'.format(fmt))
else:
fmt = rules[rule - 1]
await self.bot.say("Rule {}: \"{}\"".format(rule, fmt))
@rules.command(name='add', aliases=['create'], pass_context=True, no_pm=True)
@checks.custom_perms(manage_server=True)
async def rules_add(self, ctx, *, rule: str):
"""Adds a rule to this server's rules"""
# Nothing fancy here, just get the rules, append the rule, and save it
rules = await config.get_content('rules')
server_rules = rules.get(ctx.message.server.id) or []
server_rules.append(rule)
rules[ctx.message.server.id] = server_rules
await config.save_content('rules', rules)
r_filter = {'server_id': ctx.message.server.id}
entry = {'server_id': ctx.message.server.id,
'rules': [rule]}
update = lambda row: row['rules'].append(rule)
if not await config.update_content('rules', update, r_filter):
await config.add_content('rules', entry, r_filter)
await self.bot.say("I have just saved your new rule, use the rules command to view this server's current rules")
@rules.command(name='remove', aliases=['delete'], pass_context=True, no_pm=True)
@checks.custom_perms(manage_server=True)
async def rules_delete(self, ctx, rule: int = None):
async def rules_delete(self, ctx, rule: int):
"""Removes one of the rules from the list of this server's rules
Provide a number to delete that rule; if no number is provided
I'll print your current rules and ask for a number"""
rules = await config.get_content('rules')
server_rules = rules.get(ctx.message.server.id) or []
if server_rules is None or len(server_rules) == 0:
await self.bot.say(
"This server currently has no rules on it! Can't remove something that doesn't exist bro")
return
# Get the list of rules so that we can print it if no number was provided
# Since this is a list and not a dictionary, order is preserved, and we just need the number of the rule
list_rules = "\n".join("{}) {}".format(num + 1, rule) for num, rule in enumerate(server_rules))
if rule is None:
await self.bot.say("Your rules are:\n```\n{}```Please provide the rule number"
"you would like to remove (just the number)".format(list_rules))
# All we need for the check is to ensure that the content is just a digit, that is all we need
msg = await self.bot.wait_for_message(timeout=60.0, author=ctx.message.author, channel=ctx.message.channel,
check=lambda m: m.content.isdigit())
if msg is None:
await self.bot.say("You took too long...it's just a number, seriously? Try typing a bit quicker")
return
del server_rules[int(msg.content) - 1]
rules[ctx.message.server.id] = server_rules
await config.save_content('rules', rules)
Provide a number to delete that rule"""
r_filter = {'server_id': ctx.message.server.id}
update = {'rules': r.row['rules'].delete_at(rule - 1)}
if not await config.update_content('rules', update, r_filter):
await self.bot.say("That is not a valid rule number, try running the command again.")
else:
await self.bot.say("I have just removed that rule from your list of rules!")
return
# This check is just to ensure a number was provided within the list's range
try:
del server_rules[rule - 1]
rules[ctx.message.server.id] = server_rules
await config.save_content('rules', rules)
await self.bot.say("I have just removed that rule from your list of rules!")
except IndexError:
await self.bot.say("That is not a valid rule number, try running the command again. "
"Your current rules are:\n```\n{}```".format(list_rules))
def setup(bot):

View file

@ -38,18 +38,18 @@ class Overwatch:
async def ow_stats(self, ctx, user: discord.Member = None, hero: str = ""):
"""Prints out a basic overview of a member's stats
Provide a hero after the member to get stats for that specific hero"""
if user is None:
user = ctx.message.author
user = user or ctx.message.author
r_filter = {'member_id': user.id}
ow_stats = await config.get_content('overwatch', r_filter)
ow_stats = await config.get_content('overwatch')
bt = ow_stats.get(user.id)
if bt is None:
if ow_stats is None:
await self.bot.say("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
await self.bot.say("Searching profile information....")
bt = ow_stats[0]['battletag']
if hero == "":
# If no hero was provided, we just want the base stats for a player
async with self.session.get(base_url + "{}/stats/general".format(bt), headers=self.headers) as r:
@ -101,6 +101,7 @@ class Overwatch:
# Battletags are normally provided like name#id
# However the API needs this to be a -, so repliace # with - if it exists
bt = bt.replace("#", "-")
r_filter = {'member_id', ctx.message.author.id}
# This API sometimes takes a while to look up information, so send a message saying we're processing
await self.bot.say("Looking up your profile information....")
@ -115,25 +116,22 @@ class Overwatch:
return
# Now just save the battletag
ow = await config.get_content('overwatch')
ow[ctx.message.author.id] = bt
await config.save_content('overwatch', ow)
entry = {'member_id': ctx.message.author.id, 'battletag': bt}
# Try adding this first, if that fails, update the saved entry
if not await config.add_content('overwatch', entry, r_filter):
await config.update_content('overwatch', entry, r_filter)
await self.bot.say("I have just saved your battletag {}".format(ctx.message.author.mention))
@ow.command(pass_context=True, name="delete", aliases=['remove'], no_pm=True)
@checks.custom_perms(send_messages=True)
async def delete(self, ctx):
"""Removes your battletag from the records"""
result = await config.get_content('overwatch')
if result.get(ctx.message.author.id):
del result[ctx.message.author.id]
r_filter = {'member_id': ctx.message.author.id}
if await config.remove_content('overwatch', r_filter):
await self.bot.say("I no longer have your battletag saved {}".format(ctx.message.author.mention))
else:
await self.bot.say("I don't even have your battletag saved {}".format(ctx.message.author.mention))
del result[ctx.message.author.id]
await self.bot.say("I have just removed your battletag!")
def setup(bot):
bot.add_cog(Overwatch(bot))

View file

@ -1,10 +1,7 @@
from discord.ext import commands
from .utils import config
from .utils import checks
import re
import os
import glob
import sys
import discord
import inspect
import aiohttp
@ -45,35 +42,6 @@ class Owner:
fmt = "Nothing currently running!"
await self.bot.say(fmt)
@commands.command(pass_context=True)
@commands.check(checks.is_owner)
async def saferestart(self, ctx):
"""This commands is used to check if there is anything playing in any servers at the moment
If there is, I'll tell you not to restart, if not I'll just go ahead and restart"""
# I do not want to restart the bot if someone is playing music
# This gets all the exiting VoiceStates that are playing music right now
# If we are, say which server it
servers_playing_music = [server_id for server_id, state in self.bot.get_cog('Music').voice_states.items() if
state.is_playing()]
if len(servers_playing_music) > 0:
await self.bot.say("Sorry, it's not safe to restart. I am currently playing a song on {} servers".format(
len(servers_playing_music)))
else:
await config.save_content('restart_server', ctx.message.channel.id)
await self.bot.say("Restarting; see you in the next life {0}!".format(ctx.message.author.mention))
python = sys.executable
os.execl(python, python, *sys.argv)
@commands.command(pass_context=True)
@commands.check(checks.is_owner)
async def restart(self, ctx):
"""Forces the bot to restart"""
# This command is left in so that we can invoke it from saferestart, or we need a restart no matter what
await config.save_content('restart_server', ctx.message.channel.id)
await self.bot.say("Restarting; see you in the next life {0}!".format(ctx.message.author.mention))
python = sys.executable
os.execl(python, python, *sys.argv)
@commands.command()
@commands.check(checks.is_owner)
async def adddoggo(self, url: str):

View file

@ -2,6 +2,7 @@ import aiohttp
import asyncio
import discord
import re
import rethinkdb as r
from discord.ext import commands
from .utils import config
@ -21,8 +22,8 @@ async def online_users():
# This method is in place to just return all online_users
url = '{}/online/all?key={}'.format(base_url, key)
with aiohttp.ClientSession(headers={"User-Agent": "Bonfire/1.0.0"}) as s:
async with s.get(url) as r:
return await r.json()
async with s.get(url) as response:
return await response.json()
except:
return {}
@ -46,56 +47,59 @@ class Picarto:
await self.bot.wait_until_ready()
# This is a loop that runs every 30 seconds, checking if anyone has gone online
while not self.bot.is_closed:
picarto = await config.get_content('picarto')
r_filter = {'notifications_on': 1}
picarto = await config.get_content('picarto', r_filter)
# Get all online users before looping, so that only one request is needed
online_users_list = await online_users()
old_online_users = {m_id: data for m_id, data in picarto.items() if
data['notifications_on'] and data['live']}
old_offline_users = {m_id: data for m_id, data in picarto.items() if
data['notifications_on'] and not data['live']}
old_online_users = {data['member_id']: data for data in picarto if data['live']}
old_offline_users = {data['member_id']: data for data in picarto if not data['live']}
for m_id, r in old_offline_users.items():
for m_id, result in old_offline_users.items():
# Get their url and their user based on that url
url = r['picarto_url']
url = result['picarto_url']
user = re.search("(?<=picarto.tv/)(.*)", url).group(1)
# Check if they are online right now
if check_online(online_users_list, user):
for server_id in r['servers']:
for server_id in result['servers']:
# Get the channel to send the message to, based on the saved alert's channel
server = self.bot.get_server(server_id)
if server is None:
continue
server_alerts = await config.get_content('server_alerts')
channel_id = server_alerts.get(server_id) or server_id
server_alerts = await config.get_content('server_alerts', {'server_id': server_id})
try:
channel_id = server_alerts[0]
except IndexError:
channel_id = server_id
channel = self.bot.get_channel(channel_id)
# Get the member that has just gone live
member = discord.utils.get(server.members, id=m_id)
fmt = "{} has just gone live! View their stream at {}".format(member.display_name, url)
await self.bot.send_message(channel, fmt)
picarto[m_id]['live'] = 1
await config.save_content('picarto', picarto)
for m_id, r in old_online_users.items():
await config.update_content('picarto', {'live': 1}, {'member_id': m_id})
for m_id, result in old_online_users.items():
# Get their url and their user based on that url
url = r['picarto_url']
url = result['picarto_url']
user = re.search("(?<=picarto.tv/)(.*)", url).group(1)
# Check if they are online right now
if not check_online(online_users_list, user):
for server_id in r['servers']:
for server_id in result['servers']:
# Get the channel to send the message to, based on the saved alert's channel
server = self.bot.get_server(server_id)
if server is None:
continue
server_alerts = await config.get_content('server_alerts')
channel_id = server_alerts.get(server_id) or server_id
server_alerts = await config.get_content('server_alerts', {'server_id': server_id})
try:
channel_id = server_alerts[0]
except IndexError:
channel_id = server_id
channel = self.bot.get_channel(channel_id)
# Get the member that has just gone live
member = discord.utils.get(server.members, id=m_id)
fmt = "{} has just gone offline! Catch them next time they stream at {}".format(
member.display_name, url)
await self.bot.send_message(channel, fmt)
picarto[m_id]['live'] = 0
await config.save_content('picarto', picarto)
await config.update_content('picarto', {'live': 0}, {'member_id': m_id})
await asyncio.sleep(30)
@commands.group(pass_context=True, invoke_without_command=True)
@ -104,31 +108,32 @@ class Picarto:
"""This command can be used to view Picarto stats about a certain member"""
# If member is not given, base information on the author
member = member or ctx.message.author
picarto_urls = await config.get_content('picarto')
try:
member_url = picarto_urls.get(member.id)['picarto_url']
except:
r_filter = {'member_id': member.id}
picarto_entry = await config.get_content('picarto', r_filter)
if picarto_entry is None:
await self.bot.say("That user does not have a picarto url setup!")
return
member_url = picarto_entry[0]['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 = '{}/channel/{}?key={}'.format(base_url, stream, key)
async with self.session.get(url, headers=self.headers) as r:
data = await r.json()
async with self.session.get(url, headers=self.headers) as response:
data = await response.json()
# Not everyone has all these settings, so use this as a way to print information if it does, otherwise ignore it
things_to_print = ['channel', 'commissions_enabled', 'is_nsfw', 'program', 'tablet', 'followers',
'content_type']
# Using title and replace to provide a nice way to print the data
fmt = "\n".join(
"{}: {}".format(i.title().replace("_", " "), r) for i, r in data.items() if i in things_to_print)
"{}: {}".format(i.title().replace("_", " "), result) for i, result in data.items() if i in things_to_print)
# 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')
if social_links:
fmt2 = "\n".join("\t{}: {}".format(i.title().replace("_", " "), r) for i, r in social_links.items())
fmt2 = "\n".join(
"\t{}: {}".format(i.title().replace("_", " "), result) for i, result in social_links.items())
fmt = "{}\nSocial Links:\n{}".format(fmt, fmt2)
await self.bot.say("Picarto stats for {}: ```\n{}```".format(member.display_name, fmt))
@ -154,37 +159,30 @@ class Picarto:
api_url = '{}/channel/{}?key={}'.format(base_url, re.search("https://www.picarto.tv/(.*)", url).group(1), key)
# Check if we can find a user with the provided information, if we can't just return
async with self.session.get(api_url, headers=self.headers) as r:
if not r.status == 200:
async with self.session.get(api_url, headers=self.headers) as response:
if not response.status == 200:
await self.bot.say("That Picarto user does not exist! "
"What would be the point of adding a nonexistant Picarto user? Silly")
return
picarto_urls = await config.get_content('picarto')
result = picarto_urls.get(ctx.message.author.id)
# If information for this user already exists, override just the url, and not the information
# Otherwise create the information with notications on, and that they're not live.
# The next time it's checked, they'll go 'online'
if result is not None:
picarto_urls[ctx.message.author.id]['picarto_url'] = url
r_filter = {'member_id': ctx.message.author.id}
entry = {'picarto_url': url,
'servers': [ctx.message.server.id],
'notifications_on': 1, 'live': 0}
if await config.add_content('picarto', entry, r_filter):
await self.bot.say(
"I have just saved your Picarto URL {}, this server will now be notified when you go live".format(
ctx.message.author.mention))
else:
picarto_urls[ctx.message.author.id] = {'picarto_url': url,
'servers': [ctx.message.server.id],
'notifications_on': 1, 'live': 0}
await config.save_content('picarto', picarto_urls)
await self.bot.say(
"I have just saved your Picarto url {}, this server will now be notified when you go live".format(
ctx.message.author.mention))
await config.update_content('picarto', {'picarto_url': url}, r_filter)
await self.bot.say("I have just updated your Picarto URL")
@picarto.command(name='remove', aliases=['delete'], pass_context=True, no_pm=True)
@checks.custom_perms(send_messages=True)
async def remove_picarto_url(self, ctx):
"""Removes your picarto URL"""
picarto = await config.get_content('picarto')
if picarto.get(ctx.message.author.id) is not None:
del picarto[ctx.message.author.id]
await config.save_content('picarto', picarto)
r_filter = {'member_id': ctx.message.author.id}
if await config.remove_content('picarto', r_filter):
await self.bot.say("I am no longer saving your picarto URL {}".format(ctx.message.author.mention))
else:
await self.bot.say(
@ -196,65 +194,39 @@ class Picarto:
async def notify(self, ctx):
"""This can be used to turn picarto notifications on or off
Call this command by itself, to add this server to the list of servers to be notified"""
member = ctx.message.author
# If this user's picarto URL is not saved, no use in adding this server to the list that doesn't exist
picarto = await config.get_content('picarto')
result = picarto.get(member.id)
r_filter = {'member_id': ctx.message.author.id}
result = await config.get_content('picarto', r_filter)
# Check if this user is saved at all
if result is None:
await self.bot.say(
"I do not have your picarto URL added {}. You can save your picarto url with {}picarto add".format(
member.mention, ctx.prefix))
# Append this server's ID and save the new content
picarto[member.id]['servers'].append(ctx.message.server.id)
await config.save_content('picarto', picarto)
await self.bot.say(
"I have just changed which channel will be notified when you go live, to `{}`".format(
ctx.message.channel.name))
"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 server is already added as one to notify in
elif ctx.message.server.id in result[0]['servers']:
await self.bot.say("I am already set to notify in this server...")
else:
await config.update_content('picarto', {'servers': r.row['servers'].append(ctx.message.server.id)},
r_filter)
@notify.command(name='on', aliases=['start,yes'], pass_context=True, no_pm=True)
@checks.custom_perms(send_messages=True)
async def notify_on(self, ctx):
"""Turns picarto notifications on"""
picarto = await config.get_content('picarto')
result = picarto.get(ctx.message.author.id)
# Check if this user has saved their picarto URL first
if result is None:
await self.bot.say(
"I do not have your picarto URL added {}. You can save your picarto url with !picarto add".format(
ctx.message.author.mention))
# Next check if they are already set to notify
elif result['notifications_on']:
await self.bot.say("What do you want me to do, send two notifications? Not gonna happen {}".format(
ctx.message.author.mention))
else:
picarto[ctx.message.author.id]['notifications_on'] = 1
await config.save_content('picarto', picarto)
await self.bot.say("I will notify if you go live {}, you'll get a bajillion followers I promise c:".format(
ctx.message.author.mention))
r_filter = {'member_id': ctx.message.author.id}
await config.update_content('picarto', {'notifications_on': 1}, r_filter)
await self.bot.say("I will notify if you go live {}, you'll get a bajillion followers I promise c:".format(
ctx.message.author.mention))
@notify.command(name='off', aliases=['stop,no'], pass_context=True, no_pm=True)
@checks.custom_perms(send_messages=True)
async def notify_off(self, ctx):
"""Turns picarto notifications off"""
picarto = await config.get_content('picarto')
# Check if this user has saved their picarto URL first
if picarto.get(ctx.message.author.id) is None:
await self.bot.say(
"I do not have your picarto URL added {}. You can save your picarto url with !picarto add".format(
ctx.message.author.mention))
# Next check if they are already set to not notify
elif not picarto.get(ctx.message.author.id)['notifications_on']:
await self.bot.say("I am already set to not notify if you go live! Pay attention brah {}".format(
r_filter = {'member_id': ctx.message.author.id}
await config.update_content('picarto', {'notifications_on': 0}, r_filter)
await self.bot.say(
"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(
ctx.message.author.mention))
else:
picarto[ctx.message.author.id]['notifications_on'] = 0
await config.save_content('picarto', picarto)
await self.bot.say(
"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(
ctx.message.author.mention))
def setup(bot):

View file

@ -3,8 +3,15 @@ from .utils import config
from .utils import checks
import discord
def find_command(command):
cmd = None
class Stats:
"""Leaderboard/stats related commands"""
def __init__(self, bot):
self.bot = bot
def find_command(self, command):
cmd = None
for part in command.split():
try:
@ -16,26 +23,18 @@ def find_command(command):
cmd = None
break
if cmd is None:
await self.bot.say("That is not a valid command!")
return
class Stats:
"""Leaderboard/stats related commands"""
def __init__(self, bot):
self.bot = bot
return cmd
@commands.group(no_pm=True)
@checks.custom_perms(send_messages=True)
async def command(self):
pass
@command.command(no_pm=True, name = "stats")
@command.command(no_pm=True, name="stats")
@checks.custom_perms(send_messages=True)
async def command_stats(self, ctx, *, command):
"""This command can be used to view some usage stats about a specific command"""
cmd = find_command(command)
cmd = self.find_command(command)
if cmd is None:
await self.bot.say("`{}` is not a valid command".format(command))
@ -45,11 +44,15 @@ class Stats:
@checks.custom_perms(send_messages=True)
async def mostboops(self, ctx):
"""Shows the person you have 'booped' the most, as well as how many times"""
boops = await config.get_content('boops')
if not boops.get(ctx.message.author.id):
r_filter = {'member_id': ctx.message.author.id}
boops = await config.get_content('boops', r_filter)
if boops is None:
await self.bot.say("You have not booped anyone {} Why the heck not...?".format(ctx.message.author.mention))
return
# Just to make this easier, just pay attention to the boops data, now that we have the right entry
boops = boops[0]['boops']
# First get a list of the ID's of all members in this server, for use in list comprehension
server_member_ids = [member.id for member in ctx.message.server.members]
# Then get a sorted list, based on the amount of times they've booped the member
@ -69,15 +72,18 @@ class Stats:
@checks.custom_perms(send_messages=True)
async def listboops(self, ctx):
"""Lists all the users you have booped and the amount of times"""
boops = await config.get_content('boops')
booped_members = boops.get(ctx.message.author.id)
if booped_members is None:
r_filter = {'member_id': ctx.message.author.id}
boops = await config.get_content('boops', r_filter)
if boops is None:
await self.bot.say("You have not booped anyone {} Why the heck not...?".format(ctx.message.author.mention))
return
# Just to make this easier, just pay attention to the boops data, now that we have the right entry
boops = boops[0]['boops']
# Same concept as the mostboops method
server_member_ids = [member.id for member in ctx.message.server.members]
booped_members = {m_id: amt for m_id, amt in booped_members.items() if m_id in server_member_ids}
booped_members = {m_id: amt for m_id, amt in boops.items() if m_id in server_member_ids}
sorted_booped_members = sorted(booped_members.items(), key=lambda k: k[1], reverse=True)
output = "\n".join(
@ -89,13 +95,13 @@ class Stats:
@checks.custom_perms(send_messages=True)
async def leaderboard(self, ctx):
"""Prints a leaderboard of everyone in the server's battling record"""
battles = await config.get_content('battle_records')
# Same concept as mostboops
# 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.server.members]
server_members = {member_id: stats for member_id, stats in battles.items() if member_id in server_member_ids}
# Only real difference is the key, the key needs to be based on the rating in the member's dictionary of stats
sorted_members = sorted(server_members.items(), key=lambda k: k[1]['rating'], reverse=True)
r_filter = lambda row: row['member_id'] in server_member_ids
battles = await config.get_content('battle_records', r_filter)
# Sort the members based on their rating
sorted_members = sorted(battles, key=lambda k: k['rating'], reverse=True)
fmt = ""
count = 1
@ -109,21 +115,25 @@ class Stats:
break
await self.bot.say("Battling leaderboard for this server:```\n{}```".format(fmt))
@commands.command(pass_context=True, no_pm=True)
@commands.command(pass_context=True, no_pm=True, name="battle stats")
@checks.custom_perms(send_messages=True)
async def stats(self, ctx, member: discord.Member=None):
async def stats(self, ctx, member: discord.Member = None):
"""Prints the battling stats for you, or the user provided"""
member = member or ctx.message.author
# 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 config.get_content('battle_records')
if member.id not in all_members:
# Make a list comprehension to just check if the user has battled
if len([entry for entry in all_members if entry['member_id'] == member.id]) == 0:
await self.bot.say("That user has not battled yet!")
return
# Same concept as the leaderboard
server_member_ids = [member.id for member in ctx.message.server.members]
server_members = {member_id: stats for member_id, stats in all_members.items() if
member_id in server_member_ids}
server_members = {stats['member_id']: stats for stats in all_members if
stats['member_id'] in server_member_ids}
sorted_server_members = sorted(server_members.items(), key=lambda x: x[1]['rating'], reverse=True)
sorted_all_members = sorted(all_members.items(), key=lambda x: x[1]['rating'], reverse=True)

View file

@ -19,10 +19,12 @@ class StatsUpdate:
def __unload(self):
self.bot.loop.create_task(self.session.close())
async def update(self, data):
async def update(self):
server_count = 0
for d in data.values():
server_count += d.get('server_count')
data = await config.get_content('bot_data')
for entry in data:
server_count += entry.get('server_count')
carbon_payload = {
'key': config.carbon_key,
@ -46,31 +48,34 @@ class StatsUpdate:
log.info('bots.discord.pw statistics returned {} for {}'.format(resp.status, payload))
async def on_server_join(self, server):
data = await config.get_content('bot_data')
shard_data = data.get('shard_{}'.format(config.shard_id)) or {}
shard_data['server_count'] = len(self.bot.servers)
shard_data['member_count'] = len(set(self.bot.get_all_members()))
data['shard_{}'.format(config.shard_id)] = shard_data
await config.save_content('bot_data', data)
await self.update(data)
r_filter = {'shard_id': config.shard_id}
server_count = len(self.bot.servers)
member_count = len(set(self.bot.get_all_members()))
entry = {'server_count': server_count, 'member_count': member_count, "shard_id": config.shard_id}
# Check if this was successful, if it wasn't, that means a new shard was added and we need to add that entry
if not await config.update_content('bot_data', entry, r_filter):
await config.add_content('bot_data', entry, r_filter)
self.bot.loop.create_task(self.update())
async def on_server_leave(self, server):
data = await config.get_content('bot_data')
shard_data = data.get('shard_{}'.format(config.shard_id)) or {}
shard_data['server_count'] = len(self.bot.servers)
shard_data['member_count'] = len(set(self.bot.get_all_members()))
data['shard_{}'.format(config.shard_id)] = shard_data
await config.save_content('bot_data', data)
await self.update(data)
r_filter = {'shard_id': config.shard_id}
server_count = len(self.bot.servers)
member_count = len(set(self.bot.get_all_members()))
entry = {'server_count': server_count, 'member_count': member_count, "shard_id": config.shard_id}
# Check if this was successful, if it wasn't, that means a new shard was added and we need to add that entry
if not await config.update_content('bot_data', entry, r_filter):
await config.add_content('bot_data', entry, r_filter)
self.bot.loop.create_task(self.update())
async def on_ready(self):
data = await config.get_content('bot_data')
shard_data = data.get('shard_{}'.format(config.shard_id)) or {}
shard_data['server_count'] = len(self.bot.servers)
shard_data['member_count'] = len(set(self.bot.get_all_members()))
data['shard_{}'.format(config.shard_id)] = shard_data
await config.save_content('bot_data', data)
await self.update(data)
r_filter = {'shard_id': config.shard_id}
server_count = len(self.bot.servers)
member_count = len(set(self.bot.get_all_members()))
entry = {'server_count': server_count, 'member_count': member_count, "shard_id": config.shard_id}
# Check if this was successful, if it wasn't, that means a new shard was added and we need to add that entry
if not await config.update_content('bot_data', entry, r_filter):
await config.add_content('bot_data', entry, r_filter)
self.bot.loop.create_task(self.update())
def setup(bot):

View file

@ -8,6 +8,7 @@ import aiohttp
import re
import json
import pendulum
import rethinkdb as r
def setup(bot):
@ -35,20 +36,25 @@ class Strawpoll:
async def strawpolls(self, ctx, poll_id: str = None):
"""This command can be used to show a strawpoll setup on this server"""
# Strawpolls cannot be 'deleted' so to handle whether a poll is running or not on a server
# Just save the poll in the config file, which can then be removed when it should not be "running" anymore
all_polls = await config.get_content('strawpolls')
server_polls = all_polls.get(ctx.message.server.id) or {}
if not server_polls:
await self.bot.say("There are currently no strawpolls running on this server!")
return
# If no poll_id was provided, print a list of all current running poll's on this server
if not poll_id:
# Just save the poll, which can then be removed when it should not be "running" anymore
r_filter = {'server_id': ctx.message.server.id}
# Add to our filter the poll_id, if one was provided
if poll_id is not None:
r_filter['poll_id'] = poll_id
polls = await config.get_content('strawpolls', r_filter)
# Now output all polls if a poll ID was not provided
if poll_id is None:
if len(polls) == 0:
await self.bot.say("There are currently no strawpolls running on this server!")
return
fmt = "\n".join(
"{}: https://strawpoll.me/{}".format(data['title'], _id) for _id, data in server_polls.items())
"{}: https://strawpoll.me/{}".format(data['title'], data['poll_id']) for data in polls)
await self.bot.say("```\n{}```".format(fmt))
# Else if a valid poll_id was provided, print info about that poll
elif poll_id in server_polls.keys():
poll = server_polls[poll_id]
# Otherwise, a poll ID was provided, and we searched based on the poll ID earlier, so just check the length
elif len(polls) > 0:
# Since strawpoll should never allow us to have more than one poll with the same ID
# It's safe to assume there's only one result
poll = polls[0]
async with self.session.get("{}/{}".format(self.url, poll_id),
headers={'User-Agent': 'Bonfire/1.0.0'}) as response:
@ -59,7 +65,7 @@ class Strawpoll:
# And the votes to match it, based on the index of the option
# The rest is simple formatting
fmt_options = "\n\t".join(
"{}: {}".format(r, data['votes'][i]) for i, r in enumerate(data['options']))
"{}: {}".format(result, data['votes'][i]) for i, result in enumerate(data['options']))
author = discord.utils.get(ctx.message.server.members, id=poll['author'])
created_ago = (pendulum.utcnow() - pendulum.parse(poll['date'])).in_words()
link = "https://strawpoll.me/{}".format(poll_id)
@ -107,37 +113,33 @@ class Strawpoll:
return
# Save this strawpoll in the list of running strawpolls for a server
all_polls = await config.get_content('strawpolls')
server_polls = all_polls.get(ctx.message.server.id) or {}
server_polls[str(data['id'])] = {'author': ctx.message.author.id, 'date': str(pendulum.utcnow()), 'title': title}
all_polls[ctx.message.server.id] = server_polls
await config.save_content('strawpolls', all_polls)
poll_id = str(data['id'])
await self.bot.say("Link for your new strawpoll: https://strawpoll.me/{}".format(data['id']))
r_filter = {'server_id': ctx.message.server.id}
sub_entry = {'poll_id': poll_id,
'author': ctx.message.author.id,
'date': str(pendulum.utcnow()),
'title': title}
entry = {'server_id': ctx.message.server.id,
'polls': sub_entry}
update = {'polls': r.row['polls'].append(sub_entry)}
if not await config.update_content('strawpolls', update, r_filter):
await config.add_content('strawpolls', entry, {'poll_id': poll_id})
await self.bot.say("Link for your new strawpoll: https://strawpoll.me/{}".format(poll_id))
@strawpolls.command(name='delete', aliases=['remove', 'stop'], pass_context=True)
@checks.custom_perms(kick_members=True)
async def remove_strawpoll(self, ctx, poll_id: str = None):
"""This command can be used to delete one of the existing strawpolls
If you don't provide an ID it will print the list of polls available"""
async def remove_strawpoll(self, ctx, poll_id):
"""This command can be used to delete one of the existing strawpolls"""
r_filter = {'server_id': ctx.message.server.id}
content = await config.get_content('strawpolls', r_filter)
content = content[0]['polls']
polls = [{key: result} for key, result in content.items() if poll_id != result['poll_id']]
all_polls = await config.get_content('strawpolls')
server_polls = all_polls.get(ctx.message.server.id) or {}
# Check if a poll_id was provided, if it is then we can continue, if not print the list of current polls
if poll_id:
poll = server_polls.get(poll_id)
# Check if no poll exists with that ID, then print a list of the polls
if not poll:
fmt = "\n".join("{}: {}".format(data['title'], _poll_id) for _poll_id, data in server_polls.items())
await self.bot.say(
"There is no poll setup with that ID! Here is a list of the current polls```\n{}```".format(fmt))
else:
# Delete the poll that was just found
del server_polls[poll_id]
all_polls[ctx.message.server.id] = server_polls
await config.save_content('strawpolls', all_polls)
await self.bot.say("I have just removed the poll with the ID {}".format(poll_id))
update = {'polls': polls}
# Try to remove the poll based on the ID, if it doesn't exist, this will return false
if await config.update_content('strawpolls', update, r_filter):
await self.bot.say("I have just removed the poll with the ID {}".format(poll_id))
else:
fmt = "\n".join("{}: {}".format(data['title'], _poll_id) for _poll_id, data in server_polls.items())
await self.bot.say("Here is a list of the polls on this server:\n```\n{}```".format(fmt))
await self.bot.say("There is no poll setup with that ID!")

View file

@ -14,10 +14,9 @@ class Tags:
@checks.custom_perms(send_messages=True)
async def tags(self, ctx):
"""Prints all the custom tags that this server currently has"""
tags = await config.get_content('tags')
tags = tags['tags']
tags = await config.get_content('tags', {'server_id': ctx.message.server.id})
# Simple generator that adds a tag to the list to print, if the tag is for this server
fmt = "\n".join("{}".format(tag['tag']) for tag in tags if tag['server_id'] == ctx.message.server.id)
fmt = "\n".join("{}".format(tag['tag']) for tag in tags)
await self.bot.say('```\n{}```'.format(fmt))
@commands.group(pass_context=True, invoke_without_command=True, no_pm=True)
@ -25,15 +24,13 @@ class Tags:
async def tag(self, ctx, *, tag: str):
"""This can be used to call custom tags
The format to call a custom tag is !tag <tag>"""
tags = await config.get_content('tags')
tags = tags['tags']
# Same generator as the method for tags, other than the second check to get the tag that is provided
result = [t for t in tags if t['tag'] == tag and t['server_id'] == ctx.message.server.id]
if len(result) == 0:
r_filter = lambda row: (row['server_id'] == ctx.message.server.id) & (row['tag'] == tag)
tags = await config.get_content('tags', r_filter)
if len(tags) == 0:
await self.bot.say('That tag does not exist!')
return
# We shouldn't ever have two tags of the same name, so just get the first result
await self.bot.say("\u200B{}".format(result[0]['result']))
await self.bot.say("\u200B{}".format(tags[0]['result']))
@tag.command(name='add', aliases=['create', 'start'], pass_context=True, no_pm=True)
@checks.custom_perms(kick_members=True)
@ -55,39 +52,28 @@ class Tags:
"Please provide the format for the tag in: {}tag add <tag> - <result>".format(ctx.prefix))
return
tags = await config.get_content('tags')
tags = tags['tags']
for t in tags:
# Attempt to find a tag with that name, so that we update it instead of making a duplicate
if t['tag'] == tag and t['server_id'] == ctx.message.server.id:
t['result'] = tag_result
await self.bot.say(
"I have just updated the tag `{0}`! You can call this tag by entering !tag {0}".format(tag))
# If we haven't found one, append a new one to the list
tags.append({'server_id': ctx.message.server.id, 'tag': tag, 'result': tag_result})
await self.bot.say(
"I have just added the tag `{0}`! You can call this tag by entering !tag {0}".format(tag))
await config.save_content('tags', {'tags': tags})
entry = {'server_id': ctx.message.server.id, 'tag': tag, 'result': tag_result}
r_filter = lambda row: (row['server_id'] == ctx.message.server.id) & (row['tag'] == tag)
# Try to create new entry first, if that fails (it already exists) then we update it
if await config.add_content('tags', entry, r_filter):
await self.bot.say(
"I have just added the tag `{0}`! You can call this tag by entering !tag {0}".format(tag))
else:
await config.update_content('tags', entry, r_filter)
await self.bot.say(
"I have just updated the tag `{0}`! You can call this tag by entering !tag {0}".format(tag))
@tag.command(name='delete', aliases=['remove', 'stop'], pass_context=True, no_pm=True)
@checks.custom_perms(kick_members=True)
async def del_tag(self, ctx, *, tag: str):
"""Use this to remove a tag that from use for this server
"""Use this to remove a tag from use for this server
Format to delete a tag is !tag delete <tag>"""
tags = await config.get_content('tags')
tags = tags['tags']
# Get a list of the tags that match this server, and the name provided (should only ever be one if any)
result = [t for t in tags if t['tag'] == tag and t['server_id'] == ctx.message.server.id]
# If we haven't found one, can't delete it
if len(result) == 0:
r_filter = lambda row: (row['server_id'] == ctx.message.server.id) & (row['tag'] == tag)
if await config.remove_content('tags', r_filter):
await self.bot.say('I have just removed the tag `{}`'.format(tag))
else:
await self.bot.say(
"The tag {} does not exist! You can't remove something if it doesn't exist...".format(tag))
return
# Since there should never be more than one result due to our checks we've made, just remove the first result
tags.remove(result[0])
await self.bot.say('I have just removed the tag `{}`'.format(tag))
await config.save_content('tags', {'tags': tags})
def setup(bot):

View file

@ -98,48 +98,6 @@ class Board:
return "```\n{}```".format(_board)
async def update_records(winner, loser):
# This is the exact same formula as the battling update.
# The only difference is I use the word "match" instead of "battle"
matches = await config.get_content('tictactoe')
winner_stats = matches.get(winner.id) or {}
winner_rating = winner_stats.get('rating') or 1000
loser_stats = matches.get(loser.id) or {}
loser_rating = loser_stats.get('rating') or 1000
difference = abs(winner_rating - loser_rating)
rating_change = 0
count = 25
while count <= difference:
if count > 300:
break
rating_change += 1
count += 25
if winner_rating > loser_rating:
winner_rating += 16 - rating_change
loser_rating -= 16 - rating_change
else:
winner_rating += 16 + rating_change
loser_rating -= 16 + rating_change
winner_wins = winner_stats.get('wins') or 0
winner_losses = winner_stats.get('losses') or 0
loser_wins = loser_stats.get('wins') or 0
loser_losses = loser_stats.get('losses') or 0
winner_wins += 1
loser_losses += 1
winner_stats = {'wins': winner_wins, 'losses': winner_losses, 'rating': winner_rating}
loser_stats = {'wins': loser_wins, 'losses': loser_losses, 'rating': loser_rating}
matches[winner.id] = winner_stats
matches[loser.id] = loser_stats
await config.save_content('tictactoe', matches)
class TicTacToe:
def __init__(self, bot):
self.bot = bot
@ -230,7 +188,7 @@ class TicTacToe:
await self.bot.say("{} 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 update_records(winner, loser)
await config.update_records('tictactoe', winner, loser)
# This game has ended, delete it so another one can be made
del self.boards[ctx.message.server.id]
else:

View file

@ -1,11 +1,13 @@
from discord.ext import commands
from .utils import config
from .utils import checks
import aiohttp
import asyncio
import discord
import json
import re
import rethinkdb as r
class Twitch:
@ -23,15 +25,15 @@ class Twitch:
# Check a specific channel's data, and get the response in text format
url = "https://api.twitch.tv/kraken/streams/{}?client_id={}".format(channel, self.key)
with aiohttp.ClientSession() as s:
async with s.get(url) as r:
response = await r.text()
async with s.get(url) as response:
result = await response.text()
# For some reason Twitch's API call is not reliable, sometimes it returns stream as None
# That is what we're checking specifically, sometimes it doesn't exist in the returned JSON at all
# Sometimes it returns something that cannot be decoded with JSON
# In either error case, just assume they're offline, the next check will most likely work
try:
data = json.loads(response)
data = json.loads(result)
return data['stream'] is not None
except (KeyError, json.JSONDecodeError):
return False
@ -40,54 +42,55 @@ class Twitch:
await self.bot.wait_until_ready()
# Loop through as long as the bot is connected
while not self.bot.is_closed:
twitch = await config.get_content('twitch')
twitch = await config.get_content('twitch', {'notifications_on': 1})
# Online/offline is based on whether they are set to such, in the config file
# This means they were detected as online/offline before and we check for a change
online_users = {m_id: data for m_id, data in twitch.items() if data['notifications_on'] and data['live']}
offline_users = {m_id: data for m_id, data in twitch.items() if
data['notifications_on'] and not data['live']}
for m_id, r in offline_users.items():
online_users = {data['member_id']: data for data in twitch if data['live']}
offline_users = {data['member_id']: data for data in twitch if not data['live']}
for m_id, result in offline_users.items():
# Get their url and their user based on that url
url = r['twitch_url']
url = result['twitch_url']
user = re.search("(?<=twitch.tv/)(.*)", url).group(1)
# Check if they are online right now
if await self.channel_online(user):
for server_id in r['servers']:
for server_id in result['servers']:
# Get the channel to send the message to, based on the saved alert's channel
server = self.bot.get_server(server_id)
if server is None:
continue
server_alerts = await config.get_content('server_alerts')
channel_id = server_alerts.get(server_id) or server_id
server_alerts = await config.get_content('server_alerts', {'server_id': server_id})
channel_id = server_id
if len(server_alerts) > 0:
channel_id = server_alerts[0].get('channel_id')
channel = self.bot.get_channel(channel_id)
# Get the member that has just gone live
member = discord.utils.get(server.members, id=m_id)
fmt = "{} has just gone live! View their stream at {}".format(member.display_name, url)
await self.bot.send_message(channel, fmt)
twitch[m_id]['live'] = 1
await config.save_content('twitch', twitch)
for m_id, r in online_users.items():
await config.update_content('twitch', {'live': 1}, {'member_id': m_id})
for m_id, result in online_users.items():
# Get their url and their user based on that url
url = r['twitch_url']
url = result['twitch_url']
user = re.search("(?<=twitch.tv/)(.*)", url).group(1)
# Check if they are online right now
if not await self.channel_online(user):
for server_id in r['servers']:
for server_id in result['servers']:
# Get the channel to send the message to, based on the saved alert's channel
server = self.bot.get_server(server_id)
if server is None:
continue
server_alerts = await config.get_content('server_alerts')
channel_id = server_alerts.get(server_id) or server_id
server_alerts = await config.get_content('server_alerts', {'server_id': server_id})
channel_id = server_id
if len(server_alerts) > 0:
channel_id = server_alerts[0].get('channel_id')
channel = self.bot.get_channel(channel_id)
# Get the member that has just gone live
member = discord.utils.get(server.members, id=m_id)
fmt = "{} has just gone offline! Catch them next time they stream at {}".format(
member.display_name, url)
await self.bot.send_message(channel, fmt)
twitch[m_id]['live'] = 0
await config.save_content('twitch', twitch)
await config.update_content('twitch', {'live': 0}, {'member_id': m_id})
await asyncio.sleep(30)
@commands.group(no_pm=True, invoke_without_command=True, pass_context=True)
@ -97,8 +100,7 @@ class Twitch:
if member is None:
member = ctx.message.author
twitch_channels = await config.get_content('twitch')
result = twitch_channels.get(member.id)
result = await config.get_content('twitch', {'member_id': member.id})
if result is None:
await self.bot.say("{} has not saved their twitch URL yet!".format(member.name))
return
@ -107,8 +109,8 @@ class Twitch:
user = re.search("(?<=twitch.tv/)(.*)", url).group(1)
twitch_url = "https://api.twitch.tv/kraken/channels/{}?client_id={}".format(user, self.key)
with aiohttp.ClientSession() as s:
async with s.get(twitch_url) as r:
data = await r.json()
async with s.get(twitch_url) as response:
data = await response.json()
with open("twitch_testing", 'w') as f:
json.dump(data, f)
@ -139,100 +141,75 @@ class Twitch:
# Try to find the channel provided, we'll get a 404 response if it does not exist
with aiohttp.ClientSession() as s:
async with s.get(url) as r:
if not r.status == 200:
async with s.get(url) as response:
if not response.status == 200:
await self.bot.say("That twitch user does not exist! "
"What would be the point of adding a nonexistant twitch user? Silly")
return
twitch = await config.get_content('twitch')
result = twitch.get(ctx.message.author.id)
r_filter = {'member_id': ctx.message.author.id}
result = await config.get_content('twitch', r_filter)
# 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 result is not None:
twitch[ctx.message.author.id]['twitch_url'] = url
await config.update_content('twitch', {'twitch_url': url}, r_filter)
else:
twitch[ctx.message.author.id] = {'twitch_url': url, 'servers': [ctx.message.server.id],
'notifications_on': 1, 'live': 0}
await config.save_content('twitch', twitch)
await config.add_content('twitch', {'twitch_url': url, 'servers': [ctx.message.server.id],
'notifications_on': 1, 'live': 0}, r_filter)
await self.bot.say("I have just saved your twitch url {}".format(ctx.message.author.mention))
@twitch.command(name='remove', aliases=['delete'], pass_context=True, no_pm=True)
@checks.custom_perms(send_messages=True)
async def remove_twitch_url(self, ctx):
"""Removes your twitch URL"""
twitch = await config.get_content('twitch')
# Make sure the user exists before trying to delete them from the list
if twitch.get(ctx.message.author.id) is not None:
# Simply remove this user from the list, and save
del twitch[ctx.message.author.id]
await config.save_content('twitch', twitch)
await self.bot.say("I am no longer saving your twitch URL {}".format(ctx.message.author.mention))
else:
await self.bot.say(
"I do not have your twitch URL added {}. You can save your twitch url with !twitch add".format(
ctx.message.author.mention))
# Just try to remove it, if it doesn't exist, nothing is going to happen
await config.remove_content('twitch', {'member_id': ctx.message.author.id})
await self.bot.say("I am no longer saving your twitch URL {}".format(ctx.message.author.mention))
@twitch.group(pass_context=True, no_pm=True, invoke_without_command=True)
@checks.custom_perms(send_messages=True)
async def notify(self, ctx):
"""This can be used to modify notification settings for your twitch user
Call this command by itself to add 'this' server as one that will be notified when you on/offline"""
twitch = await config.get_content('twitch')
result = twitch.get(ctx.message.author.id)
r_filter = {'member_id': ctx.message.author.id}
result = await config.get_content('twitch', r_filter)
# Check if this user is saved at all
if result is None:
await self.bot.say(
"I do not have your twitch URL added {}. You can save your twitch url with !twitch add".format(
ctx.message.author.mention))
# Otherwise we just need to append the server's ID to the servers list
# Then check if this server is already added as one to notify in
elif ctx.message.server.id in result[0]['servers']:
await self.bot.say("I am already set to notify in this server...")
else:
twitch[ctx.message.author.id]['servers'].append(ctx.message.server.id)
await config.save_content('twitch', twitch)
await config.update_content('twitch', {'servers': r.row['servers'].append(ctx.message.server.id)}, r_filter)
@notify.command(name='on', aliases=['start,yes'], pass_context=True, no_pm=True)
@checks.custom_perms(send_messages=True)
async def notify_on(self, ctx):
"""Turns twitch notifications on"""
# Make sure this user is saved before we attempt to modify their information
twitch = await config.get_content('twitch')
result = twitch.get(ctx.message.author.id)
if result is None:
await self.bot.say(
"I do not have your twitch URL added {}. You can save your twitch url with !twitch add".format(
ctx.message.author.mention))
# Then check to see if notifications are already on
elif result['notifications_on']:
await self.bot.say("What do you want me to do, send two notifications? Not gonna happen {}".format(
ctx.message.author.mention))
# Otherwise, turn on notifications
else:
twitch[ctx.message.author.id]['notifications_on'] = 1
await config.save_content('twitch', twitch)
r_filter = {'member_id': ctx.message.author.id}
if await config.update_content('twitch', {"notifications_on": 1}, r_filter):
await self.bot.say("I will notify if you go live {}, you'll get a bajillion followers I promise c:".format(
ctx.message.author.mention))
else:
await self.bot.say("I can't notify if you go live if I don't know your twitch URL yet!")
@notify.command(name='off', aliases=['stop,no'], pass_context=True, no_pm=True)
@checks.custom_perms(send_messages=True)
async def notify_off(self, ctx):
"""Turns twitch notifications off"""
# This method is exactly the same, except for turning off notifcations instead of on
twitch = await config.get_content('twitch')
if twitch.get(ctx.message.author.id) is None:
await self.bot.say(
"I do not have your twitch URL added {}. You can save your twitch url with !twitch add".format(
ctx.message.author.mention))
elif not twitch.get(ctx.message.author.id)['notifications_on']:
await self.bot.say("I am already set to not notify if you go live! Pay attention brah {}".format(
ctx.message.author.mention))
else:
twitch[ctx.message.author.id]['notifications_on'] = 0
r_filter = {'member_id': ctx.message.author.id}
if await config.update_content('twitch', {"notifications_on": 1}, r_filter):
await self.bot.say(
"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(
ctx.message.author.mention))
else:
await self.bot.say(
"I mean, I'm already not going to notify anyone, because I don't have your twitch URL saved...")
def setup(bot):

View file

@ -4,6 +4,7 @@ import rethinkdb as r
import pendulum
loop = asyncio.get_event_loop()
global_config = {}
# Ensure that the required config.yml file actually exists
try:
@ -38,7 +39,7 @@ class Cache:
loop.create_task(self.update())
async def update(self):
self.values = await _get_content(self.key)
self.values = await get_content(self.key)
self.refreshed = pendulum.utcnow()
@ -82,18 +83,74 @@ 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', 'battling', 'battle_records', 'boops', 'server_alerts', 'user_notifications',
'nsfw_channels', 'custom_permissions', 'rules', 'overwatch', 'picarto', 'twitch', 'strawpolls', 'tags',
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 = {}
sharded_data = {}
# Populate cache with each object
for k in possible_keys:
cache[k] = Cache(k)
# 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:
cache[k] = Cache(k)"""
# We still need 'cache' for prefixes and custom permissions however, so for now, just include that
cache['prefixes'] = Cache('prefixes')
cache['custom_permissions'] = Cache('custom_permissions')
async def update_records(key, winner, loser):
# We're using the Harkness scale to rate
# http://opnetchessclub.wikidot.com/harkness-rating-system
r_filter = lambda row: (row['member_id'] == winner.id) | (row['member_id'] == loser.id)
matches = await get_content(key, r_filter)
winner_stats = {}
loser_stats = {}
for stat in matches:
if stat.get('member_id') == winner.id:
winner_stats = stat
elif stat.get('member_id') == loser.id:
loser_stats = stat
winner_rating = winner_stats.get('rating') or 1000
loser_rating = loser_stats.get('rating') or 1000
# The scale is based off of increments of 25, increasing the change by 1 for each increment
# That is all this loop does, increment the "change" for every increment of 25
# The change caps off at 300 however, so break once we are over that limit
difference = abs(winner_rating - loser_rating)
rating_change = 0
count = 25
while count <= difference:
if count > 300:
break
rating_change += 1
count += 25
# 16 is the base change, increased or decreased based on whoever has the higher current rating
if winner_rating > loser_rating:
winner_rating += 16 - rating_change
loser_rating -= 16 - rating_change
else:
winner_rating += 16 + rating_change
loser_rating -= 16 + rating_change
# Just increase wins/losses for each person, making sure it's at least 0
winner_wins = winner_stats.get('wins') or 0
winner_losses = winner_stats.get('losses') or 0
loser_wins = loser_stats.get('wins') or 0
loser_losses = loser_stats.get('losses') or 0
winner_wins += 1
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}
await update_content(key, {'member_id': winner.id}, winner_stats)
await update_content(key, {'member_id': loser.id}, loser_stats)
def command_prefix(bot, message):
@ -102,58 +159,109 @@ def command_prefix(bot, message):
# 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:
prefix = cache['prefixes'].values.get(message.server.id)
values = cache['prefixes'].values
prefix = [data['prefix'] for data in values if message.server.id == data['server_id']][0]
return prefix or default_prefix
except KeyError:
return default_prefix
async def add_content(table, filter, 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:
cursor = await r.table(table).filter(filter).run(conn)
if len(list(cursor)) > 0:
await conn.close()
return False
else:
await r.table(table).insert(content).run(conn)
await conn.close()
return True
except r.ReqlOpFailedError:
# This means the table does not exist
await r.create_table(table).run(conn)
await r.table(table).filter(filter).insert(content).run(conn)
await con.close()
return False
async def update_content(table, filter, content):
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
# Will take some testing, but I might not even need the add_content method above
await r.table(table).filter(filter).update(content).run(conn)
except r.ReqlOpFailedError:
async def add_content(table, content, r_filter=None):
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:
if r_filter is not None:
cursor = await r.table(table).filter(r_filter).run(conn)
content = await _convert_to_list(cursor)
if len(content) > 0:
await conn.close()
return False
else:
await r.table(table).insert(content).run(conn)
await conn.close()
return True
except r.ReqlOpFailedError:
# This means the table does not exist
await r.create_table(table).run(conn)
await r.table(table).insert(content).run(conn)
await conn.close()
return True
async def get_content(key: str, filter):
r.set_loop_type("asyncio")
conn = await r.connect(**db_opts)
try:
cursor = await r.table(key).filter(filter).run(conn)
content = list(cursor.items)[0]
del content['id']
except (IndexError, r.ReqlOpFailedError):
content = None
await conn.close()
return content
async def remove_content(table, r_filter=None):
if r_filter is None:
r_filter = {}
r.set_loop_type("asyncio")
conn = await r.connect(**db_opts)
try:
result = await r.table(table).filter(r_filter).delete().run(conn)
except r.ReqlOpFailedError:
result = {}
pass
await conn.close()
return result.get('deleted', 0) > 0
async def update_content(table, content, r_filter=None):
if r_filter is None:
r_filter = {}
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).filter(r_filter).update(content).run(conn)
except r.ReqlOpFailedError:
await conn.close()
result = {}
await conn.close()
return result.get('replaced', 0) > 0 or result.get('unchanged', 0) > 0
async def replace_content(table, content, r_filter=None):
# This method is here because .replace and .update can have some different functionalities
if r_filter is None:
r_filter = {}
r.set_loop_type("asyncio")
conn = await r.connect(**db_opts)
try:
result = await r.table(table).filter(r_filter).replace(content).run(conn)
except r.ReqlOpFailedError:
await conn.close()
result = {}
await conn.close()
return result.get('replaced', 0) > 0 or result.get('unchanged', 0) > 0
async def get_content(key: str, r_filter=None):
if r_filter is None:
r_filter = {}
r.set_loop_type("asyncio")
conn = await r.connect(**db_opts)
try:
cursor = await r.table(key).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