add async interface

Author: davschul

pylsp_jsonrpc/endpoint.py

@@ -4,6 +4,7 @@
 import logging
 import uuid
 import sys
+import asyncio
 
 from concurrent import futures
 from .exceptions import (JsonRpcException, JsonRpcRequestCancelled,
@@ -12,6 +13,7 @@
 log = logging.getLogger(__name__)
 JSONRPC_VERSION = '2.0'
 CANCEL_METHOD = '$/cancelRequest'
+EXIT_METHOD = 'exit'
 
 
 class Endpoint:
@@ -35,9 +37,25 @@ def __init__(self, dispatcher, consumer, id_generator=lambda: str(uuid.uuid4()),
         self._client_request_futures = {}
         self._server_request_futures = {}
         self._executor_service = futures.ThreadPoolExecutor(max_workers=max_workers)
+        self._cancelledRequests = set()
+        self._messageQueue = None
+        self._consume_task = None
+
+    def init_async(self):
+        self._messageQueue = asyncio.Queue()
+        self._consume_task = asyncio.create_task(self.consume_task())
+
+    async def consume_task(self):
+        loop = asyncio.get_running_loop()
+        while loop.is_running():
+            message = await self._messageQueue.get()
+            await asyncio.to_thread(self.consume, message)
+            self._messageQueue.task_done()
 
     def shutdown(self):
         self._executor_service.shutdown()
+        if self._consume_task is not None:
+            self._consume_task.cancel()
 
     def notify(self, method, params=None):
         """Send a JSON RPC notification to the client.
@@ -94,6 +112,21 @@ def callback(future):
                 future.set_exception(JsonRpcRequestCancelled())
         return callback
 
+    async def consume_async(self, message):
+        """Consume a JSON RPC message from the client and put it into a queue.
+
+        Args:
+            message (dict): The JSON RPC message sent by the client
+        """
+        if message['method'] == CANCEL_METHOD:
+            self._cancelledRequests.add(message.get('params')['id'])
+
+        # The exit message needs to be handled directly since the stream cannot be closed asynchronously
+        if message['method'] == EXIT_METHOD:
+            self.consume(message)
+        else:
+            await self._messageQueue.put(message)
+
     def consume(self, message):
         """Consume a JSON RPC message from the client.
 
@@ -182,6 +215,9 @@ def _handle_request(self, msg_id, method, params):
         except KeyError as e:
             raise JsonRpcMethodNotFound.of(method) from e
 
+        if msg_id in self._cancelledRequests:
+            raise JsonRpcRequestCancelled()
+
         handler_result = handler(params)
 
         if callable(handler_result):

pylsp_jsonrpc/streams.py

@@ -3,6 +3,7 @@
 
 import logging
 import threading
+import asyncio
 
 try:
     import ujson as json
@@ -65,6 +66,31 @@ def _read_message(self):
         # Grab the body
         return self._rfile.read(content_length)
 
+    async def listen_async(self, message_consumer):
+        """Blocking call to listen for messages on the rfile.
+
+        Args:
+            message_consumer (fn): function that is passed each message as it is read off the socket.
+        """
+
+        while not self._rfile.closed:
+            try:
+                request_str = await asyncio.to_thread(self._read_message)
+                log.warning("got stdin")
+            except ValueError:
+                if self._rfile.closed:
+                    return
+                log.exception("Failed to read from rfile")
+
+            if request_str is None:
+                break
+
+            try:
+                await message_consumer(json.loads(request_str.decode('utf-8')))
+            except ValueError:
+                log.exception("Failed to parse JSON message %s", request_str)
+                continue
+
     @staticmethod
     def _content_length(line):
         """Extract the content length from an input line."""