Fix traceback issue (#265)

Also: * Remove ANSI escape sequences from tests output * UI fix: hide scrollbar from test results title
empa-scientific-it · Nov 26, 2024 · 7bce206 · 7bce206
1 parent 6d6452e
commit 7bce206
Show file tree

Hide file tree

Showing 2 changed files with 65 additions and 50 deletions.
diff --git a/tutorial/tests/testsuite/helpers.py b/tutorial/tests/testsuite/helpers.py
@@ -1,4 +1,5 @@
 import html
+import re
 from dataclasses import dataclass
 from enum import Enum
 from pathlib import Path
@@ -14,6 +15,12 @@
 from .ai_helpers import AIExplanation, OpenAIWrapper
 
 
+def strip_ansi_codes(text: str) -> str:
+    """Remove ANSI escape sequences from text"""
+    ansi_escape = re.compile(r"\x1B(?:[@-Z\\-_]|\[[0-?]*[ -/]*[@-~])")
+    return ansi_escape.sub("", text)
+
+
 class TestOutcome(Enum):
     PASS = 1
     FAIL = 2
@@ -293,7 +300,7 @@ def to_html(self) -> str:
         # Exception information if test failed
         if self.exception is not None:
             exception_type = type(self.exception).__name__
-            exception_message = str(self.exception)
+            exception_message = strip_ansi_codes(str(self.exception))
 
             html_parts.append(
                 f"""
@@ -329,10 +336,10 @@ def to_html(self) -> str:
                             </div>
                             <div class="output-content">
                                 <div id="{tab_id}_output" class="output-pane active">
-                                    <pre>{html.escape(self.stdout) if self.stdout else 'No output'}</pre>
+                                    <pre>{html.escape(strip_ansi_codes(self.stdout)) if self.stdout else 'No output'}</pre>
                                 </div>
                                 <div id="{tab_id}_error" class="output-pane">
-                                    <pre>{html.escape(self.stderr) if self.stderr else 'No errors'}</pre>
+                                    <pre>{html.escape(strip_ansi_codes(self.stderr)) if self.stderr else 'No errors'}</pre>
                                 </div>
                             </div>
                         </div>
@@ -605,7 +612,7 @@ def prepare_output_cell(self) -> ipywidgets.Output:
         title = "Test Results for " if function else "Test Results "
         output_cell.append_display_data(
             HTML(
-                "<div>"
+                '<div style="overflow: hidden;">'
                 f'<h2 style="font-size: 1.5rem; margin: 0;">{title}'
                 '<code style="font-size: 1.1rem; background: #f3f4f6; padding: 0.25rem 0.5rem; border-radius: 0.25rem; font-family: ui-monospace, monospace;">'
                 f"solution_{function.name}</code></h2>"

diff --git a/tutorial/tests/testsuite/testsuite.py b/tutorial/tests/testsuite/testsuite.py
@@ -7,7 +7,7 @@
 import os
 import pathlib
 from collections import defaultdict
-from contextlib import redirect_stderr, redirect_stdout
+from contextlib import contextmanager, redirect_stderr, redirect_stdout
 from queue import Queue
 from threading import Thread
 from typing import Dict, List, Optional
@@ -248,6 +248,17 @@ def run_cell(self) -> List[IPytestResult]:
 
         return test_results
 
+    @contextmanager
+    def traceback_handling(self, debug: bool):
+        """Context manager to temporarily modify traceback behavior"""
+        original_traceback = self.shell._showtraceback
+        try:
+            if not debug:
+                self.shell._showtraceback = lambda *args, **kwargs: None
+            yield
+        finally:
+            self.shell._showtraceback = original_traceback
+
     @cell_magic
     def ipytest(self, line: str, cell: str):
         """The `%%ipytest` cell magic"""
@@ -270,56 +281,53 @@ def ipytest(self, line: str, cell: str):
             self.threaded = True
             self.test_queue = Queue()
 
-        # If debug is in the line, then we want to show the traceback
-        if self.debug:
-            self.shell._showtraceback = self._orig_traceback
-        else:
-            self.shell._showtraceback = lambda *args, **kwargs: None
-
-        # Get the module containing the test(s)
-        if (
-            module_name := get_module_name(
-                " ".join(line_contents), self.shell.user_global_ns
-            )
-        ) is None:
-            raise TestModuleNotFoundError
+        with self.traceback_handling(self.debug):
+            # Get the module containing the test(s)
+            if (
+                module_name := get_module_name(
+                    " ".join(line_contents), self.shell.user_global_ns
+                )
+            ) is None:
+                raise TestModuleNotFoundError
 
-        self.module_name = module_name
+            self.module_name = module_name
 
-        # Check that the test module file exists
-        if not (
-            module_file := pathlib.Path(f"tutorial/tests/test_{self.module_name}.py")
-        ).exists():
-            raise FileNotFoundError(module_file)
+            # Check that the test module file exists
+            if not (
+                module_file := pathlib.Path(
+                    f"tutorial/tests/test_{self.module_name}.py"
+                )
+            ).exists():
+                raise FileNotFoundError(module_file)
 
-        self.module_file = module_file
+            self.module_file = module_file
 
-        # Run the cell
-        results = self.run_cell()
+            # Run the cell
+            results = self.run_cell()
 
-        # If in debug mode, display debug information first
-        if self.debug:
-            debug_output = DebugOutput(
-                module_name=self.module_name,
-                module_file=self.module_file,
-                results=results,
-            )
-            display(HTML(debug_output.to_html()))
-
-        # Parse the AST of the test module to retrieve the solution code
-        ast_parser = AstParser(self.module_file)
-        # Display the test results and the solution code
-        for result in results:
-            solution = (
-                ast_parser.get_solution_code(result.function.name)
-                if result.function and result.function.name
-                else None
-            )
-            TestResultOutput(
-                result,
-                solution,
-                self.shell.openai_client,  # type: ignore
-            ).display_results()
+            # If in debug mode, display debug information first
+            if self.debug:
+                debug_output = DebugOutput(
+                    module_name=self.module_name,
+                    module_file=self.module_file,
+                    results=results,
+                )
+                display(HTML(debug_output.to_html()))
+
+            # Parse the AST of the test module to retrieve the solution code
+            ast_parser = AstParser(self.module_file)
+            # Display the test results and the solution code
+            for result in results:
+                solution = (
+                    ast_parser.get_solution_code(result.function.name)
+                    if result.function and result.function.name
+                    else None
+                )
+                TestResultOutput(
+                    result,
+                    solution,
+                    self.shell.openai_client,  # type: ignore
+                ).display_results()
 
 
 def load_ipython_extension(ipython):