Type annotation improvments

pygls/capabilities.py

@@ -15,11 +15,12 @@
 # limitations under the License.                                           #
 ############################################################################
 import logging
-from functools import reduce
 from typing import Any, Dict, List, Optional, Set, TypeVar, Union
 
 from lsprotocol import types
 
+from pygls.lsp._capabilities import get_capability
+
 logger = logging.getLogger(__name__)
 T = TypeVar("T")
 
@@ -32,22 +33,6 @@
 )
 
 
-def get_capability(
-    client_capabilities: types.ClientCapabilities, field: str, default: Any = None
-) -> Any:
-    """Check if ClientCapabilities has some nested value without raising
-    AttributeError.
-    e.g. get_capability('text_document.synchronization.will_save')
-    """
-    try:
-        value = reduce(getattr, field.split("."), client_capabilities)
-    except AttributeError:
-        return default
-
-    # If we reach the desired leaf value but it's None, return the default.
-    return default if value is None else value
-
-
 class ServerCapabilitiesBuilder:
     """Create `ServerCapabilities` instance depending on builtin and user registered
     features.

pygls/exceptions.py

@@ -16,6 +16,8 @@
 # See the License for the specific language governing permissions and      #
 # limitations under the License.                                           #
 ############################################################################
+from __future__ import annotations
+
 import traceback
 from typing import Any
 from typing import Set
@@ -26,16 +28,23 @@
 class JsonRpcException(Exception):
     """A class used as a base class for json rpc exceptions."""
 
+    CODE = -32603
     MESSAGE = ""
 
-    def __init__(self, message=None, code=None, data=None):
-        message = message or getattr(self.__class__, "MESSAGE")
+    def __init__(
+        self,
+        message: str | None = None,
+        code: int | None = None,
+        data: Any | None = None,
+    ):
+        message = message or self.MESSAGE
+
         super().__init__(message)
-        self.message = message
-        self.code = code or getattr(self.__class__, "CODE", -32603)
+        self.message: str = message
+        self.code: int = code or self.CODE
         self.data = data
 
-    def __eq__(self, other):
+    def __eq__(self, other: Any):
         return (
             isinstance(other, self.__class__)
             and self.code == other.code
@@ -46,7 +55,7 @@ def __hash__(self):
         return hash((self.code, self.message))
 
     @staticmethod
-    def from_error(error):
+    def from_error(error: ResponseError):
         for exc_class in _EXCEPTIONS:
             if exc_class.supports_code(error.code):
                 return exc_class(
@@ -65,7 +74,7 @@ def of(cls, exc: Any):
         )
 
     @classmethod
-    def supports_code(cls, code):
+    def supports_code(cls, code: int):
         # Defaults to UnknownErrorCode
         return getattr(cls, "CODE", -32001) == code