Simple Markov Chain in Python

markov.py

import random

def markov(string, length):
    list_of_tokens = string.split()
    markov_dict = {}
    prev = None
    for token in list_of_tokens:
        if prev is not None:
            if prev in markov_dict.keys():
                markov_dict[prev].append(token)
            else:
                markov_dict[prev] = [token]
        prev = token
    token = random.choice(list(markov_dict.keys()))
    new_text = [token]
    for i in range(length-1):
        if token in markov_dict.keys():
            choice = random.choice(markov_dict[token])
        else:
            token = random.choice(list(markov_dict.keys()))
            choice = random.choice(markov_dict[token])
        new_text.append(choice)
        token = choice
    return markov_dict, ' '.join(new_text)