def execute_sos_notebook(input_notebook,
output_notebook=None,
return_content=True):
# execute input notebook
# if input_notebook is a string, it will be loaded. Otherwise it should be a notebook object
# if output_notebook is a string, it will be used as output filename. Otherwise
# the notebook will be returned.
try:
from papermill.execute import execute_notebook
except ImportError:
raise RuntimeError(
'Please install papermill for the use of option --execute.')
if not any(entrypoint.name == 'sos'
for entrypoint in pkg_resources.iter_entry_points(
group='papermill.engine')):
raise RuntimeError(
'Please install sos-papermill for the use of option --execute.')
if isinstance(input_notebook, str):
input_file = input_notebook
else:
input_file = tempfile.NamedTemporaryFile(prefix='__tmp_input_nb',
dir=os.getcwd(),
suffix='.ipynb',
delete=False).name
with open(input_file, 'w') as notebook_file:
nbformat.write(input_notebook, notebook_file, 4)
if output_notebook is None:
output_file = tempfile.NamedTemporaryFile(prefix='__tmp_output_nb',
dir=os.getcwd(),
suffix='.ipynb',
delete=False).name
else:
output_file = output_notebook
execute_notebook(input_path=input_file,
output_path=output_file,
engine_name='sos',
kernel_name='sos')
if os.path.basename(input_file).startswith('__tmp_input_nb'):
try:
os.remove(input_file)
except Exception as e:
env.logger.warning(
f'Failed to remove temporary input file {input_file}: {e}')
if os.path.basename(output_file).startswith(
'__tmp_output_nb') and return_content:
new_nb = nbformat.read(output_file, nbformat.NO_CONVERT)
try:
os.remove(output_file)
except Exception as e:
env.logger.warning(
f'Failed to remove temporary output file {output_file}: {e}')
return new_nb
else:
return output_file
def papermill(notebook_path, output_path, parameters, parameters_raw,
parameters_file, parameters_yaml, parameters_base64, kernel):
"""Utility for executing a single notebook on a container.
Take a source notebook, apply parameters to the source notebook,
execute the notebook with the kernel specified, and save the
output in the destination notebook.
"""
if parameters_base64:
parameters_final = yaml.load(base64.b64decode(parameters_base64))
elif parameters_yaml:
parameters_final = yaml.load(parameters_yaml)
elif parameters_file:
parameters_final = read_yaml_file(parameters_file)
else:
# Read in Parameters
parameters_final = {}
for name, value in parameters:
parameters_final[name] = _resolve_type(value)
for name, value in parameters_raw:
parameters_final[name] = value
execute_notebook(notebook_path,
output_path,
parameters_final,
kernel_name=kernel)
def test(self):
nb_test1_fname = get_notebook_path('simple_execute.ipynb')
nb_test1_executed_fname = os.path.join(self.test_dir, 'test1_executed.ipynb')
execute_notebook(nb_test1_fname, nb_test1_executed_fname, {'msg': 'Hello'})
test_nb = read_notebook(nb_test1_executed_fname)
self.assertEqual(test_nb.node.cells[0].get('source'), u'# Parameters\nmsg = "Hello"\n')
self.assertEqual(test_nb.parameters, {'msg': 'Hello'})
def test(self):
path = get_notebook_path('broken.ipynb')
result_path = os.path.join(self.test_dir, 'broken.ipynb')
execute_notebook(path, result_path)
nb = read_notebook(result_path)
self.assertEqual(nb.node.cells[0].execution_count, 1)
self.assertEqual(nb.node.cells[1].execution_count, 2)
self.assertEqual(nb.node.cells[1].outputs[0].output_type, 'error')
self.assertEqual(nb.node.cells[2].execution_count, None)
def test(self):
path = get_notebook_path('broken.ipynb')
result_path = os.path.join(self.test_dir, 'broken.ipynb')
with self.assertRaises(PapermillExecutionError):
execute_notebook(path, result_path)
nb = read_notebook(result_path)
self.assertEqual(nb.node.cells[0].cell_type, "markdown")
self.assertEqual(nb.node.cells[1].execution_count, 1)
self.assertEqual(nb.node.cells[2].execution_count, 2)
self.assertEqual(nb.node.cells[2].outputs[0].output_type, 'error')
self.assertEqual(nb.node.cells[3].execution_count, None)
def papermill(notebook, output, parameters, raw_parameters, parameters_file,
parameters_yaml, parameters_base64, kernel):
"""Utility for executing a single notebook on a container."""
if parameters_base64:
parameters_final = yaml.load(base64.b64decode(parameters_base64))
elif parameters_yaml:
parameters_final = yaml.load(parameters_yaml)
elif parameters_file:
parameters_final = read_yaml_file(parameters_file)
else:
# Read in Parameters
parameters_final = {}
for name, value in parameters:
parameters_final[name] = _resolve_type(value)
for name, value in raw_parameters:
parameters_final[name] = value
# Notebook Execution
execute_notebook(notebook, output, parameters_final, kernel_name=kernel)
def main(notebook_jupytext_path, params, output_path, help_notebook=False):
"""
:param notebook_jupytext_path:
:param params:
:param output_path:
:param help_notebook:
:return:
"""
# Output name defaults to ipynb_path if the output_path is an empty string
# Convert py to ipynb
surrogate_input_path = create_ipynb_from_py(notebook_jupytext_path)
surrogate_output_path = str(
PurePath(output_path if output_path else notebook_jupytext_path).
with_suffix(IPYNB_SUFFIX))
if help_notebook:
display_notebook_help(surrogate_input_path)
sys.exit(0)
if not validate_notebook_params(surrogate_input_path, params):
display_notebook_help(surrogate_input_path)
sys.exit(1)
try:
# Pass ipynb to papermill
execute_notebook(surrogate_input_path,
surrogate_output_path,
parameters=params)
except nbclient.exceptions.DeadKernelError as e:
# Exiting with a special exit code for dead kernels
LOGGER.error(e)
sys.exit(138)
except PapermillExecutionError as e:
LOGGER.error(e)
create_html_from_ipynb(surrogate_output_path)
else:
create_html_from_ipynb(surrogate_output_path)
def parameterize(nb_input_file_path, nb_output_file_name, parameters, nb):
"""
Adds parameters to the Jupyter Notebook
- Parameters: nb_input_file_path, nb_output_file_name, parameters, nb
- Returns: notebook, boolean
"""
success = False
try:
nb = mill.execute_notebook(
nb_input_file_path,
nb_output_file_name,
parameters=parameters,
prepare_only=True,
log_output=True
)
success = True
except FileNotFoundError as err:
print(err, file=sys.stderr)
finally:
return nb, success
def get_nb(name: str, params: Dict[str, Any]) -> NotebookNode:
nb = execute_notebook(f"./examples/{name}.ipynb", "-", parameters=params)
return nb