dr.botzo/dr_botzo/markov/ircplugin.py

import logging
import re

import irc.client

from ircbot.lib import Plugin, reply_destination_for_event
import markov.lib as markovlib
from markov.models import MarkovContext, MarkovState, MarkovTarget

log = logging.getLogger('markov.ircplugin')


class Markov(Plugin):

    """Build Markov chains and reply with them."""

    def start(self):
        """Set up the handlers."""

        self.connection.add_global_handler('pubmsg', getattr(self, 'handle_chatter'), -20)
        self.connection.add_global_handler('privmsg', getattr(self, 'handle_chatter'), -20)

        self.connection.reactor.add_global_regex_handler(['pubmsg', 'privmsg'], r'^!markov\s+reply(\s+min=(\d+))?(\s+max=(\d+))?(\s+(.*)$|$)',
                                                         getattr(self, 'handle_reply'), -20)

        super(Markov, self).start()

    def stop(self):
        """Tear down handlers."""

        self.connection.remove_global_handler('pubmsg', getattr(self, 'handle_chatter'))
        self.connection.remove_global_handler('privmsg', getattr(self, 'handle_chatter'))

        self.connection.reactor.remove_global_regex_handler(['pubmsg', 'privmsg'], getattr(self, 'handle_reply'))

        super(Markov, self).stop()

    def handle_reply(self, connection, event, match):
        """Generate a reply to one line, without learning it."""

        target = reply_destination_for_event(event)

        min_size = 15
        max_size = 30
        context = markovlib.get_or_create_target_context(target)

        if match.group(2):
            min_size = int(match.group(2))
        if match.group(4):
            max_size = int(match.group(4))

        if match.group(5) != '':
            line = match.group(6)
            topics = [x for x in line.split(' ') if len(x) >= 3]

            return self.bot.reply(event, u" ".join(markovlib.generate_line(context, topics=topics,
                                                                           min_words=min_size, max_words=max_size,
                                                                           max_sentences=1)))
        else:
            return self.bot.reply(event, u" ".join(markovlib.generate_line(context, min_words=min_size,
                                                                           max_words=max_size,
                                                                           max_sentences=1)))

    def handle_chatter(self, connection, event):
        """Learn from IRC chatter."""

        what = event.arguments[0]
        my_nick = connection.get_nickname()
        trimmed_what = re.sub(r'^{0:s}[:,]\s+'.format(my_nick), '', what)
        nick = irc.client.NickMask(event.source).nick
        target = reply_destination_for_event(event)

        # learn the line
        recursing = getattr(event, '_recursing', False)
        if not recursing:
            log.debug(u"learning %s", trimmed_what)
            context = markovlib.get_or_create_target_context(target)
            markovlib.learn_line(trimmed_what, context)

        log.debug(u"searching '%s' for '%s'", what, my_nick)
        if re.search(my_nick, what, re.IGNORECASE) is not None:
            context = markovlib.get_or_create_target_context(target)

            addressed_pattern = r'^{0:s}[:,]\s+(.*)'.format(my_nick)
            addressed_re = re.compile(addressed_pattern)
            if addressed_re.match(what):
                # i was addressed directly, so respond, addressing
                # the speaker
                topics = [x for x in addressed_re.match(what).group(1).split(' ') if len(x) >= 3]

                return self.bot.reply(event, u"{0:s}: {1:s}"
                                      u"".format(nick, u" ".join(markovlib.generate_line(context,
                                                                                         topics=topics,
                                                                                         max_sentences=1))))
            else:
                # i wasn't addressed directly, so just respond
                topics = [x for x in what.split(' ') if len(x) >= 3]

                return self.bot.reply(event, u"{0:s}"
                                      u"".format(u" ".join(markovlib.generate_line(context,
                                                                                   topics=topics,
                                                                                   max_sentences=1))))


plugin = Markov