diff --git a/binder/requirements.txt b/binder/requirements.txt index 6ebde05a..6b15e076 100644 --- a/binder/requirements.txt +++ b/binder/requirements.txt @@ -3,3 +3,4 @@ pandas matplotlib ipywidgets ipykernel +ansicolors diff --git a/papermill/exceptions.py b/papermill/exceptions.py index 95c2d1ef..1e28819d 100644 --- a/papermill/exceptions.py +++ b/papermill/exceptions.py @@ -1,3 +1,6 @@ +from colors import strip_color + + class AwsError(Exception): """Raised when an AWS Exception is encountered.""" @@ -35,7 +38,7 @@ def __str__(self): # provide the same result as was produced in the past. message = f"\n{75 * '-'}\n" message += f'Exception encountered at "In [{self.exec_count}]":\n' - message += "\n".join(self.traceback) + message += strip_color("\n".join(self.traceback)) message += "\n" return message diff --git a/papermill/tests/test_execute.py b/papermill/tests/test_execute.py index 705d418c..87816cd3 100644 --- a/papermill/tests/test_execute.py +++ b/papermill/tests/test_execute.py @@ -8,6 +8,7 @@ from unittest.mock import ANY, patch import nbformat +from colors import strip_color from nbformat import validate from .. import engines, translators @@ -429,3 +430,34 @@ def test_notebook_node_input(self): 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) diff --git a/requirements.txt b/requirements.txt index 6f8ebb8a..6eb127c2 100644 --- a/requirements.txt +++ b/requirements.txt @@ -7,3 +7,4 @@ requests entrypoints tenacity >= 5.0.2 aiohttp >=3.9.0; python_version=="3.12" +ansicolors