def save_tasks_py(nb):
    """Select cells appropriate for `tasks.py`."""
    shouldnt_start_with = '# Put in a file named run-readout-scan.py'
    tasks_py = nb.copy()
    tasks_py.cells = [
        cell for cell in tasks_py.cells
        if not cell.source.startswith(shouldnt_start_with)
    ]

    exporter = PythonExporter()
    exporter.raw_template = TEMPLATE
    output, resources = exporter.from_notebook_node(tasks_py)
    with open(f'{REPO_DIR}/recirq/readout_scan/tasks.py', 'w') as f:
        f.write(output)
def save_run_py(nb):
    """Select cells appropriate for `run-readout-scan.py`."""
    should_start_with = '# Put in a file named run-readout-scan.py'
    run_py = nb.copy()
    run_py.cells = [
        cell for cell in run_py.cells
        if cell.source.startswith(should_start_with)
    ]

    # Fix-up the content
    first_cell_lines = run_py.cells[0].source.splitlines()
    new_first_cell_lines = []
    added_imports = False
    for line in first_cell_lines:
        if line.startswith(should_start_with):
            # Ignore this comment, which is out of place in the exported py file
            continue

        if line.startswith('import') and not added_imports:
            # Add imports that aren't necessary when it's in one file.
            new_first_cell_lines += [
                'from recirq.readout_scan.tasks import ReadoutScanTask, run_readout_scan'
            ]
            added_imports = True

        # Rest of the lines are unmodified
        new_first_cell_lines += [line]

    # Mutate the original cell
    run_py.cells[0].source = '\n'.join(new_first_cell_lines)

    exporter = PythonExporter()
    exporter.raw_template = TEMPLATE
    output, resources = exporter.from_notebook_node(run_py)
    with open(f'{REPO_DIR}/recirq/readout_scan/run-readout-scan.py', 'w') as f:
        f.write(output)