bring back strip_color and remove ANSI color codes from exception traceback

papermill/exceptions.py

@@ -1,3 +1,6 @@
+import re
+
+
 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
 
@@ -75,3 +78,20 @@ def missing_dep():
         )
 
     return missing_dep
+
+
+# copied from https://github.com/jonathaneunice/ansiwrap/blob/master/ansiwrap/core.py
+# papermill used to use strip_color from ansiwrap package, but removed due to dependency
+# resolution issues (ref: https://github.com/nteract/papermill/pull/681)
+def strip_color(s):
+    """
+    Remove ANSI color/style sequences from a string. The set of all
+    possibly ANSI sequences is large, so does not try to strip every
+    possible one. But does strip some outliers seen not just in text
+    generated by this module, but by other ANSI colorizers in the wild.
+    Those include `\x1b[K` (aka EL or erase to end of line) and `\x1b[m`
+    a terse version of the more common `\x1b[0m`.
+    """
+    ANSIRE = re.compile('\x1b\\[(K|.*?m)')
+
+    return ANSIRE.sub('', s)

papermill/tests/test_execute.py

@@ -11,7 +11,7 @@
 from nbformat import validate
 
 from .. import engines, translators
-from ..exceptions import PapermillExecutionError
+from ..exceptions import PapermillExecutionError, strip_color
 from ..execute import execute_notebook
 from ..iorw import load_notebook_node
 from ..log import logger
@@ -429,3 +429,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)