| 1 | from __future__ import annotations |
| 2 | |
| 3 | from dataclasses import dataclass |
| 4 | from typing import Iterable |
| 5 | |
| 6 | |
| 7 | class ContextPatchError(ValueError): |
| 8 | """Raised when a context patch cannot be parsed or applied safely.""" |
| 9 | |
| 10 | |
| 11 | @dataclass(frozen=True) |
| 12 | class ContextPatchApplyResult: |
| 13 | content: str |
| 14 | line_from: int |
| 15 | line_to: int |
| 16 | hunk_count: int |
| 17 | |
| 18 | |
| 19 | @dataclass(frozen=True) |
| 20 | class _PatchLine: |
| 21 | kind: str |
| 22 | text: str |
| 23 | |
| 24 | |
| 25 | @dataclass(frozen=True) |
| 26 | class _Hunk: |
| 27 | anchor: str |
| 28 | lines: tuple[_PatchLine, ...] |
| 29 | |
| 30 | |
| 31 | @dataclass(frozen=True) |
| 32 | class _HunkApplyResult: |
| 33 | cursor: int |
| 34 | line_from: int |
| 35 | line_to: int |
| 36 | |
| 37 | |
| 38 | _FILE_HEADERS = ( |
| 39 | "*** Update File:", |
| 40 | "*** Add File:", |
| 41 | "*** Delete File:", |
| 42 | "*** End Patch", |
| 43 | ) |
| 44 | |
| 45 | |
| 46 | def apply_context_patch(content: str, patch_text: str) -> str: |
| 47 | """Apply a PseudoPatch-inspired context patch to one text file.""" |
| 48 | return apply_context_patch_with_metadata(content, patch_text).content |
| 49 | |
| 50 | |
| 51 | def apply_context_patch_with_metadata( |
| 52 | content: str, patch_text: str |
| 53 | ) -> ContextPatchApplyResult: |
| 54 | """Apply a context patch and report the touched line range.""" |
| 55 | body = _extract_single_file_body(patch_text) |
| 56 | hunks = _parse_hunks(body) |
| 57 | if not hunks: |
| 58 | raise ContextPatchError("patch_text must contain at least one update hunk") |
| 59 | |
| 60 | lines = content.split("\n") |
| 61 | cursor = 0 |
| 62 | line_from: int | None = None |
| 63 | line_to = 1 |
| 64 | for hunk in hunks: |
| 65 | result = _apply_hunk(lines, hunk, cursor) |
| 66 | cursor = result.cursor |
| 67 | line_from = ( |
| 68 | result.line_from if line_from is None |
| 69 | else min(line_from, result.line_from) |
| 70 | ) |
| 71 | line_to = max(line_to, result.line_to) |
| 72 | |
| 73 | return ContextPatchApplyResult( |
| 74 | content="\n".join(lines), |
| 75 | line_from=line_from or 1, |
| 76 | line_to=line_to, |
| 77 | hunk_count=len(hunks), |
| 78 | ) |
| 79 | |
| 80 | |
| 81 | def _extract_single_file_body(patch_text: str) -> list[str]: |
| 82 | raw_lines = [line.rstrip("\r") for line in str(patch_text).splitlines()] |
| 83 | lines = _trim_outer_blank_lines(raw_lines) |
| 84 | if not lines: |
| 85 | raise ContextPatchError("patch_text is required") |
| 86 | |
| 87 | if not lines[0].startswith("*** Begin Patch"): |
| 88 | return lines |
| 89 | if len(lines) < 2 or lines[-1] != "*** End Patch": |
| 90 | raise ContextPatchError("patch_text missing *** End Patch") |
| 91 | |
| 92 | body: list[str] = [] |
| 93 | in_update = False |
| 94 | update_count = 0 |
| 95 | for line in lines[1:-1]: |
| 96 | if line.startswith("*** Update File:"): |
| 97 | if update_count: |
| 98 | raise ContextPatchError( |
| 99 | "patch_text may update only one file per operation" |
| 100 | ) |
| 101 | in_update = True |
| 102 | update_count += 1 |
| 103 | continue |
| 104 | if line.startswith("*** Move to:"): |
| 105 | raise ContextPatchError("patch_text does not support file moves") |
| 106 | if line.startswith(("*** Add File:", "*** Delete File:")): |
| 107 | raise ContextPatchError("patch_text supports update hunks only") |
| 108 | if in_update: |
| 109 | body.append(line) |
| 110 | |
| 111 | if not update_count: |
| 112 | raise ContextPatchError("patch_text must include an update file block") |
| 113 | return body |
| 114 | |
| 115 | |
| 116 | def _trim_outer_blank_lines(lines: list[str]) -> list[str]: |
| 117 | start = 0 |
| 118 | end = len(lines) |
| 119 | while start < end and not lines[start].strip(): |
| 120 | start += 1 |
| 121 | while end > start and not lines[end - 1].strip(): |
| 122 | end -= 1 |
| 123 | return lines[start:end] |
| 124 | |
| 125 | |
| 126 | def _parse_hunks(lines: list[str]) -> list[_Hunk]: |
| 127 | hunks: list[_Hunk] = [] |
| 128 | anchor = "" |
| 129 | current: list[_PatchLine] = [] |
| 130 | |
| 131 | def finish_current() -> None: |
| 132 | nonlocal anchor, current |
| 133 | if current: |
| 134 | hunks.append(_Hunk(anchor=anchor, lines=tuple(current))) |
| 135 | current = [] |
| 136 | anchor = "" |
| 137 | |
| 138 | for line in lines: |
| 139 | if line.startswith(_FILE_HEADERS): |
| 140 | finish_current() |
| 141 | break |
| 142 | if line.startswith("@@"): |
| 143 | finish_current() |
| 144 | anchor = line[2:].strip() |
| 145 | continue |
| 146 | if line.startswith("***"): |
| 147 | raise ContextPatchError(f"invalid patch control line: {line}") |
| 148 | |
| 149 | if line == "": |
| 150 | current.append(_PatchLine(" ", "")) |
| 151 | continue |
| 152 | if line[0] not in {" ", "+", "-"}: |
| 153 | raise ContextPatchError(f"invalid patch line prefix: {line}") |
| 154 | current.append(_PatchLine(line[0], line[1:])) |
| 155 | |
| 156 | finish_current() |
| 157 | return hunks |
| 158 | |
| 159 | |
| 160 | def _apply_hunk( |
| 161 | lines: list[str], hunk: _Hunk, cursor: int |
| 162 | ) -> _HunkApplyResult: |
| 163 | old_lines = [line.text for line in hunk.lines if line.kind in {" ", "-"}] |
| 164 | new_lines = [line.text for line in hunk.lines if line.kind in {" ", "+"}] |
| 165 | |
| 166 | if old_lines == new_lines: |
| 167 | raise ContextPatchError("patch hunk does not change content") |
| 168 | if not old_lines: |
| 169 | if not hunk.anchor: |
| 170 | raise ContextPatchError("insert-only patch hunk needs an @@ anchor") |
| 171 | insert_at = _find_anchor(lines, hunk.anchor, cursor) |
| 172 | lines[insert_at:insert_at] = new_lines |
| 173 | return _HunkApplyResult( |
| 174 | cursor=insert_at + len(new_lines), |
| 175 | line_from=insert_at + 1, |
| 176 | line_to=insert_at + max(len(new_lines), 1), |
| 177 | ) |
| 178 | |
| 179 | start = cursor |
| 180 | if hunk.anchor: |
| 181 | start = _find_anchor(lines, hunk.anchor, cursor) |
| 182 | |
| 183 | try: |
| 184 | match_index = _find_context(lines, old_lines, start, anchored=bool(hunk.anchor)) |
| 185 | except ContextPatchError: |
| 186 | if not hunk.anchor: |
| 187 | raise |
| 188 | # Models often anchor on the line they want to replace. Preserve the |
| 189 | # insert-after-anchor rule, but allow replacement context to start at |
| 190 | # the anchor line when it is not found after the anchor. |
| 191 | match_index = _find_context(lines, old_lines, start - 1, anchored=True) |
| 192 | lines[match_index : match_index + len(old_lines)] = new_lines |
| 193 | return _HunkApplyResult( |
| 194 | cursor=match_index + len(new_lines), |
| 195 | line_from=match_index + 1, |
| 196 | line_to=match_index + max(len(new_lines), 1), |
| 197 | ) |
| 198 | |
| 199 | |
| 200 | def _find_anchor(lines: list[str], anchor: str, start: int) -> int: |
| 201 | for index in range(start, len(lines)): |
| 202 | if lines[index] == anchor: |
| 203 | return index + 1 |
| 204 | stripped_anchor = anchor.strip() |
| 205 | for index in range(start, len(lines)): |
| 206 | if lines[index].strip() == stripped_anchor: |
| 207 | return index + 1 |
| 208 | raise ContextPatchError(f"anchor not found: {anchor}") |
| 209 | |
| 210 | |
| 211 | def _find_context( |
| 212 | lines: list[str], |
| 213 | context: list[str], |
| 214 | start: int, |
| 215 | *, |
| 216 | anchored: bool, |
| 217 | ) -> int: |
| 218 | matches = _matching_indexes(lines, context, start, mode="exact") |
| 219 | if not matches: |
| 220 | matches = _matching_indexes(lines, context, start, mode="rstrip") |
| 221 | |
| 222 | if not matches: |
| 223 | preview = "\n".join(context[:5]) |
| 224 | raise ContextPatchError(f"context not found:\n{preview}") |
| 225 | if len(matches) > 1 and not anchored: |
| 226 | preview = "\n".join(context[:5]) |
| 227 | raise ContextPatchError( |
| 228 | "context matched multiple locations; add an @@ anchor or more context:\n" |
| 229 | f"{preview}" |
| 230 | ) |
| 231 | return matches[0] |
| 232 | |
| 233 | |
| 234 | def _matching_indexes( |
| 235 | lines: list[str], |
| 236 | context: list[str], |
| 237 | start: int, |
| 238 | *, |
| 239 | mode: str, |
| 240 | ) -> list[int]: |
| 241 | if len(lines) < len(context): |
| 242 | return [] |
| 243 | |
| 244 | matches: list[int] = [] |
| 245 | for index in range(max(start, 0), len(lines) - len(context) + 1): |
| 246 | candidate = lines[index : index + len(context)] |
| 247 | if _lines_equal(candidate, context, mode=mode): |
| 248 | matches.append(index) |
| 249 | return matches |
| 250 | |
| 251 | |
| 252 | def _lines_equal(left: Iterable[str], right: Iterable[str], *, mode: str) -> bool: |
| 253 | if mode == "rstrip": |
| 254 | return [item.rstrip() for item in left] == [item.rstrip() for item in right] |
| 255 | return list(left) == list(right) |