Useful for courses. Replace the solution code with a NotImplementedError, so that you don't have to duplicate code.

strip_solutions.py

import click
import ast
import astor
from pathlib import Path
import logging

log = logging.getLogger()
logging.basicConfig(level=logging.INFO)


def test_process_function():
    code = ast.parse(Path("test_minimal.py").read_text())
    code = process_ast(code, "test_")
    assert len(code.body[1].body) == 2


def process_ast(code, skip_pattern=""):
    rs = ast.parse("raise NotImplementedError")

    for f in code.body:
        if f.__class__ != ast.FunctionDef:
            continue

        if f.name.startswith(skip_pattern):
            continue

        apply_suffix = "_solution"
        if not f.name.endswith(apply_suffix):
            continue

        log.info(astor.to_source(f))
        new_body = []
        # preserve comments.
        if astor.to_source(f.body[0]).startswith('"""'):
            new_body.append(f.body[0])
        new_body.append(rs)
        f.body = new_body
        f.name = f.name[: -len(apply_suffix)]
        log.info(astor.to_source(f))
    return code


def process_file(fpath, skip_pattern=""):
    log.info("Process %r %r", fpath, skip_pattern)
    code = ast.parse(Path(fpath).read_text())

    new_code = process_ast(code, skip_pattern)

    Path(fpath + ".ex.py").write_text(astor.to_source(new_code))


@click.command()
@click.argument("files", nargs=-1)
@click.option(
    "--skip-pattern", default="", help="Skip functions starting with"
)
def main(files, skip_pattern=""):
    """Replace function code with 
       raise NotImplementedError("Write me")
       for python exercises.
    """

    for fpath in files:
        process_file(fpath, skip_pattern)


if __name__ == "__main__":
    main()