diff --git a/src/serena/tools/jetbrains_plugin_client.py b/src/serena/tools/jetbrains_plugin_client.py index cd31b7b1..37d10a4d 100644 --- a/src/serena/tools/jetbrains_plugin_client.py +++ b/src/serena/tools/jetbrains_plugin_client.py @@ -181,6 +181,54 @@ class JetBrainsPluginClient(ToStringMixin): request_data = {"relativePath": relative_path} return self._make_request("POST", "/getSymbolsOverview", request_data) + def get_supertypes( + self, + name_path: str, + relative_path: str, + depth: int | None = None, + limit_children: int | None = None, + ) -> dict[str, Any]: + """ + Gets the supertypes (parent classes/interfaces) of a symbol. + + :param name_path: the name path of the symbol + :param relative_path: the relative path to the file containing the symbol + :param depth: depth limit for hierarchy traversal (None or 0 for unlimited) + :param limit_children: optional limit on children per level + :return: dictionary containing symbol info and type hierarchy + """ + request_data = { + "namePath": name_path, + "relativePath": relative_path, + "depth": depth, + "limitChildren": limit_children, + } + return self._make_request("POST", "/getSupertypes", request_data) + + def get_subtypes( + self, + name_path: str, + relative_path: str, + depth: int | None = None, + limit_children: int | None = None, + ) -> dict[str, Any]: + """ + Gets the subtypes (subclasses/implementations) of a symbol. + + :param name_path: the name path of the symbol + :param relative_path: the relative path to the file containing the symbol + :param depth: depth limit for hierarchy traversal (None or 0 for unlimited) + :param limit_children: optional limit on children per level + :return: dictionary containing symbol info and type hierarchy + """ + request_data = { + "namePath": name_path, + "relativePath": relative_path, + "depth": depth, + "limitChildren": limit_children, + } + return self._make_request("POST", "/getSubtypes", request_data) + def rename_symbol( self, name_path: str, relative_path: str, new_name: str, rename_in_comments: bool, rename_in_text_occurrences: bool ) -> None: diff --git a/src/serena/tools/jetbrains_tools.py b/src/serena/tools/jetbrains_tools.py index e2b84e6e..7fa4b553 100644 --- a/src/serena/tools/jetbrains_tools.py +++ b/src/serena/tools/jetbrains_tools.py @@ -1,3 +1,5 @@ +from typing import Literal + from serena.tools import Tool, ToolMarkerOptional, ToolMarkerSymbolicRead from serena.tools.jetbrains_plugin_client import JetBrainsPluginClient @@ -121,3 +123,98 @@ class JetBrainsGetSymbolsOverviewTool(Tool, ToolMarkerSymbolicRead, ToolMarkerOp ) result = self._to_json(response_dict) return self._limit_length(result, max_answer_chars) + + +class JetBrainsTypeHierarchyTool(Tool, ToolMarkerSymbolicRead, ToolMarkerOptional): + """ + Retrieves the type hierarchy (supertypes and/or subtypes) of a symbol using the JetBrains backend + """ + + @staticmethod + def _transform_hierarchy_nodes(nodes: list[dict] | None) -> dict[str, list]: + """ + Transform a list of TypeHierarchyNode into a file-grouped compact format. + + Returns a dict where keys are relative_paths and values are lists of either: + - "SymbolNamePath" (leaf node) + - {"SymbolNamePath": {nested_file_grouped_children}} (node with children) + """ + if not nodes: + return {} + + result: dict[str, list] = {} + + for node in nodes: + symbol = node["symbol"] + name_path = symbol["name_path"] + rel_path = symbol["relative_path"] + children = node.get("children", []) + + if rel_path not in result: + result[rel_path] = [] + + if children: + # Node with children - recurse + nested = JetBrainsTypeHierarchyTool._transform_hierarchy_nodes(children) + result[rel_path].append({name_path: nested}) + else: + # Leaf node + result[rel_path].append(name_path) + + return result + + def apply( + self, + name_path: str, + relative_path: str, + hierarchy_type: Literal["super", "sub", "both"] = "sub", + depth: int | None = 1, + max_answer_chars: int = -1, + ) -> str: + """ + Gets the type hierarchy of a symbol (supertypes, subtypes, or both). + + :param name_path: name path of the symbol for which to get the type hierarchy. + :param relative_path: the relative path to the file containing the symbol. + :param hierarchy_type: which hierarchy to retrieve: "super" for parent classes/interfaces, + "sub" for subclasses/implementations, or "both" for both directions. + :param depth: depth limit for hierarchy traversal (None or 0 for unlimited). + :param max_answer_chars: max characters for the JSON result. If exceeded, no content is returned. + -1 means the default value from the config will be used. + :return: Compact JSON with file-grouped hierarchy. Error string if not applicable. + """ + with JetBrainsPluginClient.from_project(self.project) as client: + subtypes = {} + supertypes = {} + if hierarchy_type in ("super", "both"): + supertypes_response = client.get_supertypes( + name_path=name_path, + relative_path=relative_path, + depth=depth, + ) + if error_msg := supertypes_response.get("error"): + return error_msg + + symbol = supertypes_response["symbol"] + supertypes = self._transform_hierarchy_nodes(supertypes_response.get("hierarchy")) + if hierarchy_type in ("sub", "both"): + subtypes_response = client.get_subtypes( + name_path=name_path, + relative_path=relative_path, + depth=depth, + ) + if error_msg := subtypes_response.get("error"): + return error_msg + symbol = subtypes_response["symbol"] + + subtypes = self._transform_hierarchy_nodes(subtypes_response.get("hierarchy")) + + symbol_compact = {symbol["relative_path"]: [symbol["name_path"]]} + result_dict: dict[str, dict | list] = {"symbol": symbol_compact} + if supertypes: + result_dict["supertypes"] = supertypes + if subtypes: + result_dict["subtypes"] = subtypes + + result = self._to_json(result_dict) + return self._limit_length(result, max_answer_chars)