nxy/bot/plugins/quotes.py

# -*- coding: utf-8 -*-
import re

import irc3
from docopt import Dict
from irc3.plugins.command import command
from irc3.utils import IrcString
from psycopg2 import Error

from . import DatabasePlugin
from ..utils import parse_int


@irc3.plugin
class Quotes(DatabasePlugin):
    requires = ['irc3.plugins.command',
                'bot.plugins.storage']

    def add_quote(self, mask, nick, quote, channel):
        # Parse nick from "<@foobar>" like strings
        nick = re.match(r'<?[~&@%+]?([a-zA-Z0-9_\-^`|\\\[\]{}]+)>?', nick).group(1)

        if not nick:
            self.bot.notice(mask.nick, '[Quotes] Error parsing nick')
        else:
            # Insert quote into database
            self.cur.execute('''
            INSERT INTO
              quotes (nick, item, channel, created_by)
            VALUES
              (%s, %s, %s, %s)
            ''', [nick, quote, channel, mask.nick])

    def delete_quote(self, nick, quote):
        index, order = parse_int(quote, select=False)

        if index:
            # Delete from database
            self.cur.execute('''
            with ranked_quotes as (
              select
                id,
                rank() over (partition by nick order by id {order})
              from
                quotes
              where
                nick = %s
            )

            delete from
              quotes
            where
              id = (
                select
                  id
                from
                  ranked_quotes
                where
                  rank = %s
              )
            '''.format(order=order), [nick, index])

    @command(options_first=True)
    def q(self, mask: IrcString, target: IrcString, args: Dict):
        """Get, add or delete quotes for an user.

        %%q <cmd> <nick> <quote>...
        %%q <nick> [<index>]
        """
        cmd = args.get('<cmd>')
        nick = args['<nick>']
        quote = ' '.join(args.get('<quote>'))

        if cmd and quote:
            try:
                # Anybody can add
                if cmd == 'add':
                    self.add_quote(mask, nick, quote, target)
                # But only admins can delete
                elif cmd == 'del' and self.guard.has_permission(mask, 'admin'):
                    self.delete_quote(nick, quote)

                self.con.commit()
            except Error as ex:
                # Rollback transaction on error
                print(ex)
                self.con.rollback()
        else:
            index = args.get('<index>')
            values = [nick]

            if index:
                index = parse_int(index)
                if not index:
                    return
                index, order = index
                order = 'rank {order}'.format(order=order)
                offset = 'offset %s'
                values.append(index)
            else:
                order = 'random()'
                offset = ''

            # Fetch quote from database
            self.cur.execute("""
            with ranked_quotes as (
              select
                nick,
                item,
                rank() over (partition by nick order by id),
                count(*) over (partition by nick) as total
              from
                quotes
            )

            select
              *
            from
              ranked_quotes
            where
              lower(nick) like lower(%s)
            order by
              {order}
            limit
              1
            {offset}
            """.format(order=order, offset=offset), values)
            result = self.cur.fetchone()

            if result:
                return '[{rank}/{total}] <{nick}> {item}'.format(**result)