Add support for textDocument/documentSymbol request in multilspy Lang…

…uageServer

src/monitors4codegen/multilspy/language_server.py

@@ -26,7 +26,7 @@
 from .multilspy_exceptions import MultilspyException
 from .multilspy_utils import PathUtils, FileUtils, TextUtils
 from pathlib import PurePath
-from typing import AsyncIterator, Iterator, List, Dict, Union
+from typing import AsyncIterator, Iterator, List, Dict, Union, Tuple
 from .type_helpers import ensure_all_methods_implemented
 
 
@@ -550,6 +550,51 @@ async def request_completions(
                 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]]:
+        """
+        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
+        """
+        with self.open_file(relative_file_path):
+            response = await self.server.send.document_symbol(
+                {
+                    "textDocument": {
+                        "uri": pathlib.Path(os.path.join(self.repository_root_path, relative_file_path)).as_uri()
+                    }
+                }
+            )
+
+        ret: List[multilspy_types.UnifiedSymbolInformation] = []
+        l_tree = None
+        assert isinstance(response, list)
+        for item in response:
+            assert isinstance(item, dict)
+            assert LSPConstants.NAME in item
+            assert LSPConstants.KIND in item
+
+            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]:
+                    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))
+                    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))
+            else:
+                ret.append(multilspy_types.UnifiedSymbolInformation(**item))
+
+        return ret, l_tree
+
 @ensure_all_methods_implemented(LanguageServer)
 class SyncLanguageServer:
     """
@@ -689,3 +734,17 @@ def request_completions(
             self.loop,
         ).result()
         return result
+
+    def request_document_symbols(self, relative_file_path: str) -> Tuple[List[multilspy_types.UnifiedSymbolInformation], Union[List[multilspy_types.TreeRepr], None]]:
+        """
+        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
+        """
+        result = asyncio.run_coroutine_threadsafe(
+            self.language_server.request_document_symbols(relative_file_path), self.loop
+        ).result()
+        return result

src/monitors4codegen/multilspy/lsp_protocol_handler/lsp_constants.py

@@ -48,3 +48,12 @@ class LSPConstants:
 
     # key used to represent the changes made to a document
     CONTENT_CHANGES = "contentChanges"
+
+    # key used to represent name of symbols
+    NAME = "name"
+
+    # key used to represent the kind of symbols
+    KIND = "kind"
+
+    # key used to represent children in document symbols
+    CHILDREN = "children"

src/monitors4codegen/multilspy/multilspy_types.py

@@ -2,8 +2,10 @@
 Defines wrapper objects around the types returned by LSP to ensure decoupling between LSP versions and multilspy
 """
 
+from __future__ import annotations
+
 from enum import IntEnum
-from typing import TypedDict
+from typing_extensions import NotRequired, TypedDict, List, Dict
 
 URI = str
 DocumentUri = str
@@ -123,4 +125,87 @@ class CompletionItem(TypedDict):
 
     kind: CompletionItemKind
     """ The kind of this completion item. Based of the kind
-    an icon is chosen by the editor. """
+    an icon is chosen by the editor. """
+
+class SymbolKind(IntEnum):
+    """A symbol kind."""
+
+    File = 1
+    Module = 2
+    Namespace = 3
+    Package = 4
+    Class = 5
+    Method = 6
+    Property = 7
+    Field = 8
+    Constructor = 9
+    Enum = 10
+    Interface = 11
+    Function = 12
+    Variable = 13
+    Constant = 14
+    String = 15
+    Number = 16
+    Boolean = 17
+    Array = 18
+    Object = 19
+    Key = 20
+    Null = 21
+    EnumMember = 22
+    Struct = 23
+    Event = 24
+    Operator = 25
+    TypeParameter = 26
+
+class SymbolTag(IntEnum):
+    """Symbol tags are extra annotations that tweak the rendering of a symbol.
+
+    @since 3.16"""
+
+    Deprecated = 1
+    """ Render a symbol as obsolete, usually using a strike-out. """
+
+class UnifiedSymbolInformation(TypedDict):
+    """Represents information about programming constructs like variables, classes,
+    interfaces etc."""
+
+    deprecated: NotRequired[bool]
+    """ Indicates if this symbol is deprecated.
+
+    @deprecated Use tags instead """
+    location: NotRequired[Location]
+    """ The location of this symbol. The location's range is used by a tool
+    to reveal the location in the editor. If the symbol is selected in the
+    tool the range's start information is used to position the cursor. So
+    the range usually spans more than the actual symbol's name and does
+    normally include things like visibility modifiers.
+
+    The range doesn't have to denote a node range in the sense of an abstract
+    syntax tree. It can therefore not be used to re-construct a hierarchy of
+    the symbols. """
+    name: str
+    """ The name of this symbol. """
+    kind: SymbolKind
+    """ The kind of this symbol. """
+    tags: NotRequired[List[SymbolTag]]
+    """ Tags for this symbol.
+
+    @since 3.16.0 """
+    containerName: NotRequired[str]
+    """ The name of the symbol containing this symbol. This information is for
+    user interface purposes (e.g. to render a qualifier in the user interface
+    if necessary). It can't be used to re-infer a hierarchy for the document
+    symbols. """
+
+    detail: NotRequired[str]
+    """ More detail for this symbol, e.g the signature of a function. """
+
+    range: NotRequired[Range]
+    """ The range enclosing this symbol not including leading/trailing whitespace but everything else
+    like comments. This information is typically used to determine if the clients cursor is
+    inside the symbol to reveal in the symbol in the UI. """
+    selectionRange: NotRequired[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`. """
+
+TreeRepr = Dict[int, List['TreeRepr']]

tests/multilspy/test_multilspy_java.py

@@ -206,3 +206,142 @@ async def test_multilspy_java_clickhouse_highlevel_sinker_modified():
                 completions = await lsp.request_completions(completions_filepath, 136, 23)
                 completions = [completion["completionText"] for completion in completions if completion["kind"] == CompletionItemKind.Constructor]
                 assert completions == ['ClickHouseSinkBuffer']
+
+@pytest.mark.asyncio
+async def test_multilspy_java_example_repo_document_symbols() -> None:
+    """
+    Test the working of multilspy with Java repository - clickhouse-highlevel-sinker
+    """
+    code_language = Language.JAVA
+    params = {
+        "code_language": code_language,
+        "repo_url": "https://github.com/LakshyAAAgrawal/ExampleRepo/",
+        "repo_commit": "f3762fd55a457ff9c6b0bf3b266de2b203a766ab",
+    }
+    with create_test_context(params) as context:
+        lsp = LanguageServer.create(context.config, context.logger, context.source_directory)
+
+        # All the communication with the language server must be performed inside the context manager
+        # The server process is started when the context manager is entered and is terminated when the context manager is exited.
+        async with lsp.start_server():
+            filepath = str(PurePath("Person.java"))
+            result = await lsp.request_document_symbols(filepath)
+
+            assert result == (
+                [
+                    {
+                        "name": "Person",
+                        "kind": 5,
+                        "range": {
+                            "start": {"line": 0, "character": 0},
+                            "end": {"line": 14, "character": 1},
+                        },
+                        "selectionRange": {
+                            "start": {"line": 1, "character": 22},
+                            "end": {"line": 1, "character": 28},
+                        },
+                        "detail": "",
+                    },
+                    {
+                        "name": "name",
+                        "kind": 8,
+                        "range": {
+                            "start": {"line": 2, "character": 4},
+                            "end": {"line": 3, "character": 24},
+                        },
+                        "selectionRange": {
+                            "start": {"line": 3, "character": 19},
+                            "end": {"line": 3, "character": 23},
+                        },
+                        "detail": "",
+                    },
+                    {
+                        "name": "Person(String)",
+                        "kind": 9,
+                        "range": {
+                            "start": {"line": 5, "character": 4},
+                            "end": {"line": 8, "character": 5},
+                        },
+                        "selectionRange": {
+                            "start": {"line": 6, "character": 11},
+                            "end": {"line": 6, "character": 17},
+                        },
+                        "detail": "",
+                    },
+                    {
+                        "name": "getName()",
+                        "kind": 6,
+                        "range": {
+                            "start": {"line": 10, "character": 4},
+                            "end": {"line": 13, "character": 5},
+                        },
+                        "selectionRange": {
+                            "start": {"line": 11, "character": 18},
+                            "end": {"line": 11, "character": 25},
+                        },
+                        "detail": " : String",
+                    },
+                ],
+                None,
+            )
+
+            filepath = str(PurePath("Student.java"))
+            result = await lsp.request_document_symbols(filepath)
+
+            assert result == (
+                [
+                    {
+                        "name": "Student",
+                        "kind": 5,
+                        "range": {
+                            "start": {"line": 0, "character": 0},
+                            "end": {"line": 16, "character": 1},
+                        },
+                        "selectionRange": {
+                            "start": {"line": 1, "character": 13},
+                            "end": {"line": 1, "character": 20},
+                        },
+                        "detail": "",
+                    },
+                    {
+                        "name": "id",
+                        "kind": 8,
+                        "range": {
+                            "start": {"line": 2, "character": 4},
+                            "end": {"line": 3, "character": 19},
+                        },
+                        "selectionRange": {
+                            "start": {"line": 3, "character": 16},
+                            "end": {"line": 3, "character": 18},
+                        },
+                        "detail": "",
+                    },
+                    {
+                        "name": "Student(String, int)",
+                        "kind": 9,
+                        "range": {
+                            "start": {"line": 5, "character": 4},
+                            "end": {"line": 10, "character": 5},
+                        },
+                        "selectionRange": {
+                            "start": {"line": 6, "character": 11},
+                            "end": {"line": 6, "character": 18},
+                        },
+                        "detail": "",
+                    },
+                    {
+                        "name": "getId()",
+                        "kind": 6,
+                        "range": {
+                            "start": {"line": 12, "character": 4},
+                            "end": {"line": 15, "character": 5},
+                        },
+                        "selectionRange": {
+                            "start": {"line": 13, "character": 15},
+                            "end": {"line": 13, "character": 20},
+                        },
+                        "detail": " : int",
+                    },
+                ],
+                None,
+            )