diff --git a/src/xdoctest/doctest_example.py b/src/xdoctest/doctest_example.py index 6ad17296..b1b438e6 100644 --- a/src/xdoctest/doctest_example.py +++ b/src/xdoctest/doctest_example.py @@ -93,6 +93,8 @@ def _populate_from_cli(self, ns): 'global_exec': ns['global_exec'], 'supress_import_errors': ns['supress_import_errors'], 'verbose': ns['verbose'], + 'write_outputs': ns['write_outputs'], + 'fill_missing_wants': ns['fill_missing_wants'], } return _examp_conf diff --git a/src/xdoctest/plugin.py b/src/xdoctest/plugin.py index e34cf7b0..e6916432 100644 --- a/src/xdoctest/plugin.py +++ b/src/xdoctest/plugin.py @@ -160,6 +160,16 @@ def str_lower(x): dest='xdoctest_analysis', ) + group.addoption('--xdoctest-write-outputs', '--xdoc-write-outputs', + action='store_true', default=False, + help='Write captured outputs back to source files', + dest='xdoctest_write_outputs') + + group.addoption('--xdoctest-fill-missing-wants', '--xdoc-fill-missing-wants', + action='store_true', default=False, + help='When used with --xdoctest-write-outputs, also add missing want statements', + dest='xdoctest_fill_missing_wants') + from xdoctest import doctest_example doctest_example.DoctestConfig()._update_argparse_cli( diff --git a/src/xdoctest/runner.py b/src/xdoctest/runner.py index 96634167..06c1fe0f 100644 --- a/src/xdoctest/runner.py +++ b/src/xdoctest/runner.py @@ -52,6 +52,7 @@ import sys import time +import pathlib import types import typing import warnings @@ -389,16 +390,66 @@ def doctest_module( AFTER_ALL_HOOKS = [] if insert_skip_directive_above_failures: AFTER_ALL_HOOKS.append(_auto_disable_failing_tests_hook) + + if config.get('write_outputs', False): + AFTER_ALL_HOOKS.append(_write_outputs_hook) + for hook in AFTER_ALL_HOOKS: context = { - 'enabled_example': enabled_examples, + 'enabled_examples': enabled_examples, 'run_summary': run_summary, + 'config': config, } hook(context) return run_summary +def _write_outputs_hook(context): + """ + Write captured outputs back to source files. + Updates existing want statements and optionally fills missing ones. + + Args: + context (dict): Hook context with keys: + - enabled_examples: List of DocTest objects that ran + - run_summary: Dictionary with test results + - config: Configuration dict with write_outputs and fill_missing_wants + """ + from collections import defaultdict + + enabled_examples = context['enabled_examples'] + config = context.get('config', {}) + fill_missing = config.get('fill_missing_wants', False) + + # Group modifications by file + file_modifications = defaultdict(list) + + for example in enabled_examples: + # Skip if test didn't run + if not example.anything_ran(): + continue + + # Skip if test failed with an exception other than GotWantException + # (GotWantException means output mismatch, which is what we want to fix) + if example.exc_info is not None: + from xdoctest.checker import GotWantException + exc_type, exc_value, exc_tb = example.exc_info + # Skip if it's not a GotWantException (e.g., syntax error, runtime error) + if not isinstance(exc_value, GotWantException): + continue + + # Compute modifications for this example + for partx, part in enumerate(example._parts): + mod = _compute_part_modification(example, partx, part, fill_missing) + if mod: + file_modifications[example.fpath].append(mod) + + # Apply modifications to each file + for fpath, modifications in file_modifications.items(): + _apply_modifications_to_file(fpath, modifications) + + def _auto_disable_failing_tests_hook(context): """ Experimental feature to modify code based on failing tests. @@ -443,6 +494,187 @@ def _auto_disable_failing_tests_hook(context): file.write(''.join(lines)) +def _strip_ansi_codes(text): + """ + Remove ANSI color codes from text. + + Args: + text (str): Text that may contain ANSI escape sequences + + Returns: + str: Text with ANSI codes removed + + Example: + >>> # Test with ANSI color codes + >>> text = '\x1b[31mred text\x1b[0m' + >>> _strip_ansi_codes(text) + 'red text' + >>> # Test with no ANSI codes + >>> _strip_ansi_codes('plain text') + 'plain text' + """ + import re + ansi_escape = re.compile(r'\x1B(?:[@-Z\\-_]|\[[0-?]*[ -/]*[@-~])') + return ansi_escape.sub('', text) + + +def _detect_indentation(example, part): + """ + Detect indentation from the source file. + + Args: + example (DocTest): The doctest example containing the part + part (DoctestPart): The doctest part to analyze + + Returns: + str: The indentation prefix (spaces/tabs) + """ + # Calculate the absolute line number in the source file + # where this part's exec line appears + exec_line_number = example.lineno + part.line_offset + + lines = pathlib.Path(example.fpath).read_text().split('\n') + if exec_line_number - 1 < len(lines): + source_line = lines[exec_line_number - 1] + # Extract the leading whitespace + indentation = source_line[:len(source_line) - len(source_line.lstrip())] + return indentation + + return '' + + +def _format_want_lines(output_text, indentation): + """ + Format output as doctest want lines with proper indentation. + + Args: + output_text (str): The raw output to format + indentation (str): The indentation prefix to apply + + Returns: + List[str]: Formatted lines ready to write to file + + Example: + >>> # Test basic formatting + >>> lines = _format_want_lines('hello\\nworld', ' ') + >>> len(lines) + 2 + >>> lines[0] + ' hello\\n' + >>> lines[1] + ' world\\n' + >>> # Test with blank lines + >>> lines = _format_want_lines('line1\\n\\nline3', ' ') + >>> len(lines) + 3 + >>> lines[1] + ' \\n' + """ + if not output_text: + return [] + + lines = _strip_ansi_codes(output_text).rstrip('\n').split('\n') + + formatted = [] + for line in lines: + if line.strip() == '': + line = '' + formatted.append(indentation + line + '\n') + + return formatted + + +def _compute_part_modification(example, partx, part, fill_missing): + """ + Determine what modification is needed for a doctest part. + + Args: + example (DocTest): The doctest example containing the part + partx (int): Index of the part in the example + part (DoctestPart): The doctest part to process + fill_missing (bool): Whether to fill in missing want statements + + Returns: + dict or None: Modification dict with keys: + - line_number: Line number where want is/should be + - action: 'replace' or 'insert' + - num_old_lines: Number of lines to remove (for replace) + - new_lines: List of formatted lines to insert + Returns None if no modification needed + """ + from xdoctest import constants + + got_stdout = example.logged_stdout.get(partx, '') + got_eval = example.logged_evals.get(partx, constants.NOT_EVALED) + + if got_stdout or got_eval is constants.NOT_EVALED: + output_text = got_stdout + elif got_eval is not None: + output_text = repr(got_eval) + else: + output_text = '' + + if part.want_lines: + action = 'replace' + num_old_lines = len(part.want_lines) + elif fill_missing and output_text: + action = 'insert' + num_old_lines = 0 + else: + return None # Skip this part + + want_line_number = example.lineno + part.line_offset + part.n_exec_lines + + formatted_lines = _format_want_lines(output_text, _detect_indentation(example, part)) + + if not formatted_lines: + return None + + return { + 'line_number': want_line_number, + 'action': action, + 'num_old_lines': num_old_lines, + 'new_lines': formatted_lines, + } + + +def _apply_modifications_to_file(fpath, modifications): + """ + Apply all modifications to a single file. + + Args: + fpath (str): Path to source file + modifications (list): List of modification dicts from _compute_part_modification + """ + print(f'Updating doctests in: {fpath}') + + with open(fpath, 'r') as f: + lines = f.readlines() + + # Sort by line number (descending) to preserve line numbers + modifications = sorted(modifications, + key=lambda m: m['line_number'], + reverse=True) + + for mod in modifications: + line_idx = mod['line_number'] - 1 + + if mod['action'] == 'replace': + # Remove old want lines + for _ in range(mod['num_old_lines']): + if line_idx < len(lines): + lines.pop(line_idx) + + # Insert new want lines + for new_line in reversed(mod['new_lines']): + lines.insert(line_idx, new_line) + + with open(fpath, 'w') as f: + f.writelines(lines) + + print(' Updated {} doctest part(s)'.format(len(modifications))) + + def _convert_to_test_module(enabled_examples): """ Logic for the "dumps" command. @@ -906,6 +1138,20 @@ def _update_argparse_cli(add_argument, prefix=None): help=('Same as if durations=0'), ) + add_argument(*('--write-outputs', '--update-wants'), + dest='write_outputs', + action='store_true', + default=False, + help=('Write captured outputs back to source files, ' + 'updating existing want statements')) + + add_argument(*('--fill-missing-wants',), + dest='fill_missing_wants', + action='store_true', + default=False, + help=('When used with --write-outputs, also add want ' + 'statements for parts that have none')) + add_argument_kws = [ # (['--style'], dict(dest='style', # type=str, help='choose your style', diff --git a/tests/test_write_outputs.py b/tests/test_write_outputs.py new file mode 100644 index 00000000..30d255e0 --- /dev/null +++ b/tests/test_write_outputs.py @@ -0,0 +1,451 @@ +""" +Unit tests for the --write-outputs and --fill-missing-wants feature. +""" + +import xdoctest +import textwrap +import pytest + + +@pytest.fixture +def create_test_file(tmp_path): + """ + Fixture that returns a function to create test files. + """ + + def _create_file(content): + """Create a test Python file with the given content.""" + test_file = tmp_path / "test_doctest.py" + test_file.write_text(textwrap.dedent(content)) + return test_file + + return _create_file + + +def test_write_outputs_basic(create_test_file): + """ + Test basic write-outputs functionality. + """ + + test_file = create_test_file(''' + def example(): + """ + Example: + >>> print('hello') + wrong output + """ + pass + ''') + + xdoctest.doctest_module( + str(test_file), command="all", config={"write_outputs": True} + ) + + assert test_file.read_text() == textwrap.dedent(''' + def example(): + """ + Example: + >>> print('hello') + hello + """ + pass + ''') + + +def test_write_outputs_eval(create_test_file): + """ + Test write-outputs with eval expressions. + """ + + test_file = create_test_file(''' + def example(): + """ + Example: + >>> 2 + 2 + 5 + """ + pass + ''') + + xdoctest.doctest_module( + str(test_file), command="all", config={"write_outputs": True} + ) + + assert test_file.read_text() == textwrap.dedent(''' + def example(): + """ + Example: + >>> 2 + 2 + 4 + """ + pass + ''') + + +def test_write_outputs_multiline(create_test_file): + """ + Test write-outputs with multiline output. + """ + + test_file = create_test_file(''' + def example(): + """ + Example: + >>> for i in range(3): + ... print(i) + wrong + """ + pass + ''') + + xdoctest.doctest_module( + str(test_file), command="all", config={"write_outputs": True} + ) + + assert test_file.read_text() == textwrap.dedent(''' + def example(): + """ + Example: + >>> for i in range(3): + ... print(i) + 0 + 1 + 2 + """ + pass + ''') + + +def test_write_outputs_indentation(create_test_file): + """ + Test that indentation is preserved correctly. + """ + + test_file = create_test_file(''' + class MyClass: + def method(self): + """ + Example: + >>> print('test') + wrong + """ + pass + ''') + + xdoctest.doctest_module( + str(test_file), command="all", config={"write_outputs": True} + ) + + assert test_file.read_text() == textwrap.dedent(''' + class MyClass: + def method(self): + """ + Example: + >>> print('test') + test + """ + pass + ''') + + +def test_fill_missing_wants(create_test_file): + """ + Test --fill-missing-wants flag. + """ + + test_content = ''' + def example(): + """ + Example: + >>> print('hello') + """ + pass + ''' + test_file = create_test_file(test_content) + + # Run without --fill-missing-wants + xdoctest.doctest_module( + str(test_file), + command="all", + config={"write_outputs": True}, + ) + + assert test_file.read_text() == textwrap.dedent(test_content) + + # Now run with --fill-missing-wants + xdoctest.doctest_module( + str(test_file), + command="all", + config={"write_outputs": True, "fill_missing_wants": True}, + ) + + assert test_file.read_text() == textwrap.dedent(''' + def example(): + """ + Example: + >>> print('hello') + hello + """ + pass + ''') + + +def test_write_outputs_skip_failed(create_test_file): + """ + Test that tests with actual exceptions are skipped. + """ + + test_content = ''' + def example(): + """ + Example: + >>> 1 / 0 + should not be written + """ + pass + ''' + test_file = create_test_file(test_content) + + xdoctest.doctest_module( + str(test_file), command="all", config={"write_outputs": True} + ) + + assert test_file.read_text() == textwrap.dedent(test_content) + + + + +def test_write_outputs_preserves_blanklines(create_test_file): + """ + Test that blank lines are converted to markers. + """ + test_file = create_test_file(''' + def example(): + """ + Example: + >>> for i in [1, '', 3]: + ... print(i) + wrong + """ + pass + ''') + + xdoctest.doctest_module( + str(test_file), command="all", config={"write_outputs": True} + ) + + assert test_file.read_text() == textwrap.dedent(''' + def example(): + """ + Example: + >>> for i in [1, '', 3]: + ... print(i) + 1 + + 3 + """ + pass + ''') + + +def test_write_outputs_multiple_parts(create_test_file): + """ + Test updating parts in a single doctest. + + When a part fails, execution stops, so only the first failing part is updated. + """ + test_file = create_test_file(''' + def example(): + """ + Example: + >>> x = 5 + >>> print(x + 1) + 0 + >>> print(x * 2) + wrong + """ + pass + ''') + + xdoctest.doctest_module( + str(test_file), command="all", config={"write_outputs": True} + ) + + # The second doctest was not executed because the first one failed, + # so the output is still 'wrong' + # This is expected behavior - run --write-outputs again to update next parts + assert test_file.read_text() == textwrap.dedent(''' + def example(): + """ + Example: + >>> x = 5 + >>> print(x + 1) + 6 + >>> print(x * 2) + wrong + """ + pass + ''') + + xdoctest.doctest_module( + str(test_file), command="all", config={"write_outputs": True} + ) + + assert test_file.read_text() == textwrap.dedent(''' + def example(): + """ + Example: + >>> x = 5 + >>> print(x + 1) + 6 + >>> print(x * 2) + 10 + """ + pass + ''') + + +def test_write_outputs_no_modification_on_success(create_test_file): + """ + Test that already-correct outputs are not modified. + """ + + test_content = ''' + def example(): + """ + Example: + >>> 2 + 2 + 4 + """ + pass + ''' + test_file = create_test_file(test_content) + + xdoctest.doctest_module( + str(test_file), command="all", config={"write_outputs": True} + ) + + assert test_file.read_text() == textwrap.dedent(test_content) + + +@pytest.mark.parametrize( + "test_content, expected_spaces", + [ + ( + ''' + def function(): + """ + Example: + >>> 'test' + wrong + """ + pass + ''', + 8, # Top-level function: 4 (function) + 4 (docstring) + ), + ( + ''' + class Class1: + def method(self): + """ + Example: + >>> 'test' + wrong + """ + pass + ''', + 12, # Class method: 4 (class) + 4 (method) + 4 (docstring) + ), + ], +) +def test_indentation_at_various_levels(create_test_file, test_content, expected_spaces): + """ + Test indentation preservation at various nesting levels. + """ + test_file = create_test_file(test_content) + + xdoctest.doctest_module( + str(test_file), command="all", config={"write_outputs": True} + ) + + lines = test_file.read_text().splitlines() + + test_lines = [line for line in lines if "'test'" in line and ">>>" not in line] + + test_line = test_lines[0] + indent_count = len(test_line) - len(test_line.lstrip()) + assert indent_count == expected_spaces + + +def test_write_outputs_with_eval_none(create_test_file): + """ + Test write-outputs with expressions that evaluate to None. + + Note: Currently, when an expression evaluates to None with no stdout, + the output is treated as empty and not written. This is a known limitation. + """ + test_content = ''' + def example(): + """ + Example: + >>> None + wrong + """ + pass + ''' + test_file = create_test_file(test_content) + + xdoctest.doctest_module( + str(test_file), command="all", config={"write_outputs": True} + ) + + # File should remain unchanged because got_eval=None is treated as no output + assert test_file.read_text() == textwrap.dedent(test_content) + + +def test_write_outputs_empty_output(create_test_file): + """ + Test that statements with no output don't modify the file when output is empty. + """ + test_content = ''' + def example(): + """ + Example: + >>> x = 5 + wrong output + """ + pass + ''' + test_file = create_test_file(test_content) + + xdoctest.doctest_module( + str(test_file), command="all", config={"write_outputs": True} + ) + + # When got output is empty (""), formatted_lines will be empty + # so _compute_part_modification returns None and file is unchanged + assert test_file.read_text() == textwrap.dedent(test_content) + + +def test_fill_missing_wants_no_output(create_test_file): + """ + Test that --fill-missing-wants doesn't add wants for statements with no output. + """ + test_content = ''' + def example(): + """ + Example: + >>> x = 5 + """ + pass + ''' + test_file = create_test_file(test_content) + + xdoctest.doctest_module( + str(test_file), + command="all", + config={"write_outputs": True, "fill_missing_wants": True}, + ) + + # Should not add output for assignment statements + assert test_file.read_text() == textwrap.dedent(test_content)