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)
