allow an arg to !markov to chose the first word

This commit is contained in:
jesopo 2019-09-23 11:39:43 +01:00
parent bb2590734b
commit c253d17f54

View file

@ -85,7 +85,8 @@ class Module(ModuleManager.BaseModule):
@utils.kwarg("channel_only", True) @utils.kwarg("channel_only", True)
@utils.kwarg("help", "Generate a markov chain for the current channel") @utils.kwarg("help", "Generate a markov chain for the current channel")
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])
@utils.hook("received.command.markovfor") @utils.hook("received.command.markovfor")
@utils.kwarg("min_args", 1) @utils.kwarg("min_args", 1)
@ -99,23 +100,27 @@ class Module(ModuleManager.BaseModule):
else: else:
event["stderr"].write("Unknown channel") event["stderr"].write("Unknown channel")
def _markov_for(self, channel, stdout, stderr): def _markov_for(self, channel, stdout, stderr, first_word=None):
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) out = self._generate(channel.id, first_word)
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): def _generate(self, channel_id, first_word):
first_words = self.bot.database.execute_fetchall("""SELECT third_word, if start == None:
frequency FROM markov WHERE channel_id=? AND first_word IS NULL AND first_words = self.bot.database.execute_fetchall("""SELECT
second_word IS NULL AND third_word NOT NULL""", [channel_id]) third_word, frequency FROM markov WHERE channel_id=? AND
first_word IS NULL AND second_word IS NULL AND third_word
NOT NULL""", [channel_id])
if not first_words: if not first_words:
return None return None
first_word = self._choose(first_words) first_word = self._choose(first_words)
else:
first_word = first_word.lower()
second_words = self.bot.database.execute_fetchall("""SELECT third_word, second_words = self.bot.database.execute_fetchall("""SELECT third_word,
frequency FROM markov WHERE channel_id=? AND first_word IS NULL AND frequency FROM markov WHERE channel_id=? AND first_word IS NULL AND