From 798a31b6a72447c437a1001c025597cf0c91052d Mon Sep 17 00:00:00 2001 From: Dominik Jain Date: Fri, 13 Feb 2026 18:17:05 +0100 Subject: [PATCH] Implement generic SymbolDictGrouper (and specialisations for LS & JB) It represents a reasonable compromise between pragmatism and type-safety. Soundness is checked at runtime, but if the grouper is instantiated statically, then we immediately detect misspecification. One example application is implemented in GetSymbolsOverviewTool Provides basis for #1021 --- src/serena/symbol.py | 97 +++++++++++++++++++++++++++-- src/serena/tools/jetbrains_types.py | 5 +- src/serena/tools/symbol_tools.py | 11 +++- 3 files changed, 104 insertions(+), 9 deletions(-) diff --git a/src/serena/symbol.py b/src/serena/symbol.py index e91856c1..9e7c8087 100644 --- a/src/serena/symbol.py +++ b/src/serena/symbol.py @@ -4,7 +4,7 @@ import os from abc import ABC, abstractmethod from collections.abc import Callable, Iterator, Sequence from dataclasses import asdict, dataclass -from typing import TYPE_CHECKING, Any, Self, Union +from typing import TYPE_CHECKING, Any, Generic, Literal, NotRequired, Self, TypedDict, TypeVar, Union from sensai.util.string import ToStringMixin @@ -395,6 +395,17 @@ class LanguageServerSymbol(Symbol, ToStringMixin): traverse(self) return result + class OutputDict(TypedDict): + name: str + name_path: str + location: NotRequired[dict[str, Any]] + body_location: NotRequired[dict[str, Any]] + body: NotRequired[str | None] + kind: NotRequired[str] + children: NotRequired[list["LanguageServerSymbol.OutputDict"]] + + OutputDictKey = Literal["name", "name_path", "location", "body_location", "body", "kind", "children"] + def to_dict( self, kind: bool = False, @@ -404,7 +415,7 @@ class LanguageServerSymbol(Symbol, ToStringMixin): include_children_body: bool = False, include_relative_path: bool = True, child_inclusion_predicate: Callable[[Self], bool] | None = None, - ) -> dict[str, Any]: + ) -> OutputDict: """ Converts the symbol to a dictionary. @@ -422,7 +433,7 @@ class LanguageServerSymbol(Symbol, ToStringMixin): should be included. :return: a dictionary representation of the symbol """ - result: dict[str, Any] = {"name": self.name, "name_path": self.get_name_path()} + result: LanguageServerSymbol.OutputDict = {"name": self.name, "name_path": self.get_name_path()} if kind: result["kind"] = self.kind @@ -441,7 +452,7 @@ class LanguageServerSymbol(Symbol, ToStringMixin): if child_inclusion_predicate is None: child_inclusion_predicate = lambda s: True - def included_children(s: Self) -> list[dict[str, Any]]: + def included_children(s: Self) -> list[LanguageServerSymbol.OutputDict]: children = [] for c in s.iter_children(): if not child_inclusion_predicate(c): @@ -689,7 +700,7 @@ class LanguageServerSymbolRetriever: return [ReferenceInLanguageServerSymbol.from_lsp_reference(r) for r in references] - def get_symbol_overview(self, relative_path: str, depth: int = 0) -> dict[str, list[dict]]: + def get_symbol_overview(self, relative_path: str, depth: int = 0) -> dict[str, list[LanguageServerSymbol.OutputDict]]: """ :param relative_path: the path of the file or directory for which to get the symbol overview :param depth: the depth up to which to include child symbols (0 = only top-level symbols) @@ -775,3 +786,79 @@ class JetBrainsSymbol(Symbol): 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 + + +TSymbolDict = TypeVar("TSymbolDict") +GroupedSymbolDict = dict[str, list[dict] | dict[str, dict]] + + +class SymbolDictGrouper(Generic[TSymbolDict], ABC): + """ + A utility class for grouping a list of symbol dictionaries by one or more specified keys. + + If an instance is statically initialised (upon module import), then this establishes a guarantee + that the specified keys are defined in the symbol dictionary type, ensuring at least basic type safety. + The respective ValueError will immediately be apparent. + """ + + def __init__( + self, symbol_dict_type: type[TSymbolDict], children_key: Any, group_keys: list[Any], group_children_keys: list[Any] + ) -> 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`. + """ + # 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. + if not hasattr(symbol_dict_type, "__annotations__"): + raise ValueError(f"symbol_dict_type must be a TypedDict type, got {symbol_dict_type}") + symbol_dict_keys = set(symbol_dict_type.__annotations__.keys()) + for key in group_keys + [children_key] + group_children_keys: + if key not in symbol_dict_keys: + raise ValueError(f"symbol_dict_type {symbol_dict_type} does not contain key '{key}'") + self._children_key = children_key + self._group_keys = group_keys + self._group_children_keys = group_children_keys + + def _group_by(self, l: list[dict], keys: list[str], children_keys: list[str]) -> dict[str, Any]: + key = keys[0] + grouped: dict[str, Any] = {} + for item in l: + key_value = item.pop(key) + if key_value not in grouped: + grouped[key_value] = [] + grouped[key_value].append(item) + if len(keys) > 1: + for k, group in grouped.items(): + grouped[k] = self._group_by(group, keys[1:], children_keys) + else: + 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) + return grouped + + def group(self, symbols: list[TSymbolDict]) -> GroupedSymbolDict: + """ + :param symbols: the symbols to group + :return: dictionary with the symbols grouped as defined at construction + """ + return self._group_by(symbols, self._group_keys, self._group_children_keys) # type: ignore + + +class LanguageServerSymbolDictGrouper(SymbolDictGrouper[LanguageServerSymbol.OutputDict]): + def __init__( + self, + group_keys: list[LanguageServerSymbol.OutputDictKey], + group_children_keys: list[LanguageServerSymbol.OutputDictKey], + ) -> None: + super().__init__(LanguageServerSymbol.OutputDict, "children", group_keys, group_children_keys) + + +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) diff --git a/src/serena/tools/jetbrains_types.py b/src/serena/tools/jetbrains_types.py index 80f02947..fd72bc60 100644 --- a/src/serena/tools/jetbrains_types.py +++ b/src/serena/tools/jetbrains_types.py @@ -1,4 +1,4 @@ -from typing import NotRequired, TypedDict +from typing import Literal, NotRequired, TypedDict class PluginStatusDTO(TypedDict): @@ -30,6 +30,9 @@ class SymbolDTO(TypedDict): num_usages: NotRequired[int] +SymbolDTOKey = Literal["name_path", "relative_path", "type", "body", "quick_info", "documentation", "text_range", "children", "num_usages"] + + class SymbolCollectionResponse(TypedDict): symbols: list[SymbolDTO] diff --git a/src/serena/tools/symbol_tools.py b/src/serena/tools/symbol_tools.py index ae74bbfd..59f43577 100644 --- a/src/serena/tools/symbol_tools.py +++ b/src/serena/tools/symbol_tools.py @@ -8,6 +8,7 @@ from collections.abc import Sequence from copy import copy from typing import Any +from serena.symbol import LanguageServerSymbol, LanguageServerSymbolDictGrouper from serena.tools import ( SUCCESS_RESULT, Tool, @@ -18,7 +19,8 @@ from serena.tools.tools_base import ToolMarkerOptional from solidlsp.ls_types import SymbolKind -def _sanitize_symbol_dict(symbol_dict: dict[str, Any]) -> dict[str, Any]: +# TODO: We have to get rid of this. +def _sanitize_symbol_dict(symbol_dict: LanguageServerSymbol.OutputDict) -> dict[str, Any]: """ Sanitize a symbol dictionary inplace by removing unnecessary information. """ @@ -50,6 +52,8 @@ class GetSymbolsOverviewTool(Tool, ToolMarkerSymbolicRead): Gets an overview of the top-level symbols defined in a given file. """ + symbol_dict_grouper = LanguageServerSymbolDictGrouper(["kind"], ["kind"]) + def apply(self, relative_path: str, depth: int = 0, max_answer_chars: int = -1) -> str: """ Use this tool to get a high-level understanding of the code symbols in a file. @@ -65,11 +69,12 @@ 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._transform_symbols_to_compact_format(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) - def get_symbol_overview(self, relative_path: str, depth: int = 0) -> list[dict]: + def get_symbol_overview(self, relative_path: str, depth: int = 0) -> list[LanguageServerSymbol.OutputDict]: """ :param relative_path: relative path to a source file :param depth: the depth up to which descendants shall be retrieved