support multiple first_words

This commit is contained in:
jesopo 2019-09-30 16:27:29 +01:00
parent a75802129f
commit 3c60de8164

View file

@ -90,7 +90,7 @@ class Module(ModuleManager.BaseModule):
@utils.kwarg("usage", "[first-word]") @utils.kwarg("usage", "[first-word]")
def markov(self, event): def markov(self, event):
self._markov_for(event["target"], event["stdout"], event["stderr"], self._markov_for(event["target"], event["stdout"], event["stderr"],
first_word=(event["args_split"] or [None])[0]) first_words=event["args_split"][:])
@utils.hook("received.command.markovfor") @utils.hook("received.command.markovfor")
@utils.kwarg("min_args", 1) @utils.kwarg("min_args", 1)
@ -101,21 +101,21 @@ class Module(ModuleManager.BaseModule):
if event["args_split"][0] in event["server"].channels: if event["args_split"][0] in event["server"].channels:
channel = event["server"].channels.get(event["args_split"][0]) channel = event["server"].channels.get(event["args_split"][0])
self._markov_for(channel, event["stdout"], event["stderr"], self._markov_for(channel, event["stdout"], event["stderr"],
first_word=(event["args_split"][1:] or [None])[0]) first_words=event["args_split"][1:])
else: else:
event["stderr"].write("Unknown channel") event["stderr"].write("Unknown channel")
def _markov_for(self, channel, stdout, stderr, first_word=None): def _markov_for(self, channel, stdout, stderr, first_words):
if not channel.get_setting("markov", False): if not channel.get_setting("markov", False):
stderr.write(NO_MARKOV) stderr.write(NO_MARKOV)
else: else:
out = self._generate(channel.id, first_word) out = self._generate(channel.id, first_words)
if not out == None: if not out == None:
stdout.write(out) stdout.write(out)
else: else:
stderr.write("Failed to generate markov chain") stderr.write("Failed to generate markov chain")
def _generate(self, channel_id, first_word): def _generate(self, channel_id, first_words):
if first_word == None: if first_word == None:
first_words = self.bot.database.execute_fetchall("""SELECT first_words = self.bot.database.execute_fetchall("""SELECT
third_word, frequency FROM markov WHERE channel_id=? AND third_word, frequency FROM markov WHERE channel_id=? AND
@ -135,17 +135,20 @@ class Module(ModuleManager.BaseModule):
second_word = self._choose(second_words) second_word = self._choose(second_words)
words = [first_word, second_word] words = [first_word, second_word]
else: else:
first_word = first_word.lower() if len(first_words) == 1:
second_two_words = self.bot.database.execute_fetchall("""SELECT first_word = first_word.lower()
second_word, third_word, frequency FROM markov WHERE second_two_words = self.bot.database.execute_fetchall("""SELECT
channel_id=? AND first_word=? AND second_word NOT NULL AND second_word, third_word, frequency FROM markov WHERE
third_word NOT NULL""", [channel_id, first_word]) channel_id=? AND first_word=? AND second_word NOT NULL AND
if not second_two_words: third_word NOT NULL""", [channel_id, first_word])
return None if not second_two_words:
return None
second_word, third_word = self._choose( second_word, third_word = self._choose(
[[[s, t], f] for s, t, f in second_two_words]) [[[s, t], f] for s, t, f in second_two_words])
words = [first_word, second_word, third_word] words = [first_word, second_word, third_word]
else:
words = [word.lower() for word in first_words]
for i in range(30): for i in range(30):
two_words = words[-2:] two_words = words[-2:]