From 0f18385f4c10399d6c2a032a1faba698dc7b9ff7 Mon Sep 17 00:00:00 2001 From: Michael Panchenko Date: Thu, 29 May 2025 14:10:12 +0200 Subject: [PATCH] Better handling of lines and columns at insertions --- scripts/demo_run_tools.py | 33 ++++++-- src/serena/symbol.py | 167 ++++++++++++++++++++++++++++++-------- 2 files changed, 160 insertions(+), 40 deletions(-) diff --git a/scripts/demo_run_tools.py b/scripts/demo_run_tools.py index 4d639b4..eafae78 100644 --- a/scripts/demo_run_tools.py +++ b/scripts/demo_run_tools.py @@ -10,17 +10,34 @@ from pprint import pprint from serena.agent import * -project_root = Path(__file__).parent.parent + +@dataclass +class InMemorySerenaConfig(SerenaConfigBase): + """ + In-memory implementation of Serena configuration with the GUI disabled. + """ + + gui_log_window_enabled: bool = False + web_dashboard: bool = False + if __name__ == "__main__": - agent = SerenaAgent(project_config=str(project_root / ".serena" / "project.yml")) - overview_tool = agent.get_tool(GetSymbolsOverviewTool) + project_path = str(Path("test") / "resources" / "repos" / "python" / "test_repo") + agent = SerenaAgent(project=project_path, serena_config=InMemorySerenaConfig()) find_symbol_tool = agent.get_tool(FindSymbolTool) - print("Getting an overview of the util package\n") - pprint(json.loads(overview_tool.apply("src/serena/util"))) + print("Finding the symbol 'VariableContainer'\n") + pprint(json.loads(find_symbol_tool.apply("VariableContainer", within_relative_path=str(Path("test_repo") / "variables.py")))) - print("\n\n") + # modifying with dry-run + symbol_manager = agent.symbol_manager - print("Finding the symbol 'SerenaAgent'\n") - pprint(json.loads(find_symbol_tool.apply("SerenaAgent"))) + code_diff = symbol_manager.insert_after_symbol( + "VariableContainer/modify_instance_var", + relative_file_path=str(Path("test_repo") / "variables.py"), + body="test test\nsecond line", + dry_run=True, + ) + print(f"Deleted lines: {code_diff.deleted_lines}") + print(f"Added lines: {code_diff.added_lines}") + print(f"Edited module: {code_diff.modified_content}") diff --git a/src/serena/symbol.py b/src/serena/symbol.py index 578c5dc..201cca7 100644 --- a/src/serena/symbol.py +++ b/src/serena/symbol.py @@ -3,7 +3,6 @@ import logging import os from collections.abc import Iterator, Sequence from contextlib import contextmanager, nullcontext -from copy import copy from dataclasses import asdict, dataclass, field from difflib import SequenceMatcher from typing import TYPE_CHECKING, Any, Literal, NamedTuple, Self, Union, overload @@ -690,10 +689,37 @@ class SymbolManager: return None @overload - def insert_after_symbol(self, name_path: str, relative_file_path: str, body: str, *, dry_run: Literal[False] = False) -> None: ... + def insert_after_symbol( + self, + name_path: str, + relative_file_path: str, + body: str, + *, + use_same_indentation: bool = True, + at_new_line: bool = True, + dry_run: Literal[False] = False, + ) -> None: ... @overload - def insert_after_symbol(self, name_path: str, relative_file_path: str, body: str, *, dry_run: Literal[True]) -> CodeDiff: ... - def insert_after_symbol(self, name_path: str, relative_file_path: str, body: str, *, dry_run: bool = False) -> CodeDiff | None: + def insert_after_symbol( + self, + name_path: str, + relative_file_path: str, + body: str, + *, + use_same_indentation: bool = True, + at_new_line: bool = True, + dry_run: Literal[True], + ) -> CodeDiff: ... + def insert_after_symbol( + self, + name_path: str, + relative_file_path: str, + body: str, + *, + use_same_indentation: bool = True, + at_new_line: bool = True, + dry_run: bool = False, + ) -> CodeDiff | None: """ Inserts content after the symbol with the given name in the given file. """ @@ -707,13 +733,25 @@ class SymbolManager: f"Found symbols at locations: \n" + json.dumps([s.location.to_dict() for s in symbol_candidates], indent=2) ) symbol = symbol_candidates[-1] - return self.insert_after_symbol_at_location(symbol.location, body, dry_run=dry_run) # type: ignore[call-overload] + return self.insert_after_symbol_at_location(symbol.location, body, at_new_line=at_new_line, use_same_indentation=use_same_indentation, dry_run=dry_run) # type: ignore[call-overload] @overload - def insert_after_symbol_at_location(self, location: SymbolLocation, body: str, *, dry_run: Literal[False] = False) -> None: ... + def insert_after_symbol_at_location( + self, + location: SymbolLocation, + body: str, + *, + at_new_line: bool = True, + use_same_indentation: bool = True, + dry_run: Literal[False] = False, + ) -> None: ... @overload - def insert_after_symbol_at_location(self, location: SymbolLocation, body: str, *, dry_run: Literal[True]) -> CodeDiff: ... - def insert_after_symbol_at_location(self, location: SymbolLocation, body: str, *, dry_run: bool = False) -> CodeDiff | None: + def insert_after_symbol_at_location( + self, location: SymbolLocation, body: str, *, at_new_line: bool = True, use_same_indentation: bool = True, dry_run: Literal[True] + ) -> CodeDiff: ... + def insert_after_symbol_at_location( + self, location: SymbolLocation, body: str, *, at_new_line: bool = True, use_same_indentation: bool = True, dry_run: bool = False + ) -> CodeDiff | None: """ Appends content after the given symbol @@ -724,8 +762,6 @@ class SymbolManager: # make sure body always ends with at least one newline if not body.endswith("\n"): body += "\n" - if not body.startswith("\n"): - body = "\n" + body assert location.relative_path is not None @@ -738,22 +774,74 @@ class SymbolManager: if pos is None: raise ValueError(f"Symbol at {location} does not have a defined end position.") + line, col = pos["line"], pos["character"] + if use_same_indentation: + symbol_start_pos = symbol.body_start_position + assert symbol_start_pos is not None, f"Symbol at {location=} does not have a defined start position." + symbol_identifier_col = symbol_start_pos["character"] + indent = " " * (symbol_identifier_col) + body = "\n".join(indent + line for line in body.splitlines()) + if at_new_line: + line += 1 + # IMPORTANT: without this, the insertion does the wrong thing. See implementation of insert_text_at_position in TextUtils, + # it is somewhat counterintuitive (never inserts whitespace) + # I am not 100% sure whether col=0 is always the best choice here. + # + # Without col=0, inserting after dataclass_instance in variables.py: + # > dataclass_instance = VariableDataclass(id=1, name="Test") + # > test test + # > dataclass_instancetest test + # > second line + # > .status = "active" # Reassign dataclass field + # + # With col=0: + # > dataclass_instance = VariableDataclass(id=1, name="Test") + # > test test + # > second line + # > dataclass_instance.status = "active" # Reassign dataclass field + col = 0 + if dry_run: original_content = self._get_code_file_content(location.relative_path) - modified_content, _, _ = TextUtils.insert_text_at_position(original_content, pos["line"], pos["character"], body) + modified_content, _, _ = TextUtils.insert_text_at_position(original_content, line=line, col=col, text_to_be_inserted=body) return CodeDiff(relative_path=location.relative_path, original_content=original_content, modified_content=modified_content) else: - # The _edited_symbol_location context manager handles LSP notifications like didOpen. - # We use the pre-calculated 'pos' for the insertion. with self._edited_symbol_location(location): - self.lang_server.insert_text_at_position(location.relative_path, pos["line"], pos["character"], body) + self.lang_server.insert_text_at_position(location.relative_path, line=line, column=col, text_to_be_inserted=body) return None @overload - def insert_before_symbol(self, name_path: str, relative_file_path: str, body: str, *, dry_run: Literal[False] = False) -> None: ... + def insert_before_symbol( + self, + name_path: str, + relative_file_path: str, + body: str, + *, + at_new_line: bool = True, + use_same_indentation: bool = True, + dry_run: Literal[False] = False, + ) -> None: ... @overload - def insert_before_symbol(self, name_path: str, relative_file_path: str, body: str, *, dry_run: Literal[True]) -> CodeDiff: ... - def insert_before_symbol(self, name_path: str, relative_file_path: str, body: str, *, dry_run: bool = False) -> CodeDiff | None: + def insert_before_symbol( + self, + name_path: str, + relative_file_path: str, + body: str, + *, + at_new_line: bool = True, + use_same_indentation: bool = True, + dry_run: Literal[True], + ) -> CodeDiff: ... + def insert_before_symbol( + self, + name_path: str, + relative_file_path: str, + body: str, + *, + at_new_line: bool = True, + use_same_indentation: bool = True, + dry_run: bool = False, + ) -> CodeDiff | None: """ Inserts content before the symbol with the given name in the given file. """ @@ -767,13 +855,25 @@ class SymbolManager: f"Found symbols at locations: \n" + json.dumps([s.location.to_dict() for s in symbol_candidates], indent=2) ) symbol = symbol_candidates[0] - return self.insert_before_symbol_at_location(symbol.location, body, dry_run=dry_run) # type: ignore[call-overload] + return self.insert_before_symbol_at_location(symbol.location, body, at_new_line=at_new_line, use_same_indentation=use_same_indentation, dry_run=dry_run) # type: ignore[call-overload] @overload - def insert_before_symbol_at_location(self, location: SymbolLocation, body: str, *, dry_run: Literal[False] = False) -> None: ... + def insert_before_symbol_at_location( + self, + location: SymbolLocation, + body: str, + *, + at_new_line: bool = True, + use_same_indentation: bool = True, + dry_run: Literal[False] = False, + ) -> None: ... @overload - def insert_before_symbol_at_location(self, location: SymbolLocation, body: str, *, dry_run: Literal[True]) -> CodeDiff: ... - def insert_before_symbol_at_location(self, location: SymbolLocation, body: str, *, dry_run: bool = False) -> CodeDiff | None: + def insert_before_symbol_at_location( + self, location: SymbolLocation, body: str, *, at_new_line: bool = True, use_same_indentation: bool = True, dry_run: Literal[True] + ) -> CodeDiff: ... + def insert_before_symbol_at_location( + self, location: SymbolLocation, body: str, *, at_new_line: bool = True, use_same_indentation: bool = True, dry_run: bool = False + ) -> CodeDiff | None: """ Inserts content before the given symbol @@ -781,25 +881,28 @@ class SymbolManager: :param body: the body of the entity to insert :param dry_run: if True, return a CodeDiff instead of modifying the file """ - # make sure body always ends with at least one newline - if not body.endswith("\n"): - body += "\n" - if not body.startswith("\n"): - body = "\n" + body - with self._edited_symbol_location(location) as symbol: - original_start_pos = symbol.body_start_position - if original_start_pos is None: + symbol_start_pos = symbol.body_start_position + if symbol_start_pos is None: raise ValueError(f"Symbol at {location} does not have a defined start position.") - pos = copy(original_start_pos) + line = symbol_start_pos["line"] + col = symbol_start_pos["character"] + if use_same_indentation: + indent = " " * (col) + body = "\n".join(indent + line for line in body.splitlines()) + + # similar problems as in insert_after_symbol_at_location, see comment there + if at_new_line: + col = 0 + line -= 1 assert location.relative_path is not None if dry_run: original_content = self._get_code_file_content(location.relative_path) - modified_content, _, _ = TextUtils.insert_text_at_position(original_content, pos["line"], pos["character"], body) + modified_content, _, _ = TextUtils.insert_text_at_position(original_content, line=line, col=col, text_to_be_inserted=body) return CodeDiff(relative_path=location.relative_path, original_content=original_content, modified_content=modified_content) else: - self.lang_server.insert_text_at_position(location.relative_path, pos["line"], pos["character"], body) + self.lang_server.insert_text_at_position(location.relative_path, line=line, column=col, text_to_be_inserted=body) return None @overload