"""IRC support for Markov chain learning and text generation.""" import logging import re import irc.client import markov.lib as markovlib from ircbot.lib import Plugin, reply_destination_for_event from ircbot.models import IrcChannel from markov.models import MarkovContext, 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', 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 = self.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] who = irc.client.NickMask(event.source).nick target = reply_destination_for_event(event) log.debug("what: '%s', who: '%s', target: '%s'", what, who, target) # 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 learning_what = what # don't learn the speaker's nick if this came over a bridge if channel and who == channel.discord_bridge: learning_what = ' '.join(learning_what.split(' ')[1:]) # remove our own nick and aliases from what we learn 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() learning_what = re.sub(r'^(({nicks})[:,]|@({nicks})[:,]?)\s+'.format(nicks=all_nicks), '', learning_what) recursing = getattr(event, 'recursing', False) if not recursing: log.debug("learning %s", learning_what) context = self.get_or_create_target_context(target) markovlib.learn_line(learning_what, context) log.debug("searching '%s' for '%s'", what, all_nicks) if re.search(all_nicks, what, re.IGNORECASE) is not None: context = self.get_or_create_target_context(target) addressed_pattern = r'^(({nicks})[:,]|@({nicks})[:,]?)\s+(?P.*)'.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(who, " ".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)))) def get_or_create_target_context(self, target_name): """Return the context for a provided nick/channel, creating missing ones.""" target_name = target_name.lower() # find the stuff, or create it channel, c = IrcChannel.objects.get_or_create(name=target_name, server=self.connection.server_config) try: target = MarkovTarget.objects.get(channel=channel) except MarkovTarget.DoesNotExist: # we need to create a context and a target, and we have to make the context first # make a context --- lacking a good idea, just create one with this target name until configured otherwise context, c = MarkovContext.objects.get_or_create(name=target_name) target, c = MarkovTarget.objects.get_or_create(name=target_name, context=context, channel=channel) return target.context try: return target.context except MarkovContext.DoesNotExist: # make a context --- lacking a good idea, just create one with this target name until configured otherwise context, c = MarkovContext.objects.get_or_create(name=target_name) target.context = context target.save() return target.context plugin = Markov