update

Author: observerw

src/lsap/capability/abc.py

@@ -2,6 +2,7 @@
 from typing import Protocol
 
 from attrs import define
+from lsp_client import Client
 from lsp_client.protocol import CapabilityClientProtocol
 from pydantic import BaseModel
 
@@ -10,8 +11,8 @@ class ClientProtocol(CapabilityClientProtocol, Protocol): ...
 
 
 @define
-class Capability[C: ClientProtocol, Req: BaseModel, Resp: BaseModel](ABC):
-    client: C
+class Capability[Req: BaseModel, Resp: BaseModel](ABC):
+    client: Client
 
     @abstractmethod
     async def __call__(self, req: Req) -> Resp | None: ...

src/lsap/capability/definition.py

@@ -1,42 +1,28 @@
 from __future__ import annotations
 
 from functools import cached_property
-from typing import Protocol, Sequence, override, runtime_checkable
+from typing import Sequence, override
 
 import asyncer
 from attrs import define
 from lsp_client.capability.request import (
     WithRequestDeclaration,
     WithRequestDefinition,
-    WithRequestDocumentSymbol,
-    WithRequestHover,
     WithRequestTypeDefinition,
 )
 from lsprotocol.types import Location
 
-from lsap.exception import UnsupportedCapabilityError
 from lsap.schema.definition import DefinitionRequest, DefinitionResponse
 from lsap.schema.models import SymbolCodeInfo
+from lsap.utils.capability import ensure_capability
 
-from .abc import Capability, ClientProtocol
+from .abc import Capability
 from .locate import LocateCapability
 from .symbol import SymbolCapability
 
 
-@runtime_checkable
-class DefinitionClient(
-    WithRequestDefinition,
-    WithRequestDocumentSymbol,
-    WithRequestHover,
-    ClientProtocol,
-    Protocol,
-): ...
-
-
 @define
-class DefinitionCapability(
-    Capability[DefinitionClient, DefinitionRequest, DefinitionResponse]
-):
+class DefinitionCapability(Capability[DefinitionRequest, DefinitionResponse]):
     @cached_property
     def locate(self) -> LocateCapability:
         return LocateCapability(self.client)
@@ -55,31 +41,25 @@ async def __call__(self, req: DefinitionRequest) -> DefinitionResponse | None:
         locations: Sequence[Location] | None = None
         match req.mode:
             case "definition":
-                locations = await self.client.request_definition_locations(
-                    file_path, lsp_pos
-                )
+                locations = await ensure_capability(
+                    self.client, WithRequestDefinition
+                ).request_definition_locations(file_path, lsp_pos)
             case "declaration":
-                if not isinstance(self.client, WithRequestDeclaration):
-                    raise UnsupportedCapabilityError(
-                        "Client does not support 'textDocument/declaration'. "
-                        "To find declarations, you can: "
-                        "1) Use 'definition' mode (most language servers treat them similarly); "
-                        "2) For C/C++, check corresponding header files manually."
-                    )
-                locations = await self.client.request_declaration_locations(
-                    file_path, lsp_pos
-                )
+                locations = await ensure_capability(
+                    self.client,
+                    WithRequestDeclaration,
+                    error="To find declarations, you can: "
+                    "1) Use 'definition' mode (most language servers treat them similarly); "
+                    "2) For C/C++, check corresponding header files manually.",
+                ).request_declaration_locations(file_path, lsp_pos)
             case "type_definition":
-                if not isinstance(self.client, WithRequestTypeDefinition):
-                    raise UnsupportedCapabilityError(
-                        "Client does not support 'textDocument/typeDefinition'. "
-                        "To find type definitions, you can: "
-                        "1) Use 'definition' on the type name itself if visible; "
-                        "2) Use 'hover' to see the type name and then search for it."
-                    )
-                locations = await self.client.request_type_definition_locations(
-                    file_path, lsp_pos
-                )
+                locations = await ensure_capability(
+                    self.client,
+                    WithRequestTypeDefinition,
+                    error="To find type definitions, you can: "
+                    "1) Use 'definition' on the type name itself if visible; "
+                    "2) Use 'hover' to see the type name and then search for it.",
+                ).request_type_definition_locations(file_path, lsp_pos)
 
         if not locations:
             return None

src/lsap/capability/draft/__init__.py

@@ -1,28 +1,19 @@
 from __future__ import annotations
 
 from functools import cached_property
-from typing import Protocol, runtime_checkable
 
 from attrs import define
-from lsp_client.capability.request import WithRequestDocumentSymbol, WithRequestHover
+from lsp_client.capability.request import WithRequestHover
 
 from lsap.schema.hover import HoverRequest, HoverResponse
+from lsap.utils.capability import ensure_capability
 
-from .abc import Capability, ClientProtocol
+from .abc import Capability
 from .locate import LocateCapability
 
 
-@runtime_checkable
-class HoverClient(
-    WithRequestHover,
-    WithRequestDocumentSymbol,
-    ClientProtocol,
-    Protocol,
-): ...
-
-
 @define
-class HoverCapability(Capability[HoverClient, HoverRequest, HoverResponse]):
+class HoverCapability(Capability[HoverRequest, HoverResponse]):
     @cached_property
     def locate(self) -> LocateCapability:
         return LocateCapability(self.client)
@@ -32,7 +23,9 @@ async def __call__(self, req: HoverRequest) -> HoverResponse | None:
             return None
 
         file_path, lsp_pos = loc_resp.file_path, loc_resp.position.to_lsp()
-        hover = await self.client.request_hover(file_path, lsp_pos)
+        hover = await ensure_capability(self.client, WithRequestHover).request_hover(
+            file_path, lsp_pos
+        )
 
         if hover is None:
             return None