From d719b17447f59366d7264fddfc4a6258ebf880ea Mon Sep 17 00:00:00 2001
From: Thomas Buckley-Houston <tom@tombh.co.uk>
Date: Sun, 10 Sep 2023 16:55:58 -0500
Subject: [PATCH] feat: accept `PositionEncoding` client capability

This commit is only for supporting sending preferred position encoding
from the client during startup. Even though this commit introduces the
server's ability to choose from the client's list of supported position
encodings, it doesn't actually support any other encoding from the
existing UTF-16.
---
 pygls/capabilities.py         | 31 +++++++++++++++++++++++++++++++
 pygls/protocol.py             |  5 +++++
 tests/test_feature_manager.py | 26 ++++++++++++++++++++++++++
 3 files changed, 62 insertions(+)

diff --git a/pygls/capabilities.py b/pygls/capabilities.py
index c36832b0..f34f806b 100644
--- a/pygls/capabilities.py
+++ b/pygls/capabilities.py
@@ -16,6 +16,7 @@
 ############################################################################
 from functools import reduce
 from typing import Any, Dict, List, Set, Union
+import logging
 
 from lsprotocol.types import (
     INLAY_HINT_RESOLVE,
@@ -64,6 +65,7 @@
     WORKSPACE_WILL_DELETE_FILES,
     WORKSPACE_WILL_RENAME_FILES,
     InlayHintOptions,
+    PositionEncodingKind,
 )
 from lsprotocol.types import (
     ClientCapabilities,
@@ -86,6 +88,8 @@
     WorkspaceFoldersServerCapabilities,
 )
 
+logger = logging.getLogger(__name__)
+
 
 def get_capability(
     client_capabilities: ClientCapabilities, field: str, default: Any = None
@@ -429,6 +433,32 @@ def _with_inline_value_provider(self):
             self.server_cap.inline_value_provider = value
         return self
 
+    def _with_position_encodings(self):
+        self.server_cap.position_encoding = PositionEncodingKind.Utf16
+
+        general = self.client_capabilities.general
+        if general is None:
+            return self
+
+        encodings = general.position_encodings
+        if encodings is None:
+            return self
+
+        if PositionEncodingKind.Utf16 in encodings:
+            return self
+
+        if PositionEncodingKind.Utf32 in encodings:
+            self.server_cap.position_encoding = PositionEncodingKind.Utf32
+            return self
+
+        if PositionEncodingKind.Utf8 in encodings:
+            self.server_cap.position_encoding = PositionEncodingKind.Utf8
+            return self
+
+        logger.warning(f"Unknown `PositionEncoding`s: {encodings}")
+
+        return self
+
     def _build(self):
         return self.server_cap
 
@@ -467,5 +497,6 @@ def build(self):
             ._with_workspace_capabilities()
             ._with_diagnostic_provider()
             ._with_inline_value_provider()
+            ._with_position_encodings()
             ._build()
         )
diff --git a/pygls/protocol.py b/pygls/protocol.py
index 65cd62d7..64c25ec1 100644
--- a/pygls/protocol.py
+++ b/pygls/protocol.py
@@ -73,6 +73,7 @@
     WORKSPACE_DID_CHANGE_WORKSPACE_FOLDERS,
     WORKSPACE_EXECUTE_COMMAND,
     WORKSPACE_SEMANTIC_TOKENS_REFRESH,
+    PositionEncodingKind,
 )
 from lsprotocol.types import (
     ApplyWorkspaceEditParams,
@@ -214,6 +215,10 @@ def default_converter():
 
     converter.register_structure_hook(JsonRPCNotification, _params_field_structure_hook)
 
+    converter.register_structure_hook(
+        Union[PositionEncodingKind, str, None], lambda obj, _: obj
+    )
+
     return converter
 
 
diff --git a/tests/test_feature_manager.py b/tests/test_feature_manager.py
index d9655b85..4e5c852b 100644
--- a/tests/test_feature_manager.py
+++ b/tests/test_feature_manager.py
@@ -215,12 +215,38 @@ def server_capabilities(**kwargs):
             file_operations=lsp.FileOperationOptions(),
         )
 
+    if "position_encoding" not in kwargs:
+        kwargs["position_encoding"] = lsp.PositionEncodingKind.Utf16
+
     return lsp.ServerCapabilities(**kwargs)
 
 
 @pytest.mark.parametrize(
     "method, options, capabilities, expected",
     [
+        (
+            lsp.INITIALIZE,
+            None,
+            lsp.ClientCapabilities(
+                general=lsp.GeneralClientCapabilities(
+                    position_encodings=[lsp.PositionEncodingKind.Utf8]
+                )
+            ),
+            server_capabilities(position_encoding=lsp.PositionEncodingKind.Utf8),
+        ),
+        (
+            lsp.INITIALIZE,
+            None,
+            lsp.ClientCapabilities(
+                general=lsp.GeneralClientCapabilities(
+                    position_encodings=[
+                        lsp.PositionEncodingKind.Utf8,
+                        lsp.PositionEncodingKind.Utf32,
+                    ]
+                )
+            ),
+            server_capabilities(position_encoding=lsp.PositionEncodingKind.Utf32),
+        ),
         (
             lsp.TEXT_DOCUMENT_DID_SAVE,
             lsp.SaveOptions(include_text=True),