diff --git a/tutorial/tests/testsuite/helpers.py b/tutorial/tests/testsuite/helpers.py index 21607bd2..b30d05e7 100644 --- 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:
-
{html.escape(self.stdout) if self.stdout else 'No output'}
+
{html.escape(strip_ansi_codes(self.stdout)) if self.stdout else 'No output'}
-
{html.escape(self.stderr) if self.stderr else 'No errors'}
+
{html.escape(strip_ansi_codes(self.stderr)) if self.stderr else 'No errors'}
@@ -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( - "
" + '
' f'

{title}' '' f"solution_{function.name}

" diff --git a/tutorial/tests/testsuite/testsuite.py b/tutorial/tests/testsuite/testsuite.py index ad289383..17bdcfa7 100644 --- 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):