Add support for asynchronously handling requests

pyls/__main__.py

@@ -4,16 +4,13 @@
 import logging
 import logging.config
 import sys
-
-from . import language_server
-from .python_ls import PythonLanguageServer
+from .python_ls import start_io_lang_server, start_tcp_lang_server, PythonLanguageServer
 
 LOG_FORMAT = "%(asctime)s UTC - %(levelname)s - %(name)s - %(message)s"
 
 
 def add_arguments(parser):
     parser.description = "Python Language Server"
-
     parser.add_argument(
         "--tcp", action="store_true",
         help="Use TCP server instead of stdio"
@@ -51,10 +48,10 @@ def main():
     _configure_logger(args.verbose, args.log_config, args.log_file)
 
     if args.tcp:
-        language_server.start_tcp_lang_server(args.host, args.port, PythonLanguageServer)
+        start_tcp_lang_server(args.host, args.port, PythonLanguageServer)
     else:
         stdin, stdout = _binary_stdio()
-        language_server.start_io_lang_server(stdin, stdout, PythonLanguageServer)
+        start_io_lang_server(stdin, stdout, PythonLanguageServer)
 
 
 def _binary_stdio():

pyls/json_rpc_server.py

@@ -0,0 +1,115 @@
+# Copyright 2017 Palantir Technologies, Inc.
+import json
+import logging
+import threading
+
+from jsonrpc.jsonrpc2 import JSONRPC20Response
+from jsonrpc.jsonrpc import JSONRPCRequest
+from jsonrpc.exceptions import (
+    JSONRPCInvalidRequestException,
+)
+
+
+log = logging.getLogger(__name__)
+
+
+class JSONRPCServer(object):
+    """ Read/Write JSON RPC messages """
+
+    def __init__(self, rfile, wfile):
+        self.batch_messages = {}
+        self.rfile = rfile
+        self.wfile = wfile
+        self.write_lock = threading.Lock()
+
+    def close(self):
+        with self.write_lock:
+            self.wfile.close()
+            self.rfile.close()
+
+    def get_messages(self):
+        """Generator that produces well structured JSON RPC message.
+
+        Returns:
+            message: received message
+
+        Note:
+            This method is not thread safe and should only invoked from a single thread
+        """
+        while not self.rfile.closed:
+            request_str = self._read_message()
+
+            if request_str is None:
+                break
+            if isinstance(request_str, bytes):
+                request_str = request_str.decode("utf-8")
+
+            try:
+                try:
+                    message_blob = json.loads(request_str)
+                    message = JSONRPCRequest.from_data(message_blob)
+                except JSONRPCInvalidRequestException:
+                    # work around where JSONRPC20Reponse expects _id key
+                    message_blob['_id'] = message_blob['id']
+                    message = JSONRPC20Response(**message_blob)
+            except (KeyError, ValueError):
+                log.error("Could not parse message %s", request_str)
+                continue
+
+            yield message
+
+    def write_message(self, message):
+        """ Write message to out file descriptor
+
+        Args:
+            message (dict): body of the message to send
+        """
+        with self.write_lock:
+            if self.wfile.closed:
+                return
+            log.debug("Sending %s", message)
+            body = json.dumps(message, separators=(",", ":"))
+            content_length = len(body)
+            response = (
+                "Content-Length: {}\r\n"
+                "Content-Type: application/vscode-jsonrpc; charset=utf8\r\n\r\n"
+                "{}".format(content_length, body)
+            )
+            self.wfile.write(response.encode('utf-8'))
+            self.wfile.flush()
+
+    def _read_message(self):
+        """Reads the contents of a message
+
+        Returns:
+            body of message if parsable else None
+        """
+        line = self.rfile.readline()
+
+        if not line:
+            return None
+
+        content_length = _content_length(line)
+
+        # Blindly consume all header lines
+        while line and line.strip():
+            line = self.rfile.readline()
+
+        if not line:
+            return None
+
+        # Grab the body
+        return self.rfile.read(content_length)
+
+
+def _content_length(line):
+    """Extract the content length from an input line."""
+    if line.startswith(b'Content-Length: '):
+        _, value = line.split(b'Content-Length: ')
+        value = value.strip()
+        try:
+            return int(value)
+        except ValueError:
+            raise ValueError("Invalid Content-Length header: {}".format(value))
+
+    return None