dr.botzo/markov/ircplugin.py

110 lines
4.3 KiB
Python

import logging
import re
import irc.client
from django.conf import settings
from ircbot.lib import Plugin, reply_destination_for_event
from ircbot.models import IrcChannel
import markov.lib as markovlib
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', self.handle_chatter, -20)
self.connection.add_global_handler('privmsg', self.handle_chatter, -20)
self.connection.reactor.add_global_regex_handler(['pubmsg', 'privmsg'],
r'^!markov\s+reply(\s+min=(\d+))?(\s+max=(\d+))?(\s+(.*)$|$)',
self.handle_reply, -20)
super(Markov, self).start()
def stop(self):
"""Tear down handlers."""
self.connection.remove_global_handler('pubmsg', self.handle_chatter)
self.connection.remove_global_handler('privmsg', self.handle_chatter)
self.connection.reactor.remove_global_regex_handler(['pubmsg', 'privmsg'], 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, " ".join(markovlib.generate_line(context, topics=topics,
min_words=min_size, max_words=max_size)))
else:
return self.bot.reply(event, " ".join(markovlib.generate_line(context, min_words=min_size,
max_words=max_size)))
def handle_chatter(self, connection, event):
"""Learn from IRC chatter."""
what = event.arguments[0]
all_nicks = '|'.join(settings.ADDITIONAL_NICK_MATCHES + [connection.get_nickname()])
trimmed_what = re.sub(r'^(({nicks})[:,]|@({nicks}))\s+'.format(nicks=all_nicks), '', what)
nick = irc.client.NickMask(event.source).nick
target = reply_destination_for_event(event)
# check to see whether or not we should learn from this channel
channel = None
if irc.client.is_channel(target):
channel, c = IrcChannel.objects.get_or_create(name=target)
if channel and not channel.markov_learn_from_channel:
log.debug("not learning from %s as i've been told to ignore it", channel)
else:
# learn the line
recursing = getattr(event, 'recursing', False)
if not recursing:
log.debug("learning %s", trimmed_what)
context = markovlib.get_or_create_target_context(target)
markovlib.learn_line(trimmed_what, context)
log.debug("searching '%s' for '%s'", what, all_nicks)
if re.search(all_nicks, what, re.IGNORECASE) is not None:
context = markovlib.get_or_create_target_context(target)
addressed_pattern = r'^(({nicks})[:,]|@({nicks}))\s+(?P<addressed_msg>.*)'.format(nicks=all_nicks)
match = re.match(addressed_pattern, what, re.IGNORECASE)
if match:
# i was addressed directly, so respond, addressing
# the speaker
topics = [x for x in match.group('addressed_msg').split(' ') if len(x) >= 3]
return self.bot.reply(event, "{0:s}: {1:s}"
"".format(nick, " ".join(markovlib.generate_line(context, topics=topics))))
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, "{0:s}"
"".format(" ".join(markovlib.generate_line(context, topics=topics))))
plugin = Markov