A script to remove solutions from a Jupyter notebook.

prepare_student_notebook.py

"""This script removes solutions from a given notebook and saves it as a new notebook.

Input: a Jupyter notebook with solutions 

The solutions are marked according to the nbgrader's format:
### BEGIN SOLUTION 
...
### END SOLUTION

Output: a path to a new notebook without solutions."""

import sys
import nbformat


def remove_solutions(notebook_path: str) -> nbformat.notebooknode.NotebookNode:
    """Remove solutions from a notebook and save it as a new notebook.

    The solution block starts with ### BEGIN SOLUTION and ends with ### END SOLUTION.
    We replace each block with "# YOUR SOLUTION HERE"

    There may be multiple solution blocks and the cell may not start with ### BEGIN SOLUTION."""

    notebook = nbformat.read(notebook_path, as_version=4)
    for cell in notebook.cells:
        # we can skip if the cell doesn't contain the solution
        if "### BEGIN SOLUTION" not in cell.source:
            continue

        # the cell will have the following string:
        # "... \n\s*### BEGIN SOLUTION\n ... \n\s*### END SOLUTION\n ..."
        # When we split with "### BEGIN SOLUTION\n", we get a list of strings
        # [" ... \n\s*", "<solution string> ### END SOLUTION\n ... \n\s*", ...,
        #  "<solution string> ### END SOLUTION\n ... "]
        blocks = cell.source.split("### BEGIN SOLUTION\n")
        for i in range(1, len(blocks)):
            # we want to remove the "### END SOLUTION\n" from the solution string
            blocks[i] = "# YOUR SOLUTION HERE" + blocks[i].split("### END SOLUTION")[1]
        cell.source = "".join(blocks)

    # return notebook object
    return notebook


if __name__ == "__main__":
    notebook_path = sys.argv[1]
    output_path = "student_version/" + notebook_path
    notebook = remove_solutions(notebook_path)
    # Save notebook
    nbformat.write(notebook, output_path)