nxy/bot/plugins/tell.py

# -*- coding: utf-8 -*-
from datetime import datetime

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

from . import DatabasePlugin


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

    def __init__(self, bot: irc3.IrcBot):
        super().__init__(bot)
        self.tell_queue = {}

        # Fetch tells from database
        self.cur.execute('''
        select
           to_nick, from_nick, message, created_at
        from
          tells
        ''')

        # Add tells to queue
        for res in self.cur.fetchall():
            nick = res['to_nick'].lower()

            # Create list in queue and add tell
            if nick not in self.tell_queue:
                self.tell_queue[nick] = []
            self.tell_queue[nick].append(res[1:])

    @command
    def tell(self, mask: IrcString, target: IrcString, args: DocOptDict):
        """Saves a message for nick to forward on activity

        %%tell <nick> <message>...
        """
        nick = args['<nick>']
        nick_lower = nick.lower()
        tell = [nick, mask.nick, ' '.join(args['<message>']).strip(),
                datetime.now()]

        # Create list in queue and add tell
        if nick_lower not in self.tell_queue:
            self.tell_queue[nick_lower] = []
        self.tell_queue[nick_lower].append(tell[1:])

        try:
            # Insert tell into database
            self.cur.execute('''
            insert into
              tells (to_nick, from_nick, message, created_at)
            values
              (%s, %s, %s, %s)
            ''', tell)
            self.con.commit()
        except Error:
            # Rollback transaction on error
            self.con.rollback()

    @irc3.event(r'(?i)^:(?P<mask>.*) PRIVMSG .* :.*')
    def check(self, mask: str):
        """If activ nick has tells, forward and delete them."""
        nick = IrcString(mask).nick
        nick_lower = nick.lower()

        if nick_lower in self.tell_queue:
            # Forward all tells for nick
            for tell in self.tell_queue[nick_lower]:
                # TODO: format time
                self.bot.privmsg(nick, '[Tell] Message from {nick} at {time}: '
                                       '{message}'.format(nick=tell[0],
                                                          time=tell[2],
                                                          message=tell[1]))

            # Remove nick from queue
            del self.tell_queue[nick_lower]

            try:
                # Remove tells from database
                self.cur.execute('''
                delete from
                  tells
                where
                  lower(to_nick) = lower(%s)
                ''', [nick])
                self.con.commit()
            except Error as ex:
                print(ex)
                # Rollback transaction on error
                self.con.rollback()