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)