add markov generation mode to cli

trekipsum/cli.py

@@ -3,7 +3,7 @@
 import argparse
 import logging
 
-from trekipsum import dialog
+from trekipsum import dialog, markov
 
 logger = logging.getLogger(__name__)
 
@@ -19,6 +19,8 @@ def positive(value):
 def parse_cli_args():
     """Parse command-line arguments."""
     parser = argparse.ArgumentParser(description='TrekIpsum generator')
+    parser.add_argument('-m', '--markov', action='store_true',
+                        help='use markov chain mode for generation')
     parser.add_argument('--speaker', type=str,
                         help='limit output to this speakers')
     parser.add_argument('-a', '--attribute', action='store_true',
@@ -47,7 +49,12 @@ def main_cli():
     loglevel = logging.DEBUG if args.debug else logging.CRITICAL
     logging.basicConfig(level=loglevel, format='%(asctime)s %(levelname)s: %(message)s')
     logger.setLevel(loglevel)
-    chooser = dialog.SqliteRandomChooser()
+
+    if args.markov is True:
+        chooser = markov.MarkovRandomChooser()
+    else:
+        chooser = dialog.SqliteRandomChooser()
+
     for paragraph in range(args.paragraphs):
         speaker = args.speaker
         lines = []

trekipsum/markov/__init__.py

@@ -6,6 +6,7 @@
 import six
 
 from ..dialog.sqlite import DEFAULT_SQLITE_PATH
+from ..exceptions import NoDialogFoundException
 
 SENTENCE_DELIMITER = ''  # special value for beginning/ending a sentence
 
@@ -214,3 +215,34 @@ def to_chain(self, context):
         for link in links:
             chain[link[0]].append((link[1], link[2]))
         return chain
+
+
+class MarkovRandomChooser(object):
+    """Walk Markov chains to generate dialog from datastore."""
+
+    def __init__(self):
+        """Initialize a new dialog markov chain chooser for dialog."""
+        self._datastore = DialogChainDatastore()
+        self._speaker_walker = ChainWalker(self._datastore.to_chain('speakers'))
+
+    def random_dialog(self, speaker):
+        """
+        Get random line of dialog, optionally limited to specific speaker.
+
+        Returns:
+            tuple containing (speaker name, line of dialog)
+        """
+        speaker = speaker.upper() if speaker else None
+        if speaker is None:
+            speaker = self.random_speaker()
+        dialog_walker = ChainWalker(self._datastore.to_chain(speaker))
+        try:
+            return speaker, dialog_walker.build_sentence()
+        except KeyError:
+            raise NoDialogFoundException(speaker)
+
+    def random_speaker(self, from_speaker=None):
+        """
+        Get random speaker name, optionally walking from a specific speaker.
+        """
+        return self._speaker_walker.next_word(from_speaker)