From 0c01e2b7f8025fd60e4c9d2b37b01bfa68e69775 Mon Sep 17 00:00:00 2001
From: Mehmet Bektas <mbektasgh@outlook.com>
Date: Wed, 3 Apr 2024 11:42:12 -0700
Subject: [PATCH] bring back strip_color and remove ANSI color codes from
 exception traceback

---
 papermill/exceptions.py         | 21 ++++++++++++++++++++-
 papermill/tests/test_execute.py | 32 +++++++++++++++++++++++++++++++-
 2 files changed, 51 insertions(+), 2 deletions(-)

diff --git a/papermill/exceptions.py b/papermill/exceptions.py
index 95c2d1ef..b61eeb8b 100644
--- a/papermill/exceptions.py
+++ b/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,19 @@ 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)
diff --git a/papermill/tests/test_execute.py b/papermill/tests/test_execute.py
index 705d418c..45b12454 100644
--- a/papermill/tests/test_execute.py
+++ b/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,33 @@ 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)