Markov Chain

markov_chain.py

import random

nonWord = "\n"

class Chain:
    words = {}

    def __init__(self, data=None, base=None):
        if data != None:
            self.scan(data)
        if base != None:
            self.words = base

    def scan(self, data):
        w1 = None
        w2 = None
        for word in data.split(" "):
            self.words.setdefault((w2, w1), []).append(word)
            w2, w1 = w1, word

        self.words.setdefault((w2, w1), []).append(nonWord)

    def generate(self, max_words=80):
        w1 = None
        w2 = None
        out = ""
        for i in range(max_words):
            word = random.choice(self.words[(w1,w2)])
            if word == nonWord:
                break
            out += word + " "
            w1 , w2 = w2, word
        return out