import os import shutil import tempfile import unittest from copy import deepcopy from functools import partial from pathlib import Path from unittest.mock import ANY, patch import nbformat from colors import strip_color from nbformat import validate from papermill import engines, translators from papermill.exceptions import PapermillExecutionError from papermill.execute import execute_notebook from papermill.iorw import load_notebook_node from papermill.log import logger from papermill.utils import chdir from . import get_notebook_path, kernel_name execute_notebook = partial(execute_notebook, kernel_name=kernel_name) class TestNotebookHelpers(unittest.TestCase): def setUp(self): self.test_dir = tempfile.mkdtemp() self.notebook_name = 'simple_execute.ipynb' self.notebook_path = get_notebook_path(self.notebook_name) self.nb_test_executed_fname = os.path.join(self.test_dir, f'output_{self.notebook_name}') def tearDown(self): shutil.rmtree(self.test_dir) @patch(f"{engines.__name__}.PapermillNotebookClient") def test_start_timeout(self, preproc_mock): execute_notebook(self.notebook_path, self.nb_test_executed_fname, start_timeout=123) args, kwargs = preproc_mock.call_args expected = [ ('timeout', None), ('startup_timeout', 123), ('kernel_name', kernel_name), ('log', logger), ] actual = {(key, kwargs[key]) for key in kwargs} self.assertTrue( set(expected).issubset(actual), msg=f'Expected arguments {expected} are not a subset of actual {actual}', ) @patch(f"{engines.__name__}.PapermillNotebookClient") def test_default_start_timeout(self, preproc_mock): execute_notebook(self.notebook_path, self.nb_test_executed_fname) args, kwargs = preproc_mock.call_args expected = [ ('timeout', None), ('startup_timeout', 60), ('kernel_name', kernel_name), ('log', logger), ] actual = {(key, kwargs[key]) for key in kwargs} self.assertTrue( set(expected).issubset(actual), msg=f'Expected arguments {expected} are not a subset of actual {actual}', ) def test_cell_insertion(self): execute_notebook(self.notebook_path, self.nb_test_executed_fname, {'msg': 'Hello'}) test_nb = load_notebook_node(self.nb_test_executed_fname) self.assertListEqual(test_nb.cells[1].get('source').split('\n'), ['# Parameters', 'msg = "Hello"', '']) self.assertEqual(test_nb.metadata.papermill.parameters, {'msg': 'Hello'}) def test_no_tags(self): notebook_name = 'no_parameters.ipynb' nb_test_executed_fname = os.path.join(self.test_dir, f'output_{notebook_name}') execute_notebook(get_notebook_path(notebook_name), nb_test_executed_fname, {'msg': 'Hello'}) test_nb = load_notebook_node(nb_test_executed_fname) self.assertListEqual(test_nb.cells[0].get('source').split('\n'), ['# Parameters', 'msg = "Hello"', '']) self.assertEqual(test_nb.metadata.papermill.parameters, {'msg': 'Hello'}) def test_quoted_params(self): execute_notebook(self.notebook_path, self.nb_test_executed_fname, {'msg': '"Hello"'}) test_nb = load_notebook_node(self.nb_test_executed_fname) self.assertListEqual(test_nb.cells[1].get('source').split('\n'), ['# Parameters', r'msg = "\"Hello\""', '']) self.assertEqual(test_nb.metadata.papermill.parameters, {'msg': '"Hello"'}) def test_backslash_params(self): execute_notebook(self.notebook_path, self.nb_test_executed_fname, {'foo': r'do\ not\ crash'}) test_nb = load_notebook_node(self.nb_test_executed_fname) self.assertListEqual( test_nb.cells[1].get('source').split('\n'), ['# Parameters', r'foo = "do\\ not\\ crash"', ''], ) self.assertEqual(test_nb.metadata.papermill.parameters, {'foo': r'do\ not\ crash'}) def test_backslash_quote_params(self): execute_notebook(self.notebook_path, self.nb_test_executed_fname, {'foo': r'bar=\"baz\"'}) test_nb = load_notebook_node(self.nb_test_executed_fname) self.assertListEqual( test_nb.cells[1].get('source').split('\n'), ['# Parameters', r'foo = "bar=\\\"baz\\\""', ''], ) self.assertEqual(test_nb.metadata.papermill.parameters, {'foo': r'bar=\"baz\"'}) def test_double_backslash_quote_params(self): execute_notebook(self.notebook_path, self.nb_test_executed_fname, {'foo': r'\\"bar\\"'}) test_nb = load_notebook_node(self.nb_test_executed_fname) self.assertListEqual( test_nb.cells[1].get('source').split('\n'), ['# Parameters', r'foo = "\\\\\"bar\\\\\""', ''], ) self.assertEqual(test_nb.metadata.papermill.parameters, {'foo': r'\\"bar\\"'}) def test_prepare_only(self): for example in ['broken1.ipynb', 'keyboard_interrupt.ipynb']: path = get_notebook_path(example) result_path = os.path.join(self.test_dir, example) # Should not raise as we don't execute the notebook at all execute_notebook(path, result_path, {'foo': r'do\ not\ crash'}, prepare_only=True) nb = load_notebook_node(result_path) self.assertEqual(nb.cells[0].cell_type, "code") self.assertEqual( nb.cells[0].get('source').split('\n'), ['# Parameters', r'foo = "do\\ not\\ crash"', ''], ) class TestBrokenNotebook1(unittest.TestCase): def setUp(self): self.test_dir = tempfile.mkdtemp() def tearDown(self): shutil.rmtree(self.test_dir) def test(self): path = get_notebook_path('broken1.ipynb') # check that the notebook has two existing marker cells, so that this test is sure to be # validating the removal logic (the markers are simulatin an error in the first code cell # that has since been fixed) original_nb = load_notebook_node(path) self.assertEqual(original_nb.cells[0].metadata["tags"], ["papermill-error-cell-tag"]) self.assertIn("In [1]", original_nb.cells[0].source) self.assertEqual(original_nb.cells[2].metadata["tags"], ["papermill-error-cell-tag"]) result_path = os.path.join(self.test_dir, 'broken1.ipynb') with self.assertRaises(PapermillExecutionError): execute_notebook(path, result_path) nb = load_notebook_node(result_path) self.assertEqual(nb.cells[0].cell_type, "markdown") self.assertRegex(nb.cells[0].source, r'^$') self.assertEqual(nb.cells[0].metadata["tags"], ["papermill-error-cell-tag"]) self.assertEqual(nb.cells[1].cell_type, "markdown") self.assertEqual(nb.cells[2].execution_count, 1) self.assertEqual(nb.cells[3].cell_type, "markdown") self.assertEqual(nb.cells[4].cell_type, "markdown") self.assertEqual(nb.cells[5].cell_type, "markdown") self.assertRegex(nb.cells[5].source, '') self.assertEqual(nb.cells[5].metadata["tags"], ["papermill-error-cell-tag"]) self.assertEqual(nb.cells[6].execution_count, 2) self.assertEqual(nb.cells[6].outputs[0].output_type, 'error') self.assertEqual(nb.cells[7].execution_count, None) # double check the removal (the new cells above should be the only two tagged ones) self.assertEqual(sum("papermill-error-cell-tag" in cell.metadata.get("tags", []) for cell in nb.cells), 2) class TestBrokenNotebook2(unittest.TestCase): def setUp(self): self.test_dir = tempfile.mkdtemp() def tearDown(self): shutil.rmtree(self.test_dir) def test(self): path = get_notebook_path('broken2.ipynb') result_path = os.path.join(self.test_dir, 'broken2.ipynb') with self.assertRaises(PapermillExecutionError): execute_notebook(path, result_path) nb = load_notebook_node(result_path) self.assertEqual(nb.cells[0].cell_type, "markdown") self.assertRegex(nb.cells[0].source, r'^.*In \[2\].*$') self.assertEqual(nb.cells[1].execution_count, 1) self.assertEqual(nb.cells[2].cell_type, "markdown") self.assertRegex(nb.cells[2].source, '') self.assertEqual(nb.cells[3].execution_count, 2) self.assertEqual(nb.cells[3].outputs[0].output_type, 'display_data') self.assertEqual(nb.cells[3].outputs[1].output_type, 'error') self.assertEqual(nb.cells[4].execution_count, None) class TestReportMode(unittest.TestCase): def setUp(self): self.test_dir = tempfile.mkdtemp() self.notebook_name = 'report_mode_test.ipynb' self.notebook_path = get_notebook_path(self.notebook_name) self.nb_test_executed_fname = os.path.join(self.test_dir, f'output_{self.notebook_name}') def tearDown(self): shutil.rmtree(self.test_dir) def test_report_mode(self): nb = execute_notebook(self.notebook_path, self.nb_test_executed_fname, {'a': 0}, report_mode=True) for cell in nb.cells: if cell.cell_type == 'code': self.assertEqual(cell.metadata.get('jupyter', {}).get('source_hidden'), True) class TestOutputPathNone(unittest.TestCase): def test_output_path_of_none(self): """Output path of None should return notebook node obj but not write an ipynb""" nb = execute_notebook(get_notebook_path('simple_execute.ipynb'), None, {'msg': 'Hello'}) self.assertEqual(nb.metadata.papermill.parameters, {'msg': 'Hello'}) class TestCWD(unittest.TestCase): def setUp(self): self.test_dir = tempfile.mkdtemp() self.base_test_dir = tempfile.mkdtemp() self.check_notebook_name = 'read_check.ipynb' self.check_notebook_path = os.path.join(self.base_test_dir, 'read_check.ipynb') # Setup read paths so base_test_dir has check_notebook_name shutil.copyfile(get_notebook_path(self.check_notebook_name), self.check_notebook_path) with open(os.path.join(self.test_dir, 'check.txt'), 'w', encoding='utf-8') as f: # Needed for read_check to pass f.write('exists') self.simple_notebook_name = 'simple_execute.ipynb' self.simple_notebook_path = os.path.join(self.base_test_dir, 'simple_execute.ipynb') # Setup read paths so base_test_dir has simple_notebook_name shutil.copyfile(get_notebook_path(self.simple_notebook_name), self.simple_notebook_path) self.nb_test_executed_fname = 'test_output.ipynb' def tearDown(self): shutil.rmtree(self.test_dir) shutil.rmtree(self.base_test_dir) def test_local_save_ignores_cwd_assignment(self): with chdir(self.base_test_dir): # Both paths are relative execute_notebook(self.simple_notebook_name, self.nb_test_executed_fname, cwd=self.test_dir) self.assertTrue(os.path.isfile(os.path.join(self.base_test_dir, self.nb_test_executed_fname))) def test_execution_respects_cwd_assignment(self): with chdir(self.base_test_dir): # Both paths are relative execute_notebook(self.check_notebook_name, self.nb_test_executed_fname, cwd=self.test_dir) self.assertTrue(os.path.isfile(os.path.join(self.base_test_dir, self.nb_test_executed_fname))) def test_pathlib_paths(self): # Copy of test_execution_respects_cwd_assignment but with `Path`s with chdir(self.base_test_dir): execute_notebook( Path(self.check_notebook_name), Path(self.nb_test_executed_fname), cwd=Path(self.test_dir), ) self.assertTrue(Path(self.base_test_dir).joinpath(self.nb_test_executed_fname).exists()) class TestSysExit(unittest.TestCase): def setUp(self): self.test_dir = tempfile.mkdtemp() def tearDown(self): shutil.rmtree(self.test_dir) def test_sys_exit(self): notebook_name = 'sysexit.ipynb' result_path = os.path.join(self.test_dir, f'output_{notebook_name}') execute_notebook(get_notebook_path(notebook_name), result_path) nb = load_notebook_node(result_path) self.assertEqual(nb.cells[0].cell_type, "code") self.assertEqual(nb.cells[0].execution_count, 1) self.assertEqual(nb.cells[1].execution_count, 2) self.assertEqual(nb.cells[1].outputs[0].output_type, 'error') self.assertEqual(nb.cells[1].outputs[0].ename, 'SystemExit') self.assertEqual(nb.cells[1].outputs[0].evalue, '') self.assertEqual(nb.cells[2].execution_count, None) def test_sys_exit0(self): notebook_name = 'sysexit0.ipynb' result_path = os.path.join(self.test_dir, f'output_{notebook_name}') execute_notebook(get_notebook_path(notebook_name), result_path) nb = load_notebook_node(result_path) self.assertEqual(nb.cells[0].cell_type, "code") self.assertEqual(nb.cells[0].execution_count, 1) self.assertEqual(nb.cells[1].execution_count, 2) self.assertEqual(nb.cells[1].outputs[0].output_type, 'error') self.assertEqual(nb.cells[1].outputs[0].ename, 'SystemExit') self.assertEqual(nb.cells[1].outputs[0].evalue, '0') self.assertEqual(nb.cells[2].execution_count, None) def test_sys_exit1(self): notebook_name = 'sysexit1.ipynb' result_path = os.path.join(self.test_dir, f'output_{notebook_name}') with self.assertRaises(PapermillExecutionError): execute_notebook(get_notebook_path(notebook_name), result_path) nb = load_notebook_node(result_path) self.assertEqual(nb.cells[0].cell_type, "markdown") self.assertRegex(nb.cells[0].source, r'^$') self.assertEqual(nb.cells[1].execution_count, 1) self.assertEqual(nb.cells[2].cell_type, "markdown") self.assertRegex(nb.cells[2].source, '') self.assertEqual(nb.cells[3].execution_count, 2) self.assertEqual(nb.cells[3].outputs[0].output_type, 'error') self.assertEqual(nb.cells[4].execution_count, None) def test_system_exit(self): notebook_name = 'systemexit.ipynb' result_path = os.path.join(self.test_dir, f'output_{notebook_name}') execute_notebook(get_notebook_path(notebook_name), result_path) nb = load_notebook_node(result_path) self.assertEqual(nb.cells[0].cell_type, "code") self.assertEqual(nb.cells[0].execution_count, 1) self.assertEqual(nb.cells[1].execution_count, 2) self.assertEqual(nb.cells[1].outputs[0].output_type, 'error') self.assertEqual(nb.cells[1].outputs[0].ename, 'SystemExit') self.assertEqual(nb.cells[1].outputs[0].evalue, '') self.assertEqual(nb.cells[2].execution_count, None) def test_line_magic_error(self): notebook_name = 'line_magic_error.ipynb' result_path = os.path.join(self.test_dir, f'output_{notebook_name}') with self.assertRaises(PapermillExecutionError): execute_notebook(get_notebook_path(notebook_name), result_path) nb = load_notebook_node(result_path) self.assertEqual(nb.cells[0].cell_type, "markdown") self.assertRegex(nb.cells[0].source, r'^$') self.assertEqual(nb.cells[0].metadata["tags"], ["papermill-error-cell-tag"]) self.assertEqual(nb.cells[2].cell_type, "code") self.assertEqual(nb.cells[2].execution_count, 1) self.assertEqual(nb.cells[3].execution_count, None) class TestNotebookValidation(unittest.TestCase): def setUp(self): self.test_dir = tempfile.mkdtemp() def tearDown(self): shutil.rmtree(self.test_dir) def test_from_version_4_4_upgrades(self): notebook_name = 'nb_version_4.4.ipynb' result_path = os.path.join(self.test_dir, f'output_{notebook_name}') execute_notebook(get_notebook_path(notebook_name), result_path, {'var': 'It works'}) nb = load_notebook_node(result_path) validate(nb) class TestMinimalNotebook(unittest.TestCase): def setUp(self): self.test_dir = tempfile.mkdtemp() def tearDown(self): shutil.rmtree(self.test_dir) def test_no_v3_language_backport(self): notebook_name = 'blank-vscode.ipynb' result_path = os.path.join(self.test_dir, f'output_{notebook_name}') execute_notebook(get_notebook_path(notebook_name), result_path, {'var': 'It works'}) nb = load_notebook_node(result_path) validate(nb) class TestExecuteWithCustomEngine(unittest.TestCase): class CustomEngine(engines.Engine): @classmethod def execute_managed_notebook(cls, nb_man, kernel_name, **kwargs): pass @classmethod def nb_kernel_name(cls, nb, name=None): return "my_custom_kernel" @classmethod def nb_language(cls, nb, language=None): return "my_custom_language" def setUp(self): self.test_dir = tempfile.mkdtemp() self.notebook_path = get_notebook_path('simple_execute.ipynb') self.nb_test_executed_fname = os.path.join(self.test_dir, 'output_simple_execute.ipynb') self._orig_papermill_engines = deepcopy(engines.papermill_engines) self._orig_translators = deepcopy(translators.papermill_translators) engines.papermill_engines.register("custom_engine", self.CustomEngine) translators.papermill_translators.register("my_custom_language", translators.PythonTranslator()) def tearDown(self): shutil.rmtree(self.test_dir) engines.papermill_engines = self._orig_papermill_engines translators.papermill_translators = self._orig_translators @patch.object(CustomEngine, "execute_managed_notebook", wraps=CustomEngine.execute_managed_notebook) @patch("papermill.parameterize.translate_parameters", wraps=translators.translate_parameters) def test_custom_kernel_name_and_language(self, translate_parameters, execute_managed_notebook): """Tests execute against engine with custom implementations to fetch kernel name and language from the notebook object """ execute_notebook( self.notebook_path, self.nb_test_executed_fname, engine_name="custom_engine", parameters={"msg": "fake msg"}, ) self.assertEqual(execute_managed_notebook.call_args[0], (ANY, "my_custom_kernel")) self.assertEqual(translate_parameters.call_args[0], (ANY, 'my_custom_language', {"msg": "fake msg"}, ANY)) class TestNotebookNodeInput(unittest.TestCase): def setUp(self): self.test_dir = tempfile.TemporaryDirectory() self.result_path = os.path.join(self.test_dir.name, 'output.ipynb') def tearDown(self): self.test_dir.cleanup() def test_notebook_node_input(self): input_nb = nbformat.read(get_notebook_path('simple_execute.ipynb'), as_version=4) execute_notebook(input_nb, self.result_path, {'msg': 'Hello'}) test_nb = nbformat.read(self.result_path, as_version=4) self.assertEqual(test_nb.metadata.papermill.parameters, {'msg': 'Hello'}) class TestOutputFormatting(unittest.TestCase): def setUp(self): self.test_dir = tempfile.mkdtemp() def tearDown(self): shutil.rmtree(self.test_dir) def test_output_formatting(self): notebook_name = 'sysexit1.ipynb' result_path = os.path.join(self.test_dir, f'output_{notebook_name}') try: execute_notebook(get_notebook_path(notebook_name), result_path) # exception should be thrown by now self.assertFalse(True) except PapermillExecutionError as ex: self.assertEqual(ex.traceback[1], "\x1b[0;31mSystemExit\x1b[0m\x1b[0;31m:\x1b[0m 1\n") self.assertEqual(strip_color(ex.traceback[1]), "SystemExit: 1\n") nb = load_notebook_node(result_path) self.assertEqual(nb.cells[0].cell_type, "markdown") self.assertRegex(nb.cells[0].source, r'^$') self.assertEqual(nb.cells[1].execution_count, 1) self.assertEqual(nb.cells[2].cell_type, "markdown") self.assertRegex(nb.cells[2].source, '') self.assertEqual(nb.cells[3].execution_count, 2) self.assertEqual(nb.cells[3].outputs[0].output_type, 'error') self.assertEqual(nb.cells[4].execution_count, None)