bring back strip_color and remove ANSI color codes from exception tra…

…ceback (#791)

* bring back strip_color and remove ANSI color codes from exception traceback

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* use strip_color from ansicolors package

* fix binder build, import strip_color directly

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

---------

Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>

binder/requirements.txt

@@ -3,3 +3,4 @@ pandas
 matplotlib
 ipywidgets
 ipykernel
+ansicolors

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
 

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'^<span .*<a href="#papermill-error-cell".*In \[2\].*</span>$')
+        self.assertEqual(nb.cells[1].execution_count, 1)
+
+        self.assertEqual(nb.cells[2].cell_type, "markdown")
+        self.assertRegex(nb.cells[2].source, '<span id="papermill-error-cell" .*</span>')
+        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)

requirements.txt

@@ -7,3 +7,4 @@ requests
 entrypoints
 tenacity >= 5.0.2
 aiohttp >=3.9.0; python_version=="3.12"
+ansicolors