Implement JetBrainsCodeEditor

src/serena/code_editor.py

@@ -1,18 +1,24 @@
 import json
+import logging
 import os
 from abc import ABC, abstractmethod
 from collections.abc import Iterable, Iterator, Reversible
 from contextlib import contextmanager
 from typing import TYPE_CHECKING, Generic, Optional, TypeVar
 
-from serena.symbol import AbstractSymbol, PositionInFile, Symbol, SymbolManager
+from serena.symbol import AbstractSymbol, JetBrainsSymbol, PositionInFile, Symbol, SymbolManager
 from solidlsp import SolidLanguageServer
 from solidlsp.ls import LSPFileBuffer
+from solidlsp.ls_utils import TextUtils
+
+from .config.serena_config import Project
+from .tools import JetBrainsPluginClient
 
 if TYPE_CHECKING:
     from .agent import SerenaAgent
 
 
+log = logging.getLogger(__name__)
 TSymbol = TypeVar("TSymbol", bound=AbstractSymbol)
 
 
@@ -248,3 +254,44 @@ class LanguageServerCodeEditor(CodeEditor[Symbol]):
                 "Their locations are: \n " + json.dumps([s.location.to_dict() for s in symbol_candidates], indent=2)
             )
         return symbol_candidates[0]
+
+
+class JetBrainsCodeEditor(CodeEditor[JetBrainsSymbol]):
+    def __init__(self, project: Project, agent: Optional["SerenaAgent"] = None) -> None:
+        self._project = project
+        super().__init__(project_root=project.project_root, agent=agent)
+
+    class EditedFile(CodeEditor.EditedFile):
+        def __init__(self, relative_path: str, project: Project):
+            path = os.path.join(project.project_root, relative_path)
+            log.info("Editing file: %s", path)
+            with open(path, encoding=project.project_config.encoding) as f:
+                self._content = f.read()
+
+        def get_contents(self) -> str:
+            return self._content
+
+        def delete_text_between_positions(self, start_pos: PositionInFile, end_pos: PositionInFile) -> None:
+            self._content, _ = TextUtils.delete_text_between_positions(
+                self._content, start_pos.line, start_pos.col, end_pos.line, end_pos.col
+            )
+
+        def insert_text_at_position(self, pos: PositionInFile, text: str) -> None:
+            self._content, _, _ = TextUtils.insert_text_at_position(self._content, pos.line, pos.col, text)
+
+    @contextmanager
+    def _open_file_context(self, relative_path: str) -> Iterator["CodeEditor.EditedFile"]:
+        yield self.EditedFile(relative_path, self._project)
+
+    def _find_unique_symbol(self, name_path: str, relative_file_path: str) -> JetBrainsSymbol:
+        with JetBrainsPluginClient() as client:
+            result = client.find_symbol(name_path, relative_path=relative_file_path, include_body=False, depth=0, include_location=True)
+            symbols = result["symbols"]
+            if not symbols:
+                raise ValueError(f"No symbol with name {name_path} found in file {relative_file_path}")
+            if len(symbols) > 1:
+                raise ValueError(
+                    f"Found multiple {len(symbols)} symbols with name {name_path} in file {relative_file_path}. "
+                    "Their locations are: \n " + json.dumps([s["location"] for s in symbols], indent=2)
+                )
+            return JetBrainsSymbol(symbols[0], self._project)

src/serena/symbol.py

@@ -7,11 +7,15 @@ from dataclasses import asdict, dataclass, field
 from difflib import SequenceMatcher
 from typing import TYPE_CHECKING, Any, Literal, NamedTuple, Self, Union
 
+from sensai.util.datastruct import Maybe
 from sensai.util.string import ToStringMixin
 
 from solidlsp import SolidLanguageServer
 from solidlsp.ls import ReferenceInSymbol as LSPReferenceInSymbol
 from solidlsp.ls_types import Position, SymbolKind, UnifiedSymbolInformation
+from solidlsp.ls_utils import TextUtils
+
+from .config.serena_config import Project
 
 if TYPE_CHECKING:
     from .agent import SerenaAgent
@@ -697,3 +701,55 @@ class SymbolManager:
     def _get_code_file_content(self, relative_path: str) -> str:
         """Get the content of a file using the language server."""
         return self._lang_server.language_server.retrieve_full_file_content(relative_path)
+
+
+class JetBrainsSymbol(AbstractSymbol):
+    def __init__(self, symbol_dict: dict, project: Project) -> None:
+        """
+        :param symbol_dict: dictionary as returned by the JetBrains plugin client.
+        """
+        self._project = project
+        self._dict = symbol_dict
+        self._cached_file_content: str | None = None
+        self._cached_body_start_position: PositionInFile | None = None
+        self._cached_body_end_position: PositionInFile | None = None
+
+    def get_relative_path(self) -> str:
+        return self._dict["relative_path"]
+
+    def get_file_content(self) -> str:
+        if self._cached_file_content is None:
+            path = os.path.join(self._project.project_root, self.get_relative_path())
+            with open(path, encoding=self._project.project_config.encoding) as f:
+                self._cached_file_content = f.read()
+        return self._cached_file_content
+
+    def is_position_in_file_available(self) -> bool:
+        return "text_range" in self._dict
+
+    @property
+    def _body_indices(self) -> tuple[int, int]:
+        text_range = self._dict["text_range"]
+        return text_range["start_offset"], text_range["end_offset"]
+
+    def get_body_start_position(self) -> PositionInFile | None:
+        if not self.is_position_in_file_available():
+            return None
+        if self._cached_body_start_position is None:
+            index = self._dict["text_range"]["start_offset"]
+            line, col = TextUtils.get_line_col_from_index(self.get_file_content(), index)
+            self._cached_body_start_position = PositionInFile(line=line, col=col)
+        return self._cached_body_start_position
+
+    def get_body_end_position(self) -> PositionInFile | None:
+        if not self.is_position_in_file_available():
+            return None
+        if self._cached_body_end_position is None:
+            index = self._dict["text_range"]["end_offset"]
+            line, col = TextUtils.get_line_col_from_index(self.get_file_content(), index)
+            self._cached_body_end_position = PositionInFile(line=line, col=col)
+        return self._cached_body_end_position
+
+    def is_neighbouring_definition_separated_by_empty_line(self) -> bool:
+        # NOTE: Symbol types cannot really be differentiated, because types are not handled in a language-agnostic way.
+        return False

src/serena/tools/jetbrains_plugin_client.py

@@ -87,7 +87,9 @@ class JetBrainsPluginClient:
     def heartbeat(self) -> dict[str, Any]:
         return self._make_request("GET", "/heartbeat")
 
-    def find_symbol(self, name_path: str, relative_path: str | None = None, include_body: bool = False, depth: int = 0) -> dict[str, Any]:
+    def find_symbol(
+        self, name_path: str, relative_path: str | None = None, include_body: bool = False, depth: int = 0, include_location: bool = False
+    ) -> dict[str, Any]:
         """
         Find symbols by name.
 
@@ -98,7 +100,13 @@ class JetBrainsPluginClient:
 
         :return: Dictionary containing 'symbols' list with matching symbols
         """
-        request_data = {"namePath": name_path, "relativePath": relative_path, "includeBody": include_body, "depth": depth}
+        request_data = {
+            "namePath": name_path,
+            "relativePath": relative_path,
+            "includeBody": include_body,
+            "depth": depth,
+            "includeLocation": include_location,
+        }
         return self._make_request("POST", "/findSymbol", request_data)
 
     def find_references(self, name_path: str, relative_path: str) -> dict[str, Any]: