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,
@@ -35,6 +36,17 @@ 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 = asyncio.Queue()
+
+    async def consume_task(self):
+        log.warning("starting task")
+        loop = asyncio.get_running_loop()
+        while True:
+            message = await self._messageQueue.get()
+            await loop.run_in_executor(None, self.consume, message)
+            log.warning("got message in task")
+            self._messageQueue.task_done()
 
     def shutdown(self):
         self._executor_service.shutdown()
@@ -94,7 +106,15 @@ def callback(future):
                 future.set_exception(JsonRpcRequestCancelled())
         return callback
 
+    async def consume_async(self, message):
+        log.warning("got message put in queue")
+        if message['method'] == CANCEL_METHOD:
+            self._cancelledRequests.add(message.get('params')['id'])
+        await self._messageQueue.put(message)
+
+
     def consume(self, message):
+        log.warning("consume message")
         """Consume a JSON RPC message from the client.
 
         Args:
@@ -182,6 +202,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,8 @@
 
 import logging
 import threading
+import asyncio
+from concurrent.futures import Future
 
 try:
     import ujson as json
@@ -65,6 +67,55 @@ 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 self._read_message_async()
+            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
+
+    async def _read_message_async(self):
+        """Reads the contents of a message.
+
+        Returns:
+            body of message if parsable else None
+        """
+
+        loop = asyncio.get_running_loop()
+        line = await loop.run_in_executor(None, self._rfile.readline)
+
+        if not line:
+            return None
+
+        content_length = self._content_length(line)
+
+        # Blindly consume all header lines
+        while line and line.strip():
+            line = await loop.run_in_executor(None, self._rfile.readline)
+
+        if not line:
+            return None
+
+        # Grab the body
+        return await loop.run_in_executor(None, self._rfile.read, content_length)
+
     @staticmethod
     def _content_length(line):
         """Extract the content length from an input line."""