tilde-projects/Code/irc/madlibbot/madlib.py

#!/usr/bin/python

import collections
import glob
import inflect
import random
import re

p = inflect.engine()
file_pattern = "/home/*/madlibs/*.madlib"
file_regex = re.compile(r"^/home/(.*?)/.*/([^/]*)\.madlib$")
word_regex = re.compile(r"{{(.*?)(#.*?)?(\|.*)?}}")
word_replace_regex = re.compile(r"^(.*)#(\d*)$")

# Take a file path and return a tuple with the title and original path
def munge_story(file_name):
    match = re.match(file_regex, file_name)
    count = count_words_file(file_name)
    return (
        "'" + match.group(2).replace("_", " ") + "' by ~" + match.group(1),
        match.group(0),
        count,
    )


# Retrive a list of all madlib story files
def find_stories(limit=999, shuffle=False):
    files = glob.glob(file_pattern)
    if shuffle == True:
        files = random.sample(files, max(1, min(limit, len(files))))
    else:
        files = files[:limit]
    return map(munge_story, files)


# Count the number of replacable words in the story
def count_words(story):
    # words that have a '#' part and match should only be counted once
    count = 0
    repeats = []
    for match in re.finditer(word_regex, story):
        if match.group(2) is None:  # the '#' part
            count += 1
        elif match.group(1) + match.group(2) not in repeats:
            count += 1
            repeats.append(match.group(1) + match.group(2))
    return count


# Count the number of replacable words in the story when given a file path
def count_words_file(storyPath):
    with open(storyPath, "r") as storyFile:
        story = storyFile.read()
        return count_words(story)


# Get the next replaceable word in the story. Includes full token and just the word
# If you specify rand as True, it will get a random word instead of the next one
def find_next_word(story, rand=False):
    matches = list(word_regex.finditer(story))
    if len(matches) == 0:
        return None
    match = matches[0]
    if rand is True:
        match = random.choice(matches)

    return (match.group(0), match.group(1))


# Using a query phrase, replace a word and return the entire story body
def replace_word(story, query, word):
    rquery = word_regex.search(query)
    # '{{foo#bar|baz|bang}}' => ('foo', '#bar', '|baz|bang')
    pipes = [
        p.lower()
        for p in (rquery.group(3) if rquery.group(3) is not None else "")
        .strip("|")
        .split("|")
    ]
    munged_word = process_pipes(word, pipes)
    story = story.replace(query, munged_word, 1)

    if rquery.group(2) is not None:  # if there is a '#' part we replace all instances
        print("Looking for {{" + rquery.group(1) + rquery.group(2) + ".*?}}")
        others = re.findall(r"{{" + rquery.group(1) + rquery.group(2) + ".*?}}", story)
        if len(others) > 0:
            story = replace_word(
                story, others[0], word
            )  # iteratively replace the next instance

    return story


# Modify user input based on certain modifying pipes
def process_pipes(word, pipes):
    for pipe in pipes:
        try:
            word = {
                "upper": lambda word: word.upper(),
                "lower": lambda word: word.lower(),
                "title": lambda word: word.title(),
                "numeric": lambda word: p.number_to_words(word),
                "ordinal": lambda word: p.ordinal(word),
                collections.defaultdict: lambda word: word,
            }[pipe](word)
        except Exception:
            pass  # just keep going if an error occurs processing a pipe
    return word


# A helper function that will split the story up into chat-printable lengths
def yield_lines(line, max_size):
    words = []
    words_length = 0
    # Add words to the words line untill we run out of space
    for word in line.split():
        # If the next word itself is longer thatn max_size, we'll have to chop it up
        while len(word) > max_size:
            splitsize = max_size - words_length - 1
            words.append(word[:splitsize] + "-")
            yield " ".join(words)
            words = []
            words_length = 0
            word = word[splitsize:]
        if words_length + len(word) + 1 > max_size:
            yield " ".join(words)
            words = []
            words_length = 0
        words.append(word)
        words_length += len(word) + 1  # For the space
    yield " ".join(words)  # For any words left over


# Open a story file and ask for the parts
def make_story(storyFile):
    with open(storyFile, "r") as storyFile:
        story = storyFile.read()
        match = word_regex.search(story)
        while match is not None:
            word = raw_input("Give me {}: ".format(match.group(1)))
            print("Replacing '{}' with '{}'".format(match.group(0), word))
            print(story)
            story = re.sub(match.group(0), word, story, 1)
            match = word_regex.search(story)
        print(story)


def start():
    stories = find_stories(20, True)
    if len(stories) == 0:
        print("Sorry, no stories found. :/")
        return

    input = -1
    while input < 0 or input >= len(stories):
        for idx, story in enumerate(stories):
            print("[{}] {}".format(idx, story[0]))
        input = raw_input("Which story would you like?: ")
        try:
            # Try to convert the string input into a number
            input = int(input)
        except:
            # If they put something stupid in, treat it as a -1 and ask for
            # input again
            input = -1
    # Call make_story with the file path of the selected story
    make_story(stories[input][1])


if __name__ == "__main__":
    start()