resolve imports and remove solutions to generate questions .tex

questions.py

"""Generate LaTeX that only includes solutions for a particular TeX document.

Usage:
    questions.py [--source=<src>] [--dest=<dest>]
"""

from TexSoup import TexSoup
from docopt import docopt

SOlUTION_COMMAND = 'sol'
IMPORT_COMMAND = 'subimport'


def main(tex):
    """Returns new TexSoup, resolving includes and removing solutions."""

    soup = TexSoup(tex)

    for subimport in soup.find_all(IMPORT_COMMAND):
        path = ''.join(arg.value for arg in subimport.args)
        subimport.replace(*main(open(path)).contents)

    # check resolver works
    assert len(list(soup.find_all(IMPORT_COMMAND))) == 0

    for sol in soup.find_all(SOLUTION_COMMAND):
        sol.delete()

    return soup


if __name__ == '__main__':
    arguments = docopt(__doc__, version='LaTeX Questions 1.0')
    soup = main(open(arguments['--source'] or input('File:').strip()))
    with open(arguments['--dest'] or input('Dest:').strip(), 'w') as f:
        f.write(repr(soup))