Migrated rape to pgsql

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 [<nick>...]
+        %%owe [<nick>]
         """
-        nick = ' '.join(args.get('<nick>')) or mask.nick
+        nick = args.get('<nick>') 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 <nick>
@@ -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)