diff --git a/src/serena/symbol.py b/src/serena/symbol.py index 8eb00595..6a17813a 100644 --- a/src/serena/symbol.py +++ b/src/serena/symbol.py @@ -5,6 +5,7 @@ import os from abc import ABC, abstractmethod from collections import OrderedDict from collections.abc import Callable, Iterable, Iterator, Sequence +from contextlib import contextmanager from dataclasses import asdict, dataclass from time import perf_counter from typing import Any, Generic, Literal, NotRequired, Self, TypedDict, TypeVar @@ -1261,6 +1262,7 @@ class SymbolDictGrouper(Generic[TSymbolDict], ABC): self._group_keys = group_keys self._group_children_keys = group_children_keys self._collapse_singleton = collapse_singleton + self._is_enabled = True def _group_by(self, l: list[dict], keys: list[str], children_keys: list[str], is_children: bool) -> dict[str, Any] | list[Any]: """ @@ -1321,8 +1323,24 @@ class SymbolDictGrouper(Generic[TSymbolDict], ABC): """ # avoid side effects by working on a deep-copy symbols_copy = copy.deepcopy(symbols) + + if not self._is_enabled: + return symbols_copy + return self._group_by(symbols_copy, self._group_keys, self._group_children_keys, is_children=False) + @contextmanager + def disabled_context(self) -> Iterator[None]: + """ + A context manager that temporarily disables grouping. + """ + original_state = self._is_enabled + self._is_enabled = False + try: + yield + finally: + self._is_enabled = original_state + class LanguageServerSymbolDictGrouper(SymbolDictGrouper[LanguageServerSymbol.OutputDict]): def __init__(