Add a new /lsp endpoint to blackd

setup.py

@@ -86,7 +86,7 @@ def get_long_description() -> str:
         "mypy_extensions>=0.4.3",
     ],
     extras_require={
-        "d": ["aiohttp>=3.7.4"],
+        "d": ["aiohttp>=3.7.4", "aiohttp-json-rpc>=0.13.3"],
         "colorama": ["colorama>=0.4.3"],
         "python2": ["typed-ast>=1.4.2"],
         "uvloop": ["uvloop>=0.15.2"],

src/blackd/__init__.py

@@ -6,14 +6,16 @@
 from multiprocessing import freeze_support
 from typing import Set, Tuple
 
+
 try:
     from aiohttp import web
     from .middlewares import cors
+    from .lsp import make_lsp_handler
 except ImportError as ie:
     raise ImportError(
-        f"aiohttp dependency is not installed: {ie}. "
+        f"A blackd dependency is not installed: {ie}. "
         + "Please re-install black with the '[d]' extra install "
-        + "to obtain aiohttp_cors: `pip install black[d]`"
+        + "to obtain it: `pip install black[d]`"
     ) from None
 
 import black
@@ -71,7 +73,12 @@ def make_app() -> web.Application:
         middlewares=[cors(allow_headers=(*BLACK_HEADERS, "Content-Type"))]
     )
     executor = ProcessPoolExecutor()
-    app.add_routes([web.post("/", partial(handle, executor=executor))])
+    app.add_routes(
+        [
+            web.post("/", partial(handle, executor=executor)),
+            web.view("/lsp", make_lsp_handler(executor)),
+        ]
+    )
     return app
 
 

src/blackd/lsp.py

@@ -0,0 +1,133 @@
+from functools import partial
+import asyncio
+from json.decoder import JSONDecodeError
+import os
+from pathlib import Path
+from aiohttp import web
+from aiohttp_json_rpc import JsonRpc, RpcInvalidParamsError
+from aiohttp_json_rpc.protocol import JsonRpcMsgTyp
+from aiohttp_json_rpc.rpc import JsonRpcRequest
+from concurrent.futures import Executor
+
+from typing import Awaitable, Callable, Generator, List, Optional
+from typing_extensions import TypedDict
+from difflib import SequenceMatcher
+from urllib.parse import urlparse
+from urllib.request import url2pathname
+import black
+
+# Reference: https://bit.ly/2XScAZF
+DocumentUri = str
+
+
+class TextDocumentIdentifier(TypedDict):
+    """"""
+
+    uri: DocumentUri
+
+
+class FormattingOptions(TypedDict):
+    """Reference: https://bit.ly/3CPqmvk"""
+
+    tabSize: int
+    insertSpaces: bool
+    trimTrailingWhitespace: Optional[bool]
+    insertFinalNewline: Optional[bool]
+    trimFinalNewlines: Optional[bool]
+
+
+class DocumentFormattingParams(TypedDict):
+    """Reference: https://bit.ly/3ibxWZk"""
+
+    textDocument: TextDocumentIdentifier
+    options: FormattingOptions
+
+
+class Position(TypedDict):
+    """Reference: https://bit.ly/3CQDNuX"""
+
+    line: int
+    character: int
+
+
+class Range(TypedDict):
+    """Reference: https://bit.ly/3zKxWp4"""
+
+    start: Position
+    end: Position
+
+
+class TextEdit(TypedDict):
+    """Reference: https://bit.ly/3AJCFsF"""
+
+    range: Range
+    newText: str
+
+
+def make_lsp_handler(
+    executor: Executor,
+) -> Callable[[web.Request], Awaitable[web.Response]]:
+    rpc = JsonRpc()
+
+    async def formatting_handler(request: web.Request) -> web.Response:
+        return await handle_formatting(executor, request)
+
+    rpc.add_methods(
+        ("", formatting_handler, "textDocument/formatting"),
+    )
+    return rpc.handle_request  # type: ignore
+
+
+def uri_to_path(uri_str: str) -> Path:
+    uri = urlparse(uri_str)
+    if uri.scheme != "file":
+        raise RpcInvalidParamsError(message="only file:// uri scheme is supported")
+    return Path("{0}{0}{1}{0}".format(os.path.sep, uri.netloc)) / url2pathname(uri.path)
+
+
+def format(src_path: os.PathLike) -> List[TextEdit]:
+    def gen() -> Generator[TextEdit, None, None]:
+        with open(src_path, "rb") as buf:
+            src, encoding, newline = black.decode_bytes(buf.read())
+        try:
+            formatted_str = black.format_file_contents(
+                src, fast=True, mode=black.Mode()
+            )
+        except black.NothingChanged:
+            return
+        except JSONDecodeError as e:
+            raise RpcInvalidParamsError(
+                message="File cannot be parsed as a Jupyter notebook"
+            ) from e
+        cmp = SequenceMatcher(a=src, b=formatted_str)
+        for op, i1, i2, j1, j2 in cmp.get_opcodes():
+            if op == "equal":
+                continue
+
+            rng = Range(start=offset_to_pos(i1, src), end=offset_to_pos(i2, src))
+
+            if op in {"insert", "replace"}:
+                yield TextEdit(range=rng, newText=formatted_str[j1:j2])
+            elif op == "delete":
+                yield TextEdit(range=rng, newText="")
+
+    return list(gen())
+
+
+def offset_to_pos(offset: int, src: str) -> Position:
+    line = src.count("\n", 0, offset)
+    last_nl = src.rfind("\n", 0, offset)
+    character = offset if last_nl == -1 else offset - last_nl
+    return Position(line=line, character=character)
+
+
+async def handle_formatting(
+    executor: Executor, request: JsonRpcRequest
+) -> Optional[List[TextEdit]]:
+    if request.msg.type != JsonRpcMsgTyp.REQUEST:
+        raise RpcInvalidParamsError
+
+    params: DocumentFormattingParams = request.msg.data["params"]
+    path = uri_to_path(params["textDocument"]["uri"])
+    loop = asyncio.get_event_loop()
+    return await loop.run_in_executor(executor, partial(format, path))

tests/test_blackd.py

@@ -1,21 +1,46 @@
 import re
+from pathlib import Path
+from tempfile import TemporaryDirectory
+from typing import Any, Optional, TypeVar
 from unittest.mock import patch
 
-from click.testing import CliRunner
 import pytest
+from click.testing import CliRunner
 
-from tests.util import read_data, DETERMINISTIC_HEADER
+from tests.util import DETERMINISTIC_HEADER, read_data
 
 try:
     import blackd
-    from aiohttp.test_utils import AioHTTPTestCase, unittest_run_loop
     from aiohttp import web
+    from aiohttp.client_ws import ClientWebSocketResponse
+    from aiohttp.test_utils import AioHTTPTestCase, unittest_run_loop
+    from aiohttp_json_rpc.protocol import JsonRpcMsgTyp, decode_msg, encode_request
+    from blackd.lsp import (
+        DocumentFormattingParams,
+        TextDocumentIdentifier,
+        FormattingOptions,
+    )
 except ImportError:
     has_blackd_deps = False
 else:
     has_blackd_deps = True
 
 
+if has_blackd_deps:
+
+    def make_formatting_options() -> FormattingOptions:
+        return FormattingOptions(
+            tabSize=999,
+            insertSpaces=True,
+            trimTrailingWhitespace=None,
+            insertFinalNewline=None,
+            trimFinalNewlines=None,
+        )
+
+
+JsonRpcParams = TypeVar("JsonRpcParams")
+
+
 @pytest.mark.blackd
 class BlackDTestCase(AioHTTPTestCase):
     def test_blackd_main(self) -> None:
@@ -185,3 +210,81 @@ async def test_cors_headers_present(self) -> None:
         response = await self.client.post("/", headers={"Origin": "*"})
         self.assertIsNotNone(response.headers.get("Access-Control-Allow-Origin"))
         self.assertIsNotNone(response.headers.get("Access-Control-Expose-Headers"))
+
+    async def _call_json_rpc(
+        self,
+        ws: ClientWebSocketResponse,
+        method: str,
+        params: Optional[JsonRpcParams] = None,
+        expect_success: bool = True,
+    ) -> Any:
+        await ws.send_str(encode_request(method, id=0, params=params))
+        raw_msg = await ws.receive_str(timeout=10.0)
+        msg = decode_msg(raw_msg)
+        if expect_success:
+            self.assertEqual(
+                msg.type,
+                JsonRpcMsgTyp.RESULT,
+                msg="Didn't receive RESULT response",
+            )
+            return msg.data["result"]
+        else:
+            return msg.data
+
+    @unittest_run_loop
+    async def test_lsp_formatting_method_is_available(self) -> None:
+        async with self.client.ws_connect("/lsp") as ws:
+            self.assertIn(
+                "textDocument/formatting", await self._call_json_rpc(ws, "get_methods")
+            )
+
+    @unittest_run_loop
+    async def test_lsp_formatting_changes(self) -> None:
+        with TemporaryDirectory() as dir:
+            inputfile = Path(dir) / "somefile.py"
+            inputfile.write_bytes(b"print('hello world')")
+            async with self.client.ws_connect("/lsp") as ws:
+                resp = await self._call_json_rpc(
+                    ws,
+                    "textDocument/formatting",
+                    DocumentFormattingParams(
+                        textDocument=TextDocumentIdentifier(uri=inputfile.as_uri()),
+                        options=make_formatting_options(),
+                    ),
+                )
+                self.assertIsInstance(resp, list)
+                self.assertGreaterEqual(len(resp), 1)
+
+    @unittest_run_loop
+    async def test_lsp_formatting_no_changes(self) -> None:
+        with TemporaryDirectory() as dir:
+            inputfile = Path(dir) / "somefile.py"
+            inputfile.write_bytes(b'print("hello world")\n')
+            async with self.client.ws_connect("/lsp") as ws:
+                resp = await self._call_json_rpc(
+                    ws,
+                    "textDocument/formatting",
+                    DocumentFormattingParams(
+                        textDocument=TextDocumentIdentifier(uri=inputfile.as_uri()),
+                        options=make_formatting_options(),
+                    ),
+                )
+                self.assertIsInstance(resp, list)
+                self.assertEqual(resp, [])
+
+    @unittest_run_loop
+    async def test_lsp_formatting_syntax_error(self) -> None:
+        with TemporaryDirectory() as dir:
+            inputfile = Path(dir) / "somefile.py"
+            inputfile.write_bytes(b"print(")
+            async with self.client.ws_connect("/lsp") as ws:
+                resp = await self._call_json_rpc(
+                    ws,
+                    "textDocument/formatting",
+                    DocumentFormattingParams(
+                        textDocument=TextDocumentIdentifier(uri=inputfile.as_uri()),
+                        options=make_formatting_options(),
+                    ),
+                    expect_success=False,
+                )
+                self.assertIn("error", resp)