diff --git a/nxy/plugins/rape.py b/nxy/plugins/rape.py index cf92b52..838925e 100644 --- a/nxy/plugins/rape.py +++ b/nxy/plugins/rape.py @@ -4,37 +4,34 @@ import random from docopt import Dict as DocOptDict from irc3.plugins.command import command from irc3.utils import IrcString +import irc3 from . import DatabasePlugin +@irc3.plugin class Rape(DatabasePlugin): requires = ['irc3.plugins.command', 'nxy.plugins.database'] @command - def owe(self, mask: IrcString, channel: IrcString, args: DocOptDict): + def owe(self, mask: IrcString, target: IrcString, args: DocOptDict): """Shows how much a nick owes. - %%owe [...] + %%owe [] """ - nick = ' '.join(args.get('')) or mask.nick + nick = args.get('') or mask.nick - if nick.startswith('@'): - nick = nick.lstrip('@').strip() - - self.cur.execute('select amount from owes where nick = ?', [nick]) + # Fetch result from database + self.cur.execute('select amount from owes where nick = %s', [nick]) owes = self.cur.fetchone() - if owes: - total = '4${total}'.format(total=owes['amount']) - else: - total = '3$0' - + # Colorize owe amount and return string + total = '4${total}'.format(total=owes['amount']) if owes else '3$0' return '{nick} owes: \x03{total}\x03'.format(nick=nick, total=total) @command - def rape(self, mask: IrcString, channel: IrcString, args: DocOptDict): + def rape(self, mask: IrcString, target: IrcString, args: DocOptDict): """Rapes a nick and eventually charge for it. %%rape @@ -44,19 +41,12 @@ class Rape(DatabasePlugin): rand = random.randint(0, 3) if rand in (0, 1): - self.cur.execute('select amount from owes where nick = ?', [nick]) - owe = self.cur.fetchone() - - if owe: - total = owe['amount'] + fine - self.cur.execute('update owes set amount = ? where nick = ?', - [total, nick]) - else: - total = fine - self.cur.execute('insert into owes (nick, amount) values ' - '(?, ?)', [nick, total]) - self.con.commit() + self.cur.execute('''insert into owes (nick, amount) values (%s, %s) + on conflict (nick) do update set amount = owes.amount + + excluded.amount returning amount''', [nick, fine]) + self.con.commit() + total = self.cur.fetchone()['amount'] reason = ('raping', 'being too lewd and getting raped')[rand] action = 'fines {nick} \x02${fine}\x02 for {reason}. You owe: ' \ '\x0304${total}\x03'.format(nick=nick, @@ -65,4 +55,4 @@ class Rape(DatabasePlugin): total=total) else: action = 'rapes {nick}'.format(nick=nick) - self.bot.action(channel, action) + self.bot.action(target, action)