From 69ab44eb74481b50dee9ac68778cb50b766f9c86 Mon Sep 17 00:00:00 2001 From: Michael Panchenko Date: Tue, 25 Mar 2025 00:55:03 +0100 Subject: [PATCH] Document symbols can now include body and retrieve the document-structure of a file --- src/multilspy/language_server.py | 76 +++++++++++++------ .../lsp_protocol_handler/lsp_types.py | 2 - src/multilspy/multilspy_types.py | 6 ++ test/multilspy/test_symbol_retrieval.py | 34 +++++++++ 4 files changed, 92 insertions(+), 26 deletions(-) diff --git a/src/multilspy/language_server.py b/src/multilspy/language_server.py index 22a5e1af..746ab937 100644 --- a/src/multilspy/language_server.py +++ b/src/multilspy/language_server.py @@ -31,7 +31,7 @@ from .multilspy_config import MultilspyConfig, Language from .multilspy_exceptions import MultilspyException from .multilspy_utils import PathUtils, FileUtils, TextUtils from pathlib import Path, PurePath -from typing import AsyncIterator, Iterator, List, Dict, Optional, Union, Tuple +from typing import AsyncIterator, Iterator, List, Dict, Optional, TypedDict, Union, Tuple, cast from .type_helpers import ensure_all_methods_implemented @@ -44,6 +44,8 @@ from .type_helpers import ensure_all_methods_implemented # it makes sense to have more content-related utils directly in it. +GenericDocumentSymbol = Union[LSPTypes.DocumentSymbol, LSPTypes.SymbolInformation, multilspy_types.UnifiedSymbolInformation] + @dataclasses.dataclass class LSPFileBuffer: """ @@ -187,7 +189,7 @@ class LanguageServer: self.open_file_buffers: Dict[str, LSPFileBuffer] = {} # --------------------------------- MODIFICATIONS BY MISCHA --------------------------------- - self._document_symbols_cache: dict[str, Tuple[str, Tuple[List[multilspy_types.UnifiedSymbolInformation], Optional[List[multilspy_types.TreeRepr]]]]] = {} + self._document_symbols_cache: dict[str, Tuple[str, Tuple[List[multilspy_types.UnifiedSymbolInformation], List[multilspy_types.UnifiedSymbolInformation]]]] = {} """Maps file paths to a tuple of (file_content_hash, result_of_request_document_symbols)""" self.load_cache() self._cache_has_changed = bool @@ -651,18 +653,21 @@ class LanguageServer: for json_repr in set([json.dumps(item, sort_keys=True) for item in completions_list]) ] - async def request_document_symbols(self, relative_file_path: str) -> Tuple[List[multilspy_types.UnifiedSymbolInformation], Union[List[multilspy_types.TreeRepr], None]]: + async def request_document_symbols(self, relative_file_path: str, include_body: bool = False) -> Tuple[List[multilspy_types.UnifiedSymbolInformation], List[multilspy_types.UnifiedSymbolInformation]]: """ Raise a [textDocument/documentSymbol](https://microsoft.github.io/language-server-protocol/specifications/lsp/3.17/specification/#textDocument_documentSymbol) request to the Language Server to find symbols in the given file. Wait for the response and return the result. :param relative_file_path: The relative path of the file that has the symbols - - :return Tuple[List[multilspy_types.UnifiedSymbolInformation], Union[List[multilspy_types.TreeRepr], None]]: A list of symbols in the file, and the tree representation of the symbols - """ + :param include_body: whether to include the body of the symbols in the result. + :return: A list of symbols in the file, and a list of root symbols that represent the tree structure of the symbols. Each symbol in hierarchy starting from the roots has a children attribute. + """ self.logger.log(f"Requesting document symbols for {relative_file_path} for the first time", logging.DEBUG) + # TODO: it's kinda dumb to not use the cache if include_body is False after include_body was True once + # Should be fixed in the future, it's a small performance optimization + cache_key = f"{relative_file_path}-{include_body}" with self.open_file(relative_file_path) as file_data: - file_hash_and_result = self._document_symbols_cache.get(relative_file_path) + file_hash_and_result = self._document_symbols_cache.get(cache_key) if file_hash_and_result is not None: file_hash, result = file_hash_and_result if file_hash == file_data.content_hash: @@ -679,11 +684,31 @@ class LanguageServer: } } ) + + def turn_item_into_symbol_with_children(item: GenericDocumentSymbol): + item = cast(multilspy_types.UnifiedSymbolInformation, item) + if "location" not in item: + absolute_path = os.path.join(self.repository_root_path, relative_file_path) + uri = pathlib.Path(absolute_path).as_uri() + assert "range" in item + tree_location = multilspy_types.Location( + uri=uri, + range=item['range'], + absolutePath=absolute_path, + relativePath=relative_file_path, + ) + item['location'] = tree_location + if include_body: + item['body'] = self.retrieve_symbol_body(item) + item[LSPConstants.CHILDREN] = item.get(LSPConstants.CHILDREN, []) - ret: List[multilspy_types.UnifiedSymbolInformation] = [] - l_tree = None + symbols_without_children: List[multilspy_types.UnifiedSymbolInformation] = [] assert isinstance(response, list) + root_nodes: List[multilspy_types.UnifiedSymbolInformation] = [] for item in response: + turn_item_into_symbol_with_children(item) + item = cast(multilspy_types.UnifiedSymbolInformation, item) + root_nodes.append(item) assert isinstance(item, dict) assert LSPConstants.NAME in item assert LSPConstants.KIND in item @@ -691,26 +716,29 @@ class LanguageServer: if LSPConstants.CHILDREN in item: # TODO: l_tree should be a list of TreeRepr. Define the following function to return TreeRepr as well - def visit_tree_nodes_and_build_tree_repr(tree: LSPTypes.DocumentSymbol) -> List[multilspy_types.UnifiedSymbolInformation]: + def visit_tree_nodes_and_build_tree_repr(node: GenericDocumentSymbol) -> List[multilspy_types.UnifiedSymbolInformation]: + node = cast(multilspy_types.UnifiedSymbolInformation, node) l: List[multilspy_types.UnifiedSymbolInformation] = [] - children = tree['children'] if 'children' in tree else [] - if 'children' in tree: - del tree['children'] - l.append(multilspy_types.UnifiedSymbolInformation(**tree)) + turn_item_into_symbol_with_children(node) + assert LSPConstants.CHILDREN in node + children = node[LSPConstants.CHILDREN] + node_without_children = node.copy() + del node_without_children[LSPConstants.CHILDREN] + l.append(node_without_children) for child in children: l.extend(visit_tree_nodes_and_build_tree_repr(child)) return l - ret.extend(visit_tree_nodes_and_build_tree_repr(item)) + symbols_without_children.extend(visit_tree_nodes_and_build_tree_repr(item)) else: - ret.append(multilspy_types.UnifiedSymbolInformation(**item)) + symbols_without_children.append(multilspy_types.UnifiedSymbolInformation(**item)) - result = ret, l_tree + result = symbols_without_children, root_nodes self.logger.log(f"Caching document symbols for {relative_file_path}", logging.DEBUG) - self._document_symbols_cache[relative_file_path] = (file_data.content_hash, result) + self._document_symbols_cache[cache_key] = (file_data.content_hash, result) self._cache_has_changed = True return result - + async def request_hover(self, relative_file_path: str, line: int, column: int) -> Union[multilspy_types.Hover, None]: """ Raise a [textDocument/hover](https://microsoft.github.io/language-server-protocol/specifications/lsp/3.17/specification/#textDocument_hover) request to the Language Server @@ -744,7 +772,7 @@ class LanguageServer: # ----------------------------- FROM HERE ON MODIFICATIONS BY MISCHA -------------------- - def retrieve_symbol_body(self, symbol: multilspy_types.UnifiedSymbolInformation) -> str: + def retrieve_symbol_body(self, symbol: multilspy_types.UnifiedSymbolInformation | LSPTypes.DocumentSymbol | LSPTypes.SymbolInformation) -> str: """ Load the body of the given symbol. If the body is already contained in the symbol, just return it. """ @@ -1241,17 +1269,17 @@ class SyncLanguageServer: ).result() return result - def request_document_symbols(self, relative_file_path: str) -> Tuple[List[multilspy_types.UnifiedSymbolInformation], Union[List[multilspy_types.TreeRepr], None]]: + def request_document_symbols(self, relative_file_path: str, include_body: bool = False) -> Tuple[List[multilspy_types.UnifiedSymbolInformation], List[multilspy_types.UnifiedSymbolInformation]]: """ Raise a [textDocument/documentSymbol](https://microsoft.github.io/language-server-protocol/specifications/lsp/3.17/specification/#textDocument_documentSymbol) request to the Language Server to find symbols in the given file. Wait for the response and return the result. :param relative_file_path: The relative path of the file that has the symbols - - :return Tuple[List[multilspy_types.UnifiedSymbolInformation], Union[List[multilspy_types.TreeRepr], None]]: A list of symbols in the file, and the tree representation of the symbols + :param include_body: whether to include the body of the symbols in the result. + :return: A list of symbols in the file, and a list of root symbols that represent the tree structure of the symbols. Each symbol in hierarchy starting from the roots has a children attribute. """ result = asyncio.run_coroutine_threadsafe( - self.language_server.request_document_symbols(relative_file_path), self.loop + self.language_server.request_document_symbols(relative_file_path, include_body), self.loop ).result() return result diff --git a/src/multilspy/lsp_protocol_handler/lsp_types.py b/src/multilspy/lsp_protocol_handler/lsp_types.py index 98f70b33..a6476164 100644 --- a/src/multilspy/lsp_protocol_handler/lsp_types.py +++ b/src/multilspy/lsp_protocol_handler/lsp_types.py @@ -2416,8 +2416,6 @@ class DocumentSymbol(TypedDict): selectionRange: "Range" """ The range that should be selected and revealed when this symbol is being picked, e.g the name of a function. Must be contained by the `range`. """ - children: NotRequired[List["DocumentSymbol"]] - """ Children of this symbol, e.g. properties of a class. """ class DocumentSymbolRegistrationOptions(TypedDict): diff --git a/src/multilspy/multilspy_types.py b/src/multilspy/multilspy_types.py index eae28088..55ec0d7e 100644 --- a/src/multilspy/multilspy_types.py +++ b/src/multilspy/multilspy_types.py @@ -214,6 +214,12 @@ class UnifiedSymbolInformation(TypedDict): body: NotRequired[str] """ The body of the symbol. """ + + children: NotRequired[List[UnifiedSymbolInformation]] + """ The children of the symbol. + Added to be compatible with `lsp_types.DocumentSymbol`, + since it is sometimes useful to have the children of the symbol as a user-facing feature.""" + TreeRepr = Dict[int, List['TreeRepr']] diff --git a/test/multilspy/test_symbol_retrieval.py b/test/multilspy/test_symbol_retrieval.py index 1ef68fe9..2ca9256a 100644 --- a/test/multilspy/test_symbol_retrieval.py +++ b/test/multilspy/test_symbol_retrieval.py @@ -318,3 +318,37 @@ class TestLanguageServerSymbols: import warnings warnings.warn("Could not verify container hierarchy - implementation detail") + + def test_symbol_tree_structure(self, language_server: SyncLanguageServer, repo_path: Path): + """Test the symbol tree structure.""" + file_path = str(repo_path / "test_repo" / "services.py") + symbols, root_nodes = language_server.request_document_symbols(file_path) + assert len(symbols) > 0 + assert {root["name"] for root in root_nodes} == { + "UserService", + "ItemService", + "create_service_container", + "user_var_str", + "user_service", + } + user_service_root = next(root for root in root_nodes if root["name"] == "UserService") + assert user_service_root + assert "children" in user_service_root + assert {child["name"] for child in user_service_root["children"] if child["kind"] != SymbolKind.Variable} == { + "__init__", + "create_user", + "get_user", + "list_users", + "delete_user", + } + + # Now recursively flatten the tree into a set of names and assert that it coincides with the set of symbol names + all_names_in_tree = set() + + def flatten_tree(nodes): + for node in nodes: + all_names_in_tree.add(node["name"]) + flatten_tree(node.get("children", [])) + + flatten_tree(root_nodes) + assert all_names_in_tree == {symbol["name"] for symbol in symbols}