From b0331b70d410ffbf3a58fab05689dbe820d3fb19 Mon Sep 17 00:00:00 2001
From: Stephen Macke <stephen.macke@gmail.com>
Date: Thu, 7 Sep 2023 15:15:20 -0700
Subject: [PATCH] Allow method handlers to return json rpc errors (#20)

---
 pylsp_jsonrpc/endpoint.py | 28 ++++++++++++++++++++++------
 1 file changed, 22 insertions(+), 6 deletions(-)

diff --git a/pylsp_jsonrpc/endpoint.py b/pylsp_jsonrpc/endpoint.py
index 9bbe4c8..0c9e62a 100644
--- a/pylsp_jsonrpc/endpoint.py
+++ b/pylsp_jsonrpc/endpoint.py
@@ -4,6 +4,7 @@
 import logging
 import uuid
 import sys
+from typing import Any, Dict, Mapping
 
 from concurrent import futures
 from .exceptions import (JsonRpcException, JsonRpcRequestCancelled,
@@ -175,6 +176,17 @@ def _handle_cancel_notification(self, msg_id):
         if request_future.cancel():
             log.debug("Cancelled request with id %s", msg_id)
 
+    @staticmethod
+    def _make_response_payload(header: Dict[str, Any], result: Any) -> Mapping[str, Any]:
+        # return type of 'Mapping' because it should not be mutated
+        # further from here
+        response = dict(header)
+        if isinstance(result, dict) and ('result' in result or 'error' in result):
+            response.update(result)
+        else:
+            response['result'] = result
+        return response
+
     def _handle_request(self, msg_id, method, params):
         """Handle a request from the client."""
         try:
@@ -195,11 +207,14 @@ def _handle_request(self, msg_id, method, params):
             handler_result.add_done_callback(self._request_callback(msg_id))
         else:
             log.debug("Got result from synchronous request handler: %s", handler_result)
-            self._consumer({
-                'jsonrpc': JSONRPC_VERSION,
-                'id': msg_id,
-                'result': handler_result
-            })
+            response = self._make_response_payload(
+                {
+                    'jsonrpc': JSONRPC_VERSION,
+                    'id': msg_id,
+                },
+                handler_result,
+            )
+            self._consumer(response)
 
     def _request_callback(self, request_id):
         """Construct a request callback for the given request ID."""
@@ -216,7 +231,8 @@ def callback(future):
             }
 
             try:
-                message['result'] = future.result()
+                result = future.result()
+                message = self._make_response_payload(message, result)
             except JsonRpcException as e:
                 log.exception("Failed to handle request %s", request_id)
                 message['error'] = e.to_dict()