diff --git a/python/sglang/srt/debug_utils/source_patcher/__init__.py b/python/sglang/srt/debug_utils/source_patcher/__init__.py new file mode 100644 index 000000000..c853fad17 --- /dev/null +++ b/python/sglang/srt/debug_utils/source_patcher/__init__.py @@ -0,0 +1,12 @@ +from sglang.srt.debug_utils.source_patcher.code_patcher import ( + CodePatcher, + apply_patches_from_config, + patch_function, +) +from sglang.srt.debug_utils.source_patcher.types import ( + EditSpec, + PatchApplicationError, + PatchConfig, + PatchSpec, + PatchState, +) diff --git a/python/sglang/srt/debug_utils/source_patcher/code_patcher.py b/python/sglang/srt/debug_utils/source_patcher/code_patcher.py new file mode 100644 index 000000000..421f2ccd2 --- /dev/null +++ b/python/sglang/srt/debug_utils/source_patcher/code_patcher.py @@ -0,0 +1,188 @@ +import importlib +import inspect +import textwrap +import types +from collections.abc import Callable +from typing import Any, Optional + +import yaml + +from sglang.srt.debug_utils.source_patcher.source_editor import apply_edits +from sglang.srt.debug_utils.source_patcher.types import ( + EditSpec, + PatchConfig, + PatchSpec, + PatchState, +) + + +def apply_patches_from_config( + yaml_content: str, + *, + extra_imports: Optional[list[str]] = None, +) -> list[PatchState]: + """Parse a YAML config string and apply all patches. + + Args: + yaml_content: YAML string with patch specifications. + extra_imports: Import lines inserted once at the top of each patched + function body (e.g. ["from pkg import foo"]). The caller (dumper) + uses this so users don't have to write boilerplate in YAML. + """ + raw: dict[str, Any] = yaml.safe_load(yaml_content) + config: PatchConfig = PatchConfig(**raw) + + if extra_imports: + config = _inject_preamble(config=config, extra_imports=extra_imports) + + return _apply_specs(config.patches) + + +class CodePatcher: + """Context manager that patches functions on enter and restores on exit.""" + + def __init__(self, *, patches: list[PatchSpec]) -> None: + self._patches = patches + self._states: list[PatchState] = [] + + def __enter__(self) -> "CodePatcher": + self._states = _apply_specs(self._patches) + return self + + def __exit__( + self, + exc_type: Optional[type], + exc_val: Optional[BaseException], + exc_tb: Optional[Any], + ) -> None: + for state in reversed(self._states): + state.restore() + self._states.clear() + + +def patch_function( + *, + target: Callable[..., Any], + edits: list[EditSpec], + preamble: str = "", +) -> PatchState: + """Patch a function by modifying its source and replacing __code__. + + 1. inspect.getsource -> get original source + 2. apply_edits -> modify source text + 3. optionally prepend preamble (e.g. import lines) inside the function body + 4. compile + exec -> get new code object + 5. replace target.__code__ + + Returns PatchState that can restore the original code. + """ + original_code: types.CodeType = target.__code__ + + source: str = inspect.getsource(target) + modified_source: str = apply_edits(source=source, edits=edits) + modified_source = textwrap.dedent(modified_source) + + if preamble.strip(): + modified_source = _insert_preamble(source=modified_source, preamble=preamble) + + code: types.CodeType = compile(modified_source, inspect.getfile(target), "exec") + temp_namespace: dict[str, Any] = {} + exec(code, target.__globals__, temp_namespace) + + new_fn: Any = temp_namespace[target.__name__] + target.__code__ = new_fn.__code__ + + return PatchState(target_fn=target, original_code=original_code) + + +# --------------------------------- private --------------------------------- + + +def _apply_specs(specs: list[PatchSpec]) -> list[PatchState]: + states: list[PatchState] = [] + for spec in specs: + target_fn: Callable[..., Any] = _resolve_target(spec.target) + print(f"[source_patcher] patching {spec.target}") + state: PatchState = patch_function( + target=target_fn, edits=spec.edits, preamble=spec.preamble + ) + states.append(state) + return states + + +def _inject_preamble(*, config: PatchConfig, extra_imports: list[str]) -> PatchConfig: + """Set preamble on every PatchSpec so imports are inserted once at function top.""" + import_block: str = "\n".join(extra_imports) + new_patches: list[PatchSpec] = [] + + for spec in config.patches: + existing: str = spec.preamble + combined: str = ( + existing + "\n" + import_block if existing.strip() else import_block + ) + new_patches.append( + PatchSpec(target=spec.target, edits=spec.edits, preamble=combined) + ) + + return PatchConfig(patches=new_patches) + + +def _insert_preamble(*, source: str, preamble: str) -> str: + """Insert preamble lines right after the function signature (and optional docstring).""" + lines: list[str] = source.splitlines() + + signature_end: int = _find_signature_end(lines) + + body_start: int = signature_end + 1 + body_indent: str = "" + for i in range(body_start, len(lines)): + if lines[i].strip(): + body_indent = " " * (len(lines[i]) - len(lines[i].lstrip())) + body_start = i + break + + preamble_lines: list[str] = [ + body_indent + pl for pl in preamble.strip().splitlines() + ] + return "\n".join(lines[:body_start] + preamble_lines + lines[body_start:]) + + +def _find_signature_end(lines: list[str]) -> int: + """Find the line index where the function signature ends (the line with trailing colon).""" + for i, line in enumerate(lines): + if line.rstrip().endswith(":"): + return i + return 0 + + +def _resolve_target(qualified_name: str) -> Callable[..., Any]: + """Resolve 'pkg.mod.Class.method' to the actual function object. + + Tries progressively shorter module paths from right to left, + then uses getattr for the remaining attribute chain. + """ + parts: list[str] = qualified_name.split(".") + + target: Any = None + for split_idx in range(len(parts), 0, -1): + module_path: str = ".".join(parts[:split_idx]) + try: + target = importlib.import_module(module_path) + attr_parts: list[str] = parts[split_idx:] + break + except ImportError: + continue + else: + raise ImportError(f"could not import any module prefix of '{qualified_name}'") + + for attr_name in attr_parts: + target = getattr(target, attr_name) + + if isinstance(target, classmethod): + target = target.__func__ + if not callable(target): + raise TypeError( + f"resolved target '{qualified_name}' is not callable: {type(target)}" + ) + + return target diff --git a/python/sglang/srt/debug_utils/source_patcher/source_editor.py b/python/sglang/srt/debug_utils/source_patcher/source_editor.py new file mode 100644 index 000000000..0f4b0805a --- /dev/null +++ b/python/sglang/srt/debug_utils/source_patcher/source_editor.py @@ -0,0 +1,108 @@ +from sglang.srt.debug_utils.source_patcher.types import EditSpec, PatchApplicationError + + +def apply_edits(*, source: str, edits: list[EditSpec]) -> str: + """Apply a sequence of match/replacement edits to source text. + + Each edit is applied sequentially so later edits see the result of earlier ones. + """ + result: str = source + for edit in edits: + result = _apply_single_edit(source=result, edit=edit) + return result + + +def _apply_single_edit(*, source: str, edit: EditSpec) -> str: + """Apply a single match/replacement edit to the source text.""" + match_text: str = edit.match.strip() + if not match_text: + raise PatchApplicationError("empty match text") + + source_lines: list[str] = source.splitlines() + match_lines: list[str] = match_text.splitlines() + + start_idx: int = _find_match(source_lines=source_lines, match_lines=match_lines) + match_len: int = len(match_lines) + + original_indent: int = _leading_spaces(source_lines[start_idx]) + + effective_replacement: str = _resolve_replacement(edit=edit, match_text=match_text) + replacement_lines: list[str] = ( + effective_replacement.splitlines() if effective_replacement else [] + ) + aligned: list[str] = _realign_replacement( + replacement_lines=replacement_lines, original_indent=original_indent + ) + new_lines: list[str] = ( + source_lines[:start_idx] + aligned + source_lines[start_idx + match_len :] + ) + + trailing_newline: str = "\n" if source.endswith("\n") else "" + return "\n".join(new_lines) + trailing_newline + + +def _resolve_replacement(*, edit: EditSpec, match_text: str) -> str: + """Return the effective replacement text, handling replacement, prepend, and append modes.""" + if edit.prepend.strip(): + return edit.prepend.strip() + "\n" + match_text + if edit.append.strip(): + return match_text + "\n" + edit.append.strip() + return edit.replacement.strip() + + +def _find_match(*, source_lines: list[str], match_lines: list[str]) -> int: + """Find the start index of match_lines in source_lines (strip-compared). + + Returns the index of the first matching line. + Raises PatchApplicationError if not found or found multiple times. + """ + stripped_source: list[str] = [line.strip() for line in source_lines] + stripped_match: list[str] = [line.strip() for line in match_lines] + match_len: int = len(stripped_match) + + found_indices: list[int] = [ + i + for i in range(len(stripped_source) - match_len + 1) + if stripped_source[i : i + match_len] == stripped_match + ] + + if len(found_indices) == 0: + preview: str = "\n".join(match_lines) + raise PatchApplicationError(f"match text not found in source:\n{preview}") + if len(found_indices) > 1: + preview = "\n".join(match_lines) + raise PatchApplicationError( + f"match text found multiple times ({len(found_indices)} occurrences) in source:\n{preview}" + ) + + return found_indices[0] + + +def _realign_replacement( + *, replacement_lines: list[str], original_indent: int +) -> list[str]: + """Realign replacement lines to the original indentation level. + + Strategy: + - Take the leading spaces of the first non-empty replacement line as base_indent + - For each replacement line: remove base_indent, add original_indent + """ + non_empty: list[str] = [line for line in replacement_lines if line.strip()] + if not non_empty: + return [] + + base_indent: int = _leading_spaces(non_empty[0]) + result: list[str] = [] + + for line in replacement_lines: + if not line.strip(): + result.append("") + else: + stripped = line[min(base_indent, len(line) - len(line.lstrip())) :] + result.append(" " * original_indent + stripped) + + return result + + +def _leading_spaces(line: str) -> int: + return len(line) - len(line.lstrip(" ")) diff --git a/python/sglang/srt/debug_utils/source_patcher/types.py b/python/sglang/srt/debug_utils/source_patcher/types.py new file mode 100644 index 000000000..9ff44cba6 --- /dev/null +++ b/python/sglang/srt/debug_utils/source_patcher/types.py @@ -0,0 +1,63 @@ +import types +from collections.abc import Callable +from typing import Any + +from pydantic import BaseModel, ConfigDict, model_validator + + +class PatchApplicationError(Exception): + """match text not found or not unique in source.""" + + +class _StrictBase(BaseModel): + model_config = ConfigDict(extra="forbid") + + +class EditSpec(_StrictBase): + """Specify one edit: replace, prepend before, or append after the matched text. + + Use ``replacement`` to substitute the matched text (empty string = delete). + Use ``prepend`` to keep the matched text and add lines before it. + Use ``append`` to keep the matched text and add lines after it. + Only one of ``replacement``, ``prepend``, and ``append`` may be set. + """ + + match: str + replacement: str = "" + prepend: str = "" + append: str = "" + + @model_validator(mode="after") + def _check_modes_mutually_exclusive(self) -> "EditSpec": + active: list[str] = [ + name + for name in ("replacement", "prepend", "append") + if getattr(self, name).strip() + ] + if len(active) > 1: + raise ValueError( + f"only one of 'replacement', 'prepend', 'append' may be set, " + f"got: {', '.join(active)}" + ) + return self + + +class PatchSpec(_StrictBase): + target: str + edits: list[EditSpec] + preamble: str = "" + + +class PatchConfig(_StrictBase): + patches: list[PatchSpec] + + +class PatchState: + def __init__( + self, *, target_fn: Callable[..., Any], original_code: types.CodeType + ) -> None: + self.target_fn = target_fn + self.original_code = original_code + + def restore(self) -> None: + self.target_fn.__code__ = self.original_code diff --git a/test/registered/debug_utils/source_patcher/conftest.py b/test/registered/debug_utils/source_patcher/conftest.py new file mode 100644 index 000000000..f585c5838 --- /dev/null +++ b/test/registered/debug_utils/source_patcher/conftest.py @@ -0,0 +1,66 @@ +"""Shared fixtures for source_patcher tests. + +The sample module is defined as an inline string and written to a temp file +at test time, avoiding CI complaints about fixture files without test registration. +""" + +import importlib.util +import sys +import tempfile +from pathlib import Path +from types import ModuleType + +import pytest + +SAMPLE_MODULE_NAME = "_source_patcher_test_fixtures.sample_module" + +SAMPLE_MODULE_SOURCE = '''\ +GLOBAL_VAR = "global_value" + + +class HelperClass: + """Utility class referenced by SampleClass to test cross-class calls.""" + + @staticmethod + def format_value(value: str) -> str: + return f"[{value}]" + + +class SampleClass: + def greet(self, name: str) -> str: + greeting = f"hello {name}" + return greeting + + def compute(self, x: int) -> int: + result = x * 2 + 1 + return result + + def uses_global(self) -> str: + return f"value={GLOBAL_VAR}" + + def uses_helper(self, value: str) -> str: + return HelperClass.format_value(value) + + +def standalone_function(a: int, b: int) -> int: + return a + b +''' + + +@pytest.fixture(scope="session") +def sample_module() -> ModuleType: + """Load the sample module from a temp file and register it in sys.modules.""" + if SAMPLE_MODULE_NAME in sys.modules: + return sys.modules[SAMPLE_MODULE_NAME] + + tmpdir = tempfile.mkdtemp(prefix="source_patcher_fixtures_") + module_path = Path(tmpdir) / "sample_module.py" + module_path.write_text(SAMPLE_MODULE_SOURCE) + + spec = importlib.util.spec_from_file_location(SAMPLE_MODULE_NAME, module_path) + assert spec is not None + assert spec.loader is not None + module = importlib.util.module_from_spec(spec) + sys.modules[SAMPLE_MODULE_NAME] = module + spec.loader.exec_module(module) + return module diff --git a/test/registered/debug_utils/source_patcher/test_code_patcher.py b/test/registered/debug_utils/source_patcher/test_code_patcher.py new file mode 100644 index 000000000..56e0ab721 --- /dev/null +++ b/test/registered/debug_utils/source_patcher/test_code_patcher.py @@ -0,0 +1,253 @@ +from types import ModuleType + +import pytest + +from sglang.srt.debug_utils.source_patcher.code_patcher import ( + CodePatcher, + _resolve_target, + patch_function, +) +from sglang.srt.debug_utils.source_patcher.types import EditSpec, PatchSpec + +SAMPLE_MODULE_NAME = "_source_patcher_test_fixtures.sample_module" + + +class TestPatchFunction: + def test_basic_patch_changes_behavior(self, sample_module: ModuleType) -> None: + cls = sample_module.SampleClass + obj = cls() + assert obj.greet("world") == "hello world" + + state = patch_function( + target=cls.greet, + edits=[ + EditSpec( + match='greeting = f"hello {name}"', + replacement='greeting = f"patched {name}"', + ) + ], + ) + try: + assert obj.greet("world") == "patched world" + finally: + state.restore() + + assert obj.greet("world") == "hello world" + + def test_globals_preserved_after_patch(self, sample_module: ModuleType) -> None: + cls = sample_module.SampleClass + obj = cls() + assert obj.uses_global() == "value=global_value" + + state = patch_function( + target=cls.uses_global, + edits=[ + EditSpec( + match='return f"value={GLOBAL_VAR}"', + replacement='return f"patched_value={GLOBAL_VAR}"', + ) + ], + ) + try: + assert obj.uses_global() == "patched_value=global_value" + finally: + state.restore() + + def test_function_identity_preserved(self, sample_module: ModuleType) -> None: + cls = sample_module.SampleClass + fn_id_before = id(cls.greet) + + state = patch_function( + target=cls.greet, + edits=[ + EditSpec( + match='greeting = f"hello {name}"', + replacement='greeting = f"patched {name}"', + ) + ], + ) + try: + assert id(cls.greet) == fn_id_before + finally: + state.restore() + + def test_patch_standalone_function(self, sample_module: ModuleType) -> None: + fn = sample_module.standalone_function + assert fn(2, 3) == 5 + + state = patch_function( + target=fn, + edits=[ + EditSpec( + match="return a + b", + replacement="return a * b", + ) + ], + ) + try: + assert fn(2, 3) == 6 + finally: + state.restore() + + assert fn(2, 3) == 5 + + def test_patched_code_can_reference_global_variable( + self, sample_module: ModuleType + ) -> None: + """Replacement code that references a module-level global should work.""" + cls = sample_module.SampleClass + obj = cls() + + state = patch_function( + target=cls.greet, + edits=[ + EditSpec( + match='greeting = f"hello {name}"', + replacement='greeting = f"{GLOBAL_VAR} {name}"', + ) + ], + ) + try: + assert obj.greet("world") == "global_value world" + finally: + state.restore() + + def test_patched_code_can_call_another_class_method( + self, sample_module: ModuleType + ) -> None: + """Replacement code that calls HelperClass.format_value should work.""" + cls = sample_module.SampleClass + obj = cls() + + state = patch_function( + target=cls.greet, + edits=[ + EditSpec( + match='greeting = f"hello {name}"', + replacement="greeting = HelperClass.format_value(name)", + ) + ], + ) + try: + assert obj.greet("world") == "[world]" + finally: + state.restore() + + def test_patched_code_uses_helper_via_existing_method( + self, sample_module: ModuleType + ) -> None: + """The uses_helper method already calls HelperClass; verify it survives patching.""" + cls = sample_module.SampleClass + obj = cls() + assert obj.uses_helper("test") == "[test]" + + state = patch_function( + target=cls.uses_helper, + edits=[ + EditSpec( + match="return HelperClass.format_value(value)", + replacement='return HelperClass.format_value("patched_" + value)', + ) + ], + ) + try: + assert obj.uses_helper("test") == "[patched_test]" + finally: + state.restore() + + assert obj.uses_helper("test") == "[test]" + + +class TestResolveTarget: + def test_resolve_class_method(self, sample_module: ModuleType) -> None: + target = _resolve_target(f"{SAMPLE_MODULE_NAME}.SampleClass.greet") + assert target is sample_module.SampleClass.greet + + def test_resolve_standalone_function(self, sample_module: ModuleType) -> None: + target = _resolve_target(f"{SAMPLE_MODULE_NAME}.standalone_function") + assert target is sample_module.standalone_function + + def test_resolve_nonexistent_raises(self, sample_module: ModuleType) -> None: + with pytest.raises((ImportError, AttributeError)): + _resolve_target(f"{SAMPLE_MODULE_NAME}.NonexistentClass.method") + + +class TestCodePatcher: + def test_context_manager_patches_and_restores( + self, sample_module: ModuleType + ) -> None: + cls = sample_module.SampleClass + obj = cls() + assert obj.greet("world") == "hello world" + + patches = [ + PatchSpec( + target=f"{SAMPLE_MODULE_NAME}.SampleClass.greet", + edits=[ + EditSpec( + match='greeting = f"hello {name}"', + replacement='greeting = f"ctx_patched {name}"', + ) + ], + ) + ] + + with CodePatcher(patches=patches): + assert obj.greet("world") == "ctx_patched world" + + assert obj.greet("world") == "hello world" + + def test_context_manager_multiple_patches(self, sample_module: ModuleType) -> None: + cls = sample_module.SampleClass + obj = cls() + + patches = [ + PatchSpec( + target=f"{SAMPLE_MODULE_NAME}.SampleClass.greet", + edits=[ + EditSpec( + match='greeting = f"hello {name}"', + replacement='greeting = f"p1 {name}"', + ) + ], + ), + PatchSpec( + target=f"{SAMPLE_MODULE_NAME}.SampleClass.compute", + edits=[ + EditSpec( + match="result = x * 2 + 1", + replacement="result = x * 100", + ) + ], + ), + ] + + with CodePatcher(patches=patches): + assert obj.greet("world") == "p1 world" + assert obj.compute(5) == 500 + + assert obj.greet("world") == "hello world" + assert obj.compute(5) == 11 + + def test_restores_on_exception(self, sample_module: ModuleType) -> None: + cls = sample_module.SampleClass + obj = cls() + + patches = [ + PatchSpec( + target=f"{SAMPLE_MODULE_NAME}.SampleClass.greet", + edits=[ + EditSpec( + match='greeting = f"hello {name}"', + replacement='greeting = f"err_patched {name}"', + ) + ], + ) + ] + + with pytest.raises(RuntimeError): + with CodePatcher(patches=patches): + assert obj.greet("world") == "err_patched world" + raise RuntimeError("test error") + + assert obj.greet("world") == "hello world" diff --git a/test/registered/debug_utils/source_patcher/test_source_editor.py b/test/registered/debug_utils/source_patcher/test_source_editor.py new file mode 100644 index 000000000..10510732f --- /dev/null +++ b/test/registered/debug_utils/source_patcher/test_source_editor.py @@ -0,0 +1,289 @@ +import pytest +from pydantic import ValidationError + +from sglang.srt.debug_utils.source_patcher.source_editor import apply_edits +from sglang.srt.debug_utils.source_patcher.types import EditSpec, PatchApplicationError + + +class TestApplyEdits: + """Tests for the apply_edits() source text transformation function.""" + + def test_single_line_match_to_multiline_replacement(self) -> None: + source = "def foo():\n" " x = compute()\n" " return x\n" + edits = [ + EditSpec( + match="x = compute()", + replacement="x = compute()\nprint(x)", + ) + ] + result = apply_edits(source=source, edits=edits) + assert result == ( + "def foo():\n" " x = compute()\n" " print(x)\n" " return x\n" + ) + + def test_pure_insertion(self) -> None: + source = "def foo():\n" " a = 1\n" " b = 2\n" + edits = [ + EditSpec( + match="a = 1", + replacement="a = 1\nprint(a)", + ) + ] + result = apply_edits(source=source, edits=edits) + assert result == ("def foo():\n" " a = 1\n" " print(a)\n" " b = 2\n") + + def test_pure_deletion_via_empty_replacement(self) -> None: + source = "def foo():\n" " debug_log()\n" " return 42\n" + edits = [ + EditSpec( + match="debug_log()", + replacement="", + ) + ] + result = apply_edits(source=source, edits=edits) + assert result == ("def foo():\n" " return 42\n") + + def test_deletion_fewer_lines(self) -> None: + source = "def foo():\n" " a = 1\n" " b = 2\n" " c = 3\n" + edits = [ + EditSpec( + match="a = 1\nb = 2", + replacement="ab = 3", + ) + ] + result = apply_edits(source=source, edits=edits) + assert result == ("def foo():\n" " ab = 3\n" " c = 3\n") + + def test_multiline_match_to_multiline_replacement(self) -> None: + source = ( + "def foo():\n" + " result = self.attn(\n" + " q=q,\n" + " k=k,\n" + " )\n" + " return result\n" + ) + edits = [ + EditSpec( + match="result = self.attn(\n q=q,\n k=k,\n)", + replacement="result = self.attn(\n q=q,\n k=k,\n v=v,\n)", + ) + ] + result = apply_edits(source=source, edits=edits) + assert result == ( + "def foo():\n" + " result = self.attn(\n" + " q=q,\n" + " k=k,\n" + " v=v,\n" + " )\n" + " return result\n" + ) + + def test_indent_alignment_deep_nesting(self) -> None: + source = ( + "class Foo:\n" + " class Bar:\n" + " def method(self):\n" + " x = compute()\n" + " return x\n" + ) + edits = [ + EditSpec( + match="x = compute()", + replacement="x = compute()\nprint(x)", + ) + ] + result = apply_edits(source=source, edits=edits) + assert result == ( + "class Foo:\n" + " class Bar:\n" + " def method(self):\n" + " x = compute()\n" + " print(x)\n" + " return x\n" + ) + + def test_match_not_found_raises(self) -> None: + source = "def foo():\n return 1\n" + edits = [EditSpec(match="nonexistent_call()", replacement="replaced()")] + with pytest.raises(PatchApplicationError, match="not found"): + apply_edits(source=source, edits=edits) + + def test_match_found_multiple_times_raises(self) -> None: + source = "def foo():\n" " print(1)\n" " print(1)\n" + edits = [EditSpec(match="print(1)", replacement="print(2)")] + with pytest.raises(PatchApplicationError, match="multiple"): + apply_edits(source=source, edits=edits) + + def test_multiple_edits_applied_sequentially(self) -> None: + source = "def foo():\n" " a = 1\n" " b = 2\n" " return a + b\n" + edits = [ + EditSpec(match="a = 1", replacement="a = 10"), + EditSpec(match="b = 2", replacement="b = 20"), + ] + result = apply_edits(source=source, edits=edits) + assert result == ( + "def foo():\n" " a = 10\n" " b = 20\n" " return a + b\n" + ) + + def test_strip_matching_ignores_leading_trailing_whitespace(self) -> None: + source = "def foo():\n" " x = compute()\n" " return x\n" + edits = [ + EditSpec( + match=" x = compute() ", + replacement="x = replaced()", + ) + ] + result = apply_edits(source=source, edits=edits) + assert result == ("def foo():\n" " x = replaced()\n" " return x\n") + + def test_replacement_indented_text_realigned(self) -> None: + """replacement text with its own indentation gets realigned to match source.""" + source = "def foo():\n" " x = compute()\n" " return x\n" + edits = [ + EditSpec( + match="x = compute()", + replacement="x = compute()\nprint(x)", + ) + ] + result = apply_edits(source=source, edits=edits) + assert result == ( + "def foo():\n" + " x = compute()\n" + " print(x)\n" + " return x\n" + ) + + def test_replacement_with_existing_indent_realigned(self) -> None: + """replacement text already has indentation that should be rebased.""" + source = "def foo():\n" " if True:\n" " x = 1\n" " return x\n" + edits = [ + EditSpec( + match="x = 1", + replacement="x = 1\nif x > 0:\n print(x)", + ) + ] + result = apply_edits(source=source, edits=edits) + assert result == ( + "def foo():\n" + " if True:\n" + " x = 1\n" + " if x > 0:\n" + " print(x)\n" + " return x\n" + ) + + def test_append_keeps_match_and_adds_after(self) -> None: + source = "def foo():\n" " x = compute()\n" " return x\n" + edits = [EditSpec(match="x = compute()", append="print(x)")] + result = apply_edits(source=source, edits=edits) + assert result == ( + "def foo():\n" " x = compute()\n" " print(x)\n" " return x\n" + ) + + def test_append_multiline_match(self) -> None: + source = ( + "def foo():\n" + " result = call(\n" + " a=1,\n" + " b=2,\n" + " )\n" + " return result\n" + ) + edits = [ + EditSpec( + match="result = call(\n a=1,\n b=2,\n)", + append="dumper.dump('result', result)", + ) + ] + result = apply_edits(source=source, edits=edits) + assert result == ( + "def foo():\n" + " result = call(\n" + " a=1,\n" + " b=2,\n" + " )\n" + " dumper.dump('result', result)\n" + " return result\n" + ) + + def test_prepend_adds_before_match(self) -> None: + source = "def foo():\n" " x = compute()\n" " return x\n" + edits = [EditSpec(match="x = compute()", prepend="print('before')")] + result = apply_edits(source=source, edits=edits) + assert result == ( + "def foo():\n" + " print('before')\n" + " x = compute()\n" + " return x\n" + ) + + def test_prepend_multiline(self) -> None: + source = "def foo():\n" " return x\n" + edits = [EditSpec(match="return x", prepend="a = 1\nb = 2")] + result = apply_edits(source=source, edits=edits) + assert result == ("def foo():\n" " a = 1\n" " b = 2\n" " return x\n") + + def test_prepend_deep_indent(self) -> None: + source = ( + "class Foo:\n" + " class Bar:\n" + " def method(self):\n" + " return x\n" + ) + edits = [EditSpec(match="return x", prepend="dumper.dump('x', x)")] + result = apply_edits(source=source, edits=edits) + assert result == ( + "class Foo:\n" + " class Bar:\n" + " def method(self):\n" + " dumper.dump('x', x)\n" + " return x\n" + ) + + def test_prepend_multiline_match(self) -> None: + source = ( + "def foo():\n" + " result = call(\n" + " a=1,\n" + " )\n" + " return result\n" + ) + edits = [ + EditSpec( + match="result = call(\n a=1,\n)", + prepend="dumper.dump('before', x)", + ) + ] + result = apply_edits(source=source, edits=edits) + assert result == ( + "def foo():\n" + " dumper.dump('before', x)\n" + " result = call(\n" + " a=1,\n" + " )\n" + " return result\n" + ) + + def test_replacement_and_append_mutually_exclusive(self) -> None: + with pytest.raises(ValidationError, match="only one of"): + EditSpec(match="x = 1", replacement="x = 2", append="print(x)") + + def test_replacement_and_prepend_mutually_exclusive(self) -> None: + with pytest.raises(ValidationError, match="only one of"): + EditSpec(match="x = 1", replacement="x = 2", prepend="print(x)") + + def test_prepend_and_append_mutually_exclusive(self) -> None: + with pytest.raises(ValidationError, match="only one of"): + EditSpec(match="x = 1", prepend="a()", append="b()") + + def test_second_edit_sees_result_of_first(self) -> None: + """Edits are applied sequentially; second edit matches modified source.""" + source = "def foo():\n" " x = 1\n" " return x\n" + edits = [ + EditSpec(match="x = 1", replacement="x = 1\ny = 2"), + EditSpec(match="y = 2", replacement="y = 20"), + ] + result = apply_edits(source=source, edits=edits) + assert result == ("def foo():\n" " x = 1\n" " y = 20\n" " return x\n")