1 files changed, 103 insertions, 20 deletions
diff --git a/src/debputy/lsp/lsp_dispatch.py b/src/debputy/lsp/lsp_dispatch.py
index 41e9111..b7b744c 100644
--- a/src/debputy/lsp/lsp_dispatch.py
+++ b/src/debputy/lsp/lsp_dispatch.py
@@ -1,5 +1,15 @@
 import asyncio
-from typing import Dict, Sequence, Union, Optional
+from typing import (
+    Dict,
+    Sequence,
+    Union,
+    Optional,
+    Any,
+    TypeVar,
+    Callable,
+    Mapping,
+    List,
+)
 
 from lsprotocol.types import (
     DidOpenTextDocumentParams,
@@ -11,6 +21,19 @@ from lsprotocol.types import (
     CompletionItem,
     CompletionParams,
     TEXT_DOCUMENT_HOVER,
+    TEXT_DOCUMENT_FOLDING_RANGE,
+    FoldingRange,
+    FoldingRangeParams,
+    TEXT_DOCUMENT_SEMANTIC_TOKENS_FULL,
+    SemanticTokensParams,
+    SemanticTokens,
+    Hover,
+    TEXT_DOCUMENT_CODE_ACTION,
+    Command,
+    CodeAction,
+    TextDocumentCodeActionRequest,
+    CodeActionParams,
+    SemanticTokensRegistrationOptions,
 )
 
 from debputy import __version__
@@ -18,6 +41,9 @@ from debputy.lsp.lsp_features import (
     DIAGNOSTIC_HANDLERS,
     COMPLETER_HANDLERS,
     HOVER_HANDLERS,
+    SEMANTIC_TOKENS_FULL_HANDLERS,
+    CODE_ACTION_HANDLERS,
+    SEMANTIC_TOKENS_LEGEND,
 )
 from debputy.util import _info
 
@@ -37,6 +63,10 @@ except ImportError:
     DEBPUTY_LANGUAGE_SERVER = Mock()
 
 
+P = TypeVar("P")
+R = TypeVar("R")
+
+
 def is_doc_at_version(uri: str, version: int) -> bool:
     dv = _DOCUMENT_VERSION_TABLE.get(uri)
     return dv == version
@@ -88,22 +118,12 @@ def _completions(
     ls: "LanguageServer",
     params: CompletionParams,
 ) -> Optional[Union[CompletionList, Sequence[CompletionItem]]]:
-    doc_uri = params.text_document.uri
-    doc = ls.workspace.get_text_document(doc_uri)
-
-    handler = COMPLETER_HANDLERS.get(doc.language_id)
-    if handler is None:
-        _info(
-            f"Complete request for document: {doc.path} ({doc.language_id}) - no handler"
-        )
-        return
-    _info(
-        f"Complete request for document: {doc.path} ({doc.language_id}) - delegating to handler"
-    )
-
-    return handler(
+    return _dispatch_standard_handler(
         ls,
+        params.text_document.uri,
         params,
+        COMPLETER_HANDLERS,
+        "Complete request",
     )
 
 
@@ -111,18 +131,81 @@ def _completions(
 def _hover(
     ls: "LanguageServer",
     params: CompletionParams,
-) -> Optional[Union[CompletionList, Sequence[CompletionItem]]]:
-    doc_uri = params.text_document.uri
+) -> Optional[Hover]:
+    return _dispatch_standard_handler(
+        ls,
+        params.text_document.uri,
+        params,
+        HOVER_HANDLERS,
+        "Hover doc request",
+    )
+
+
+@DEBPUTY_LANGUAGE_SERVER.feature(TEXT_DOCUMENT_CODE_ACTION)
+def _code_actions(
+    ls: "LanguageServer",
+    params: CodeActionParams,
+) -> Optional[List[Union[Command, CodeAction]]]:
+    return _dispatch_standard_handler(
+        ls,
+        params.text_document.uri,
+        params,
+        CODE_ACTION_HANDLERS,
+        "Code action request",
+    )
+
+
+@DEBPUTY_LANGUAGE_SERVER.feature(TEXT_DOCUMENT_FOLDING_RANGE)
+def _folding_ranges(
+    ls: "LanguageServer",
+    params: FoldingRangeParams,
+) -> Optional[Sequence[FoldingRange]]:
+    return _dispatch_standard_handler(
+        ls,
+        params.text_document.uri,
+        params,
+        HOVER_HANDLERS,
+        "Folding range request",
+    )
+
+
+@DEBPUTY_LANGUAGE_SERVER.feature(
+    TEXT_DOCUMENT_SEMANTIC_TOKENS_FULL,
+    SemanticTokensRegistrationOptions(
+        SEMANTIC_TOKENS_LEGEND,
+        full=True,
+    ),
+)
+def _semantic_tokens_full(
+    ls: "LanguageServer",
+    params: SemanticTokensParams,
+) -> Optional[SemanticTokens]:
+    return _dispatch_standard_handler(
+        ls,
+        params.text_document.uri,
+        params,
+        SEMANTIC_TOKENS_FULL_HANDLERS,
+        "Semantic tokens request",
+    )
+
+
+def _dispatch_standard_handler(
+    ls: "LanguageServer",
+    doc_uri: str,
+    params: P,
+    handler_table: Mapping[str, Callable[["LanguageServer", P], R]],
+    request_type: str,
+) -> R:
     doc = ls.workspace.get_text_document(doc_uri)
 
-    handler = HOVER_HANDLERS.get(doc.language_id)
+    handler = handler_table.get(doc.language_id)
     if handler is None:
         _info(
-            f"Hover request for document: {doc.path} ({doc.language_id}) - no handler"
+            f"{request_type} for document: {doc.path} ({doc.language_id}) - no handler"
         )
         return
     _info(
-        f"Hover request for document: {doc.path} ({doc.language_id}) - delegating to handler"
+        f"{request_type} for document: {doc.path} ({doc.language_id}) - delegating to handler"
     )
 
     return handler(