dr.botzo/markov/ircplugin.py
Brian S. Stephan 43f2b09057 don't add the empty string to additional nicks
thinko on my part, this was making the regex for matching all nicks to
'|nick' when the field is '', because of split producing ['']. in
particular this was making markov trigger on every line
2021-04-25 21:00:34 -05:00

110 lines
4.6 KiB
Python

"""IRC support for Markov chain learning and text generation."""
import logging
import re
import irc.client
from django.conf import settings
import markov.lib as markovlib
from ircbot.lib import Plugin, reply_destination_for_event
from ircbot.models import IrcChannel
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]
if connection.server_config.additional_addressed_nicks:
all_nicks = '|'.join(connection.server_config.additional_addressed_nicks.split('\n') +
[connection.get_nickname()])
else:
all_nicks = 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, server=connection.server_config)
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