diff --git a/pr_split/diff_ops/reconstructor.py b/pr_split/diff_ops/reconstructor.py index 74d57a4..e645afe 100644 --- a/pr_split/diff_ops/reconstructor.py +++ b/pr_split/diff_ops/reconstructor.py @@ -124,35 +124,43 @@ def apply_hunks(base_content: str, patch_file: PatchedFile, assigned_indices: li return "".join(lines) +def _assigned_hunk_indices( + patch_file: PatchedFile, assignments: list[GroupAssignment] +) -> list[int]: + """Union of the hunks every assignment for one file claims.""" + covered: set[int] = set() + for assignment in assignments: + if assignment.assignment_type is AssignmentType.WHOLE_FILE: + covered.update(range(len(patch_file))) + else: + covered.update(assignment.hunk_indices) + return sorted(covered) + + def materialize_group_files( parsed_diff: ParsedDiff, group: Group, ref: str ) -> dict[str, str | None]: - logger.info(logs.MATERIALIZING_FILES.format(count=len(group.assignments), group=group.id)) pf_map = {pf.path: pf for pf in parsed_diff.patch_set} - result: dict[str, str | None] = {} + # Several assignments may name the same file (e.g. merged across diff + # chunks); each file is written once from the union of their hunks. + assignments_by_path: dict[str, list[GroupAssignment]] = {} for assignment in group.assignments: - patch_file = pf_map.get(assignment.file_path) + assignments_by_path.setdefault(assignment.file_path, []).append(assignment) + logger.info(logs.MATERIALIZING_FILES.format(count=len(assignments_by_path), group=group.id)) + result: dict[str, str | None] = {} + for file_path, assignments in assignments_by_path.items(): + patch_file = pf_map.get(file_path) if patch_file is None: continue if patch_file.is_removed_file: - result[assignment.file_path] = None + result[file_path] = None continue + indices = _assigned_hunk_indices(patch_file, assignments) if patch_file.is_added_file: - all_indices = list(range(len(patch_file))) - match assignment.assignment_type: - case AssignmentType.WHOLE_FILE: - indices = all_indices - case AssignmentType.PARTIAL_HUNKS: - indices = assignment.hunk_indices - result[assignment.file_path] = "".join( + result[file_path] = "".join( "".join(_hunk_target_lines(patch_file[idx])) for idx in indices ) continue - base_content = _get_base_file_content(assignment.file_path, ref) - match assignment.assignment_type: - case AssignmentType.WHOLE_FILE: - indices = list(range(len(patch_file))) - case AssignmentType.PARTIAL_HUNKS: - indices = assignment.hunk_indices - result[assignment.file_path] = apply_hunks(base_content, patch_file, indices) + base_content = _get_base_file_content(file_path, ref) + result[file_path] = apply_hunks(base_content, patch_file, indices) return result diff --git a/pr_split/planner/client.py b/pr_split/planner/client.py index 0c933c8..68531df 100644 --- a/pr_split/planner/client.py +++ b/pr_split/planner/client.py @@ -227,12 +227,38 @@ def _parse_groups(raw: RawToolOutput) -> list[Group]: return groups +def _merge_assignment(existing: GroupAssignment, incoming: GroupAssignment) -> GroupAssignment: + if ( + existing.assignment_type is AssignmentType.WHOLE_FILE + or incoming.assignment_type is AssignmentType.WHOLE_FILE + ): + assignment_type = AssignmentType.WHOLE_FILE + else: + assignment_type = AssignmentType.PARTIAL_HUNKS + return GroupAssignment( + file_path=existing.file_path, + assignment_type=assignment_type, + hunk_indices=sorted(set(existing.hunk_indices) | set(incoming.hunk_indices)), + ) + + def _merge_chunk_groups(accumulated: list[Group], chunk_groups: list[Group]) -> list[Group]: acc_map = {g.id: g for g in accumulated} for cg in chunk_groups: if cg.id in acc_map: existing = acc_map[cg.id] - existing.assignments.extend(cg.assignments) + # A later chunk may assign more hunks of a file this group already + # holds. Keep one assignment per path: materialization writes a + # file once per assignment, so duplicates would drop hunks. + by_path: dict[str, GroupAssignment] = {} + for assignment in [*existing.assignments, *cg.assignments]: + if assignment.file_path in by_path: + by_path[assignment.file_path] = _merge_assignment( + by_path[assignment.file_path], assignment + ) + else: + by_path[assignment.file_path] = assignment + existing.assignments = list(by_path.values()) for dep in cg.depends_on: if dep not in existing.depends_on: existing.depends_on.append(dep) diff --git a/tests/test_client.py b/tests/test_client.py index 1bdfffa..a685f91 100644 --- a/tests/test_client.py +++ b/tests/test_client.py @@ -986,3 +986,105 @@ def test_large_diff_uses_chunking( result = _plan_split_with_llm(parsed, settings) assert len(result) == 1 mock_chunked.assert_called_once() + + +class TestMergeChunkGroupsSameFile: + def test_same_file_in_later_chunk_is_merged_into_one_assignment(self) -> None: + first = Group( + id="pr-1", + title="t", + description="d", + assignments=[ + GroupAssignment( + file_path="f.py", + assignment_type=AssignmentType.PARTIAL_HUNKS, + hunk_indices=[0], + ) + ], + ) + later = Group( + id="pr-1", + title="t", + description="d", + assignments=[ + GroupAssignment( + file_path="f.py", + assignment_type=AssignmentType.PARTIAL_HUNKS, + hunk_indices=[1], + ) + ], + ) + (merged,) = _merge_chunk_groups([first], [later]) + assert len(merged.assignments) == 1 + assert merged.assignments[0].hunk_indices == [0, 1] + assert merged.assignments[0].assignment_type == AssignmentType.PARTIAL_HUNKS + + def test_whole_file_wins_when_merging(self) -> None: + partial = GroupAssignment( + file_path="f.py", assignment_type=AssignmentType.PARTIAL_HUNKS, hunk_indices=[0] + ) + whole = GroupAssignment( + file_path="f.py", assignment_type=AssignmentType.WHOLE_FILE, hunk_indices=[] + ) + first = Group(id="pr-1", title="t", description="d", assignments=[partial]) + later = Group(id="pr-1", title="t", description="d", assignments=[whole]) + (merged,) = _merge_chunk_groups([first], [later]) + assert len(merged.assignments) == 1 + assert merged.assignments[0].assignment_type == AssignmentType.WHOLE_FILE + assert merged.assignments[0].hunk_indices == [0] + + def test_other_files_untouched(self) -> None: + first = Group( + id="pr-1", + title="t", + description="d", + assignments=[ + GroupAssignment( + file_path="a.py", + assignment_type=AssignmentType.PARTIAL_HUNKS, + hunk_indices=[0], + ) + ], + ) + later = Group( + id="pr-1", + title="t", + description="d", + assignments=[ + GroupAssignment( + file_path="b.py", + assignment_type=AssignmentType.PARTIAL_HUNKS, + hunk_indices=[2], + ) + ], + ) + (merged,) = _merge_chunk_groups([first], [later]) + assert [(a.file_path, a.hunk_indices) for a in merged.assignments] == [ + ("a.py", [0]), + ("b.py", [2]), + ] + + def test_pre_existing_duplicates_in_accumulated_group_are_kept(self) -> None: + def _pa(path: str, hunks: list[int]) -> GroupAssignment: + return GroupAssignment( + file_path=path, + assignment_type=AssignmentType.PARTIAL_HUNKS, + hunk_indices=hunks, + ) + + first = Group( + id="pr-1", title="t", description="d", assignments=[_pa("f.py", [0]), _pa("f.py", [1])] + ) + later = Group(id="pr-1", title="t", description="d", assignments=[_pa("g.py", [0])]) + (merged,) = _merge_chunk_groups([first], [later]) + assert [(a.file_path, a.hunk_indices) for a in merged.assignments] == [ + ("f.py", [0, 1]), + ("g.py", [0]), + ] + + first = Group( + id="pr-1", title="t", description="d", assignments=[_pa("f.py", [0]), _pa("f.py", [1])] + ) + later = Group(id="pr-1", title="t", description="d", assignments=[_pa("f.py", [2])]) + (merged,) = _merge_chunk_groups([first], [later]) + assert [(a.file_path, a.hunk_indices) for a in merged.assignments] == [("f.py", [0, 1, 2])] diff --git a/tests/test_reconstructor.py b/tests/test_reconstructor.py index 71b4480..42b39ae 100644 --- a/tests/test_reconstructor.py +++ b/tests/test_reconstructor.py @@ -382,6 +382,66 @@ def test_added_file_content_is_not_double_spaced(self) -> None: assert result["new_file.py"] == 'def hello():\n return "world"\n\n' +class TestMaterializeDuplicateAssignments: + @patch("pr_split.diff_ops.reconstructor._get_base_file_content") + def test_two_assignments_for_one_file_apply_both_hunks(self, mock_base: MagicMock) -> None: + mock_base.return_value = _base_content() + parsed = parse_diff(PATCH_TEXT) + group = Group( + id="pr-1", + title="t", + description="d", + assignments=[ + GroupAssignment( + file_path="example.py", + assignment_type=AssignmentType.PARTIAL_HUNKS, + hunk_indices=[0], + ), + GroupAssignment( + file_path="example.py", + assignment_type=AssignmentType.PARTIAL_HUNKS, + hunk_indices=[1], + ), + ], + ) + with patch("pr_split.diff_ops.reconstructor.logger.info") as mock_log: + result = materialize_group_files(parsed, group, "main") + content = result["example.py"] + assert content is not None + assert "inserted_after_1" in content + assert "inserted_after_11" in content + assert mock_base.call_count == 1 + assert "Materializing 1 file" in mock_log.call_args[0][0] + + def test_duplicate_assignments_on_new_file_apply_both_hunks(self) -> None: + # Two hunks in a new file (unidiff splits them when the context gap + # is large enough), each claimed by a separate PARTIAL assignment. + parsed = parse_diff( + "diff --git a/n.py b/n.py\nnew file mode 100644\n--- /dev/null\n+++ b/n.py\n" + "@@ -0,0 +1,2 @@\n+one\n+two\n" + "@@ -0,0 +10,1 @@\n+ten\n" + ) + assert len(parsed.patch_set[0]) == 2 + group = Group( + id="pr-1", + title="t", + description="d", + assignments=[ + GroupAssignment( + file_path="n.py", + assignment_type=AssignmentType.PARTIAL_HUNKS, + hunk_indices=[0], + ), + GroupAssignment( + file_path="n.py", + assignment_type=AssignmentType.PARTIAL_HUNKS, + hunk_indices=[1], + ), + ], + ) + assert materialize_group_files(parsed, group, "main")["n.py"] == "one\ntwo\nten\n" + + class TestMissingTrailingNewline: NO_EOL_PATCH = """\ --- a/f.txt