Add start_ws method to pygls' LanguageClient

pygls/client.py

@@ -20,11 +20,13 @@
 import json
 import logging
 import re
+import sys
 import typing
 from threading import Event
 
 from pygls.exceptions import PyglsError, JsonRpcException, JsonRpcInternalError
 from pygls.protocol import JsonRPCProtocol, default_converter
+from pygls.server import WebSocketTransportAdapter
 
 if typing.TYPE_CHECKING:
     from typing import Any
@@ -35,6 +37,7 @@
 
     from cattrs import Converter
 
+    from websockets.asyncio.client import ClientConnection
 
 logger = logging.getLogger(__name__)
 
@@ -116,6 +119,26 @@ async def start_tcp(self, host: str, port: int):
 
         self._async_tasks.extend([connection])
 
+    async def start_ws(self, host: str, port: int):
+        """Start communicating with a server over WebSockets."""
+
+        try:
+            from websockets.asyncio.client import connect
+        except ImportError:
+            logger.exception(
+                "Run `pip install pygls[ws]` to install dependencies required for websockets."
+            )
+            sys.exit(1)
+
+        uri = f"ws://{host}:{port}"
+        websocket = await connect(uri)
+
+        self.protocol._send_only_body = True
+        self.protocol.connection_made(WebSocketTransportAdapter(websocket))  # type: ignore
+
+        connection = asyncio.create_task(self.run_websocket(websocket))
+        self._async_tasks.extend([connection])
+
     async def run_async(self, reader: asyncio.StreamReader):
         """Run the main message processing loop, asynchronously"""
 
@@ -154,6 +177,31 @@ async def run_async(self, reader: asyncio.StreamReader):
                 # Reset
                 content_length = 0
 
+    async def run_websocket(self, websocket: ClientConnection):
+        """Run the main message processing loop, over websockets."""
+
+        try:
+            from websockets.exceptions import ConnectionClosedOK
+        except ImportError:
+            logger.exception(
+                "Run `pip install pygls[ws]` to install dependencies required for websockets."
+            )
+            return
+
+        while not self._stop_event.is_set():
+            try:
+                data = await websocket.recv(decode=False)
+            except ConnectionClosedOK:
+                self._stop_event.set()
+                break
+
+            try:
+                message = json.loads(data, object_hook=self.protocol.structure_message)
+                self.protocol.handle_message(message)
+            except Exception as exc:
+                logger.exception("Unable to handle message")
+                self._report_server_error(exc, JsonRpcInternalError)
+
     async def _server_exit(self):
         """Cleanup handler that runs when the server process managed by the client exits"""
         if self._server is None:

pygls/server.py

@@ -130,13 +130,12 @@ class WebSocketTransportAdapter:
     Write method sends data via the WebSocket interface.
     """
 
-    def __init__(self, ws, loop):
+    def __init__(self, ws):
         self._ws = ws
-        self._loop = loop
 
     def close(self) -> None:
         """Stop the WebSocket server."""
-        self._ws.close()
+        asyncio.ensure_future(self._ws.close())
 
     def write(self, data: Any) -> None:
         """Create a task to write specified data into a WebSocket."""
@@ -290,7 +289,7 @@ def start_ws(self, host: str, port: int) -> None:
 
         async def connection_made(websocket, _):
             """Handle new connection wrapped in the WebSocket."""
-            self.protocol.transport = WebSocketTransportAdapter(websocket, self.loop)
+            self.protocol.transport = WebSocketTransportAdapter(websocket)
             async for message in websocket:
                 self.protocol.handle_message(
                     json.loads(message, object_hook=self.protocol.structure_message)

pyproject.toml

@@ -21,7 +21,7 @@ readme = "README.md"
 python = ">=3.9"
 cattrs = ">=23.1.2"
 lsprotocol = "2024.0.0a2"
-websockets = { version = ">=11.0.3", optional = true }
+websockets = { version = ">=13.0", optional = true }
 
 [tool.poetry.extras]
 ws = ["websockets"]
@@ -68,6 +68,7 @@ poetry_lock_check = "poetry check"
 sequence = [
      { cmd = "pytest --cov" },
      { cmd = "pytest tests/e2e --lsp-transport tcp" },
+     { cmd = "pytest tests/e2e --lsp-transport websockets" },
 ]
 ignore_fail = "return_non_zero"
 

tests/conftest.py

@@ -133,7 +133,7 @@ def pytest_addoption(parser):
         dest="lsp_transport",
         action="store",
         default="stdio",
-        choices=("stdio", "tcp"),
+        choices=("stdio", "tcp", "websockets"),
         help="Choose the transport to use with servers under test.",
     )
 
@@ -212,6 +212,15 @@ async def fn(
             await asyncio.sleep(1)
             await client.start_tcp(host, port)
 
+        elif transport == "websockets":
+            # TODO: Make host/port configurable?
+            host, port = "localhost", 8888
+            server_cmd.extend(["--ws", "--host", host, "--port", f"{port}"])
+
+            server = await asyncio.create_subprocess_exec(*server_cmd)
+            await asyncio.sleep(1)
+            await client.start_ws(host, port)
+
         else:
             raise NotImplementedError(f"Unsupported transport: {transport!r}")
 

tests/e2e/test_threaded_handlers.py

@@ -130,7 +130,7 @@ async def test_countdown_threaded(
 ):
     """Ensure that the countdown threaded command is working as expected."""
 
-    if IS_WIN and transport == "tcp":
+    if (IS_WIN and transport == "tcp") or transport == "websockets":
         pytest.skip("see https://github.com/openlawlibrary/pygls/issues/502")
 
     client, initialize_result = threaded_handlers