Provide a way for acting on responses to client->server requests

pyls/server.py

@@ -3,7 +3,7 @@
 import logging
 import uuid
 
-import jsonrpc
+from jsonrpc import jsonrpc2, JSONRPCResponseManager
 
 log = logging.getLogger(__name__)
 
@@ -14,6 +14,8 @@ class JSONRPCServer(object):
     def __init__(self, rfile, wfile):
         self.rfile = rfile
         self.wfile = wfile
+
+        self._callbacks = {}
         self._shutdown = False
 
     def exit(self):
@@ -27,56 +29,68 @@ def shutdown(self):
         log.debug("Server shut down, awaiting exit notification")
 
     def handle(self):
-        # VSCode wants us to keep the connection open, so let's handle messages in a loop
         while True:
             try:
                 data = self._read_message()
                 log.debug("Got message: %s", data)
 
                 if self._shutdown:
                     # Handle only the exit notification when we're shut down
-                    jsonrpc.JSONRPCResponseManager.handle(data, {'exit': self.exit})
+                    JSONRPCResponseManager.handle(data, {'exit': self.exit})
                     break
 
-                response = jsonrpc.JSONRPCResponseManager.handle(data, self)
-
-                if response is not None:
-                    self._write_message(response.data)
+                if isinstance(data, bytes):
+                    data = data.decode("utf-8")
+
+                msg = json.loads(data)
+                if 'method' in msg:
+                    # It's a notification or request
+                    # Dispatch to the thread pool for handling
+                    response = JSONRPCResponseManager.handle(data, self)
+                    if response is not None:
+                        self._write_message(response.data)
+                else:
+                    # Otherwise, it's a response message
+                    on_result, on_error = self._callbacks.pop(msg['id'])
+                    if 'result' in msg and on_result:
+                        on_result(msg['result'])
+                    elif 'error' in msg and on_error:
+                        on_error(msg['error'])
             except Exception:
                 log.exception("Language server exiting due to uncaught exception")
                 break
 
-    def call(self, method, params=None):
-        """ Call a method on the client. TODO: return the result. """
-        log.debug("Sending request %s: %s", method, params)
-        req = jsonrpc.jsonrpc2.JSONRPC20Request(method=method, params=params)
-        req._id = str(uuid.uuid4())
+    def call(self, method, params=None, on_result=None, on_error=None):
+        """Call a method on the client."""
+        msg_id = str(uuid.uuid4())
+        log.debug("Sending request %s: %s: %s", msg_id, method, params)
+        req = jsonrpc2.JSONRPC20Request(method=method, params=params)
+        req._id = msg_id
+
+        def _default_on_error(error):
+            log.error("Call to %s failed with %s", method, error)
+
+        if not on_error:
+            on_error = _default_on_error
+
+        self._callbacks[msg_id] = (on_result, on_error)
         self._write_message(req.data)
 
     def notify(self, method, params=None):
         """ Send a notification to the client, expects no response. """
         log.debug("Sending notification %s: %s", method, params)
-        req = jsonrpc.jsonrpc2.JSONRPC20Request(
+        req = jsonrpc2.JSONRPC20Request(
             method=method, params=params, is_notification=True
         )
         self._write_message(req.data)
 
-    def _content_length(self, 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))
-
     def _read_message(self):
         line = self.rfile.readline()
 
         if not line:
             raise EOFError()
 
-        content_length = self._content_length(line)
+        content_length = _content_length(line)
 
         # Blindly consume all header lines
         while line and line.strip():
@@ -98,3 +112,14 @@ def _write_message(self, msg):
         )
         self.wfile.write(response.encode('utf-8'))
         self.wfile.flush()
+
+
+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))

pyls/workspace.py

@@ -60,9 +60,11 @@ def update_document(self, doc_uri, change, version=None):
         self._docs[doc_uri].apply_change(change)
         self._docs[doc_uri].version = version
 
-    def apply_edit(self, edit):
-        # Note that lang_server.call currently doesn't return anything
-        return self._lang_server.call(self.M_APPLY_EDIT, {'edit': edit})
+    def apply_edit(self, edit, on_result=None, on_error=None):
+        return self._lang_server.call(
+            self.M_APPLY_EDIT, {'edit': edit},
+            on_result=on_result, on_error=on_error
+        )
 
     def publish_diagnostics(self, doc_uri, diagnostics):
         params = {'uri': doc_uri, 'diagnostics': diagnostics}