From d41b9ec964425d88277de44ce01a51f44de2a65d Mon Sep 17 00:00:00 2001 From: Dominik Jain Date: Fri, 13 Feb 2026 22:31:32 +0100 Subject: [PATCH] Implement dict transformation in symbol overview tools using SymbolDictGrouper, adding support for the necessary transformations of the final dictionaries --- src/serena/symbol.py | 62 ++++++++++++++++++++++++++--- src/serena/tools/jetbrains_tools.py | 36 ++--------------- src/serena/tools/symbol_tools.py | 36 +---------------- 3 files changed, 61 insertions(+), 73 deletions(-) diff --git a/src/serena/symbol.py b/src/serena/symbol.py index ed0f15fc..e1050b46 100644 --- a/src/serena/symbol.py +++ b/src/serena/symbol.py @@ -802,13 +802,19 @@ class SymbolDictGrouper(Generic[TSymbolDict], ABC): """ def __init__( - self, symbol_dict_type: type[TSymbolDict], children_key: Any, group_keys: list[Any], group_children_keys: list[Any] + self, + symbol_dict_type: type[TSymbolDict], + children_key: Any, + group_keys: list[Any], + group_children_keys: list[Any], + collapse_singleton: bool, ) -> None: """ :param symbol_dict_type: the TypedDict type that represents the type of the symbol dictionaries to be grouped :param children_key: the key in the symbol dictionaries that contains the list of child symbols. This is necessary for recursive grouping. :param group_keys: keys by which to group the symbol dictionaries. Must be a subset of the keys of `symbol_dict_type`. :param group_children_keys: keys by which to group the child symbol dictionaries. Must be a subset of the keys of `symbol_dict_type`. + :param collapse_singleton: whether to collapse dictionaries containing a single entry after regrouping to just the entry's value """ # symbol_dict_type must be a TypedDict type. # We check whether the type contains all the keys specified in `keys` and raise an error if not. @@ -821,27 +827,52 @@ class SymbolDictGrouper(Generic[TSymbolDict], ABC): self._children_key = children_key self._group_keys = group_keys self._group_children_keys = group_children_keys + self._collapse_singleton = collapse_singleton def _group_by(self, l: list[dict], keys: list[str], children_keys: list[str]) -> dict[str, Any]: - key = keys[0] + assert len(keys) > 0, "keys must not be empty" + # group by the first key grouped: dict[str, Any] = {} for item in l: - key_value = item.pop(key) + key_value = item.pop(keys[0], "unknown") if key_value not in grouped: grouped[key_value] = [] grouped[key_value].append(item) if len(keys) > 1: + # continue grouping by the remaining keys for k, group in grouped.items(): grouped[k] = self._group_by(group, keys[1:], children_keys) else: + # grouping is complete; now group the children if necessary if children_keys: for k, group in grouped.items(): for item in group: if self._children_key in item: children = item[self._children_key] item[self._children_key] = self._group_by(children, children_keys, children_keys) + # post-process final group items + grouped = {k: [self._transform_item(i) for i in v] for k, v in grouped.items()} return grouped + def _transform_item(self, item: dict) -> dict: + """ + Post-processes a final group item (which has been regrouped, i.e. some keys may have been removed), + collapsing singleton items (and items containing only a single non-children key) + """ + if self._collapse_singleton: + if len(item) == 1: + # {"name": "foo"} -> "foo" + # if there is only a single entry, collapse the dictionary to just the value of that entry + return next(iter(item.values())) + elif len(item) == 2 and self._children_key in item: + # {"name": "foo", "children": {...}} -> {"foo": {...}} + # if there are exactly two entries and one of them is the children key, + # convert to {other_value: children} if there are exactly two entries and one of them is the children key + other_key = next(k for k in item.keys() if k != self._children_key) + new_item = {item[other_key]: item[self._children_key]} + return new_item + return item + def group(self, symbols: list[TSymbolDict]) -> GroupedSymbolDict: """ :param symbols: the symbols to group @@ -855,10 +886,29 @@ class LanguageServerSymbolDictGrouper(SymbolDictGrouper[LanguageServerSymbol.Out self, group_keys: list[LanguageServerSymbol.OutputDictKey], group_children_keys: list[LanguageServerSymbol.OutputDictKey], + collapse_singleton: bool = False, ) -> None: - super().__init__(LanguageServerSymbol.OutputDict, "children", group_keys, group_children_keys) + super().__init__(LanguageServerSymbol.OutputDict, "children", group_keys, group_children_keys, collapse_singleton) class JetBrainsSymbolDictGrouper(SymbolDictGrouper[jb.SymbolDTO]): - def __init__(self, group_keys: list[jb.SymbolDTOKey], group_children_keys: list[jb.SymbolDTOKey]) -> None: - super().__init__(jb.SymbolDTO, "children", group_keys, group_children_keys) + def __init__( + self, + group_keys: list[jb.SymbolDTOKey], + group_children_keys: list[jb.SymbolDTOKey], + collapse_singleton: bool = False, + map_name_path_to_name: bool = False, + ) -> None: + super().__init__(jb.SymbolDTO, "children", group_keys, group_children_keys, collapse_singleton) + self._map_name_path_to_name = map_name_path_to_name + + def _transform_item(self, item: dict) -> dict: + if self._map_name_path_to_name: + # {"name_path: "Class/myMethod"} -> {"name: "myMethod"} + new_item = dict(item) + if "name_path" in item: + name_path = new_item.pop("name_path") + new_item["name"] = name_path.split("/")[-1] + return super()._transform_item(new_item) + else: + return super()._transform_item(item) diff --git a/src/serena/tools/jetbrains_tools.py b/src/serena/tools/jetbrains_tools.py index ed9b53cc..736fa435 100644 --- a/src/serena/tools/jetbrains_tools.py +++ b/src/serena/tools/jetbrains_tools.py @@ -1,9 +1,9 @@ import logging -from collections import defaultdict from typing import Any, Literal import serena.jetbrains.jetbrains_types as jb from serena.jetbrains.jetbrains_plugin_client import JetBrainsPluginClient +from serena.symbol import JetBrainsSymbolDictGrouper from serena.tools import Tool, ToolMarkerOptional, ToolMarkerSymbolicRead log = logging.getLogger(__name__) @@ -128,37 +128,7 @@ class JetBrainsGetSymbolsOverviewTool(Tool, ToolMarkerSymbolicRead, ToolMarkerOp """ USE_COMPACT_FORMAT = True - - @staticmethod - def _transform_symbols_to_compact_format(symbols: list[jb.SymbolDTO]) -> dict[str, list]: - """ - Transform symbol overview from verbose format to compact grouped format. - - Groups symbols by kind and uses names instead of full symbol objects. - For symbols with children, creates nested dictionaries. - - The name_path can be inferred from the hierarchical structure: - - Top-level symbols: name_path = name - - Nested symbols: name_path = parent_name + "/" + name - For example, "convert" under class "ProjectType" has name_path "ProjectType/convert". - """ - result = defaultdict(list) - - for symbol in symbols: - kind = symbol.get("type", "Unknown") - name_path = symbol["name_path"] - name = name_path.split("/")[-1] - children = symbol.get("children", []) - - if children: - # Symbol has children: create nested dict {name: children_dict} - children_dict = JetBrainsGetSymbolsOverviewTool._transform_symbols_to_compact_format(children) - result[kind].append({name: children_dict}) - else: - # Symbol has no children: just add the name - result[kind].append(name) # type: ignore - - return result + symbol_dict_grouper = JetBrainsSymbolDictGrouper(["type"], ["type"], collapse_singleton=True, map_name_path_to_name=True) def apply( self, @@ -187,7 +157,7 @@ class JetBrainsGetSymbolsOverviewTool(Tool, ToolMarkerSymbolicRead, ToolMarkerOp ) if self.USE_COMPACT_FORMAT: symbols = symbol_overview["symbols"] - result: dict[str, Any] = {"symbols": self._transform_symbols_to_compact_format(symbols)} + result: dict[str, Any] = {"symbols": self.symbol_dict_grouper.group(symbols)} documentation = symbol_overview.pop("documentation", None) if documentation: result["docstring"] = documentation diff --git a/src/serena/tools/symbol_tools.py b/src/serena/tools/symbol_tools.py index 2c876966..a22c1eef 100644 --- a/src/serena/tools/symbol_tools.py +++ b/src/serena/tools/symbol_tools.py @@ -3,7 +3,6 @@ Language server-related tools """ import os -from collections import defaultdict from collections.abc import Sequence from serena.symbol import LanguageServerSymbol, LanguageServerSymbolDictGrouper @@ -33,7 +32,7 @@ class GetSymbolsOverviewTool(Tool, ToolMarkerSymbolicRead): Gets an overview of the top-level symbols defined in a given file. """ - symbol_dict_grouper = LanguageServerSymbolDictGrouper(["kind"], ["kind"]) + symbol_dict_grouper = LanguageServerSymbolDictGrouper(["kind"], ["kind"], collapse_singleton=True) def apply(self, relative_path: str, depth: int = 0, max_answer_chars: int = -1) -> str: """ @@ -50,8 +49,7 @@ class GetSymbolsOverviewTool(Tool, ToolMarkerSymbolicRead): :return: a JSON object containing symbols grouped by kind in a compact format. """ result = self.get_symbol_overview(relative_path, depth=depth) - compact_result = self._transform_symbols_to_compact_format(result) - # compact_result = self.symbol_dict_grouper.group(result) + compact_result = self.symbol_dict_grouper.group(result) result_json_str = self._to_json(compact_result) return self._limit_length(result_json_str, max_answer_chars) @@ -91,36 +89,6 @@ class GetSymbolsOverviewTool(Tool, ToolMarkerSymbolicRead): ) return symbol_dicts - @staticmethod - def _transform_symbols_to_compact_format(symbols: list[LanguageServerSymbol.OutputDict]) -> dict[str, list]: - """ - Transform symbol overview from verbose format to compact grouped format. - - Groups symbols by kind and uses names instead of full symbol objects. - For symbols with children, creates nested dictionaries. - - The name_path can be inferred from the hierarchical structure: - - Top-level symbols: name_path = name - - Nested symbols: name_path = parent_name + "/" + name - For example, "convert" under class "ProjectType" has name_path "ProjectType/convert". - """ - result = defaultdict(list) - - for symbol in symbols: - kind = symbol.get("kind", "Unknown") - name = symbol.get("name", "unknown") - children = symbol.get("children", []) - - if children: - # Symbol has children: create nested dict {name: children_dict} - children_dict = GetSymbolsOverviewTool._transform_symbols_to_compact_format(children) - result[kind].append({name: children_dict}) - else: - # Symbol has no children: just add the name - result[kind].append(name) - - return result - class FindSymbolTool(Tool, ToolMarkerSymbolicRead): """