Experimental support for starting a generic LSP server (#217)

Issue #216

docs/guide.rst

@@ -88,6 +88,7 @@ Integrations
 - `GitHub Actions`_
 - `pre-commit`_
 - `Visual Studio Code`_
+- `Generic LSP Server`_
 
 
 GitHub Actions
@@ -252,3 +253,26 @@ For more details, or to install the extension, see the Visual Studio Marketplace
 .. image:: https://img.shields.io/badge/-Install%20Now-107C10?style=for-the-badge&logo=visualstudiocode
     :alt: Install VS Code extension now
     :target: vscode:extension/omnilib.ufmt
+
+
+Generic LSP Server
+~~~~~~~~~~~~~~~~~~
+
+**Experimental:**
+µfmt includes a generic LSP formatting server that can be used with any editor
+or IDE that supports the `Language Server Protocol`__.
+
+.. __: https://microsoft.github.io/language-server-protocol/specifications/specification-current/
+
+Extra dependencies are needed to run the µfmt LSP, and can be installed with
+the ``[lsp]`` package extras with pip:
+
+.. code-block:: shell-session
+
+    $ pip install ufmt[lsp]
+
+The generic LSP can then be started using the ``lsp`` subcommand:
+
+.. code-block:: shell-session
+
+    $ ufmt lsp [--tcp | --ws] [--port <port>]

pyproject.toml

@@ -30,6 +30,9 @@ dependencies = [
 ]
 
 [project.optional-dependencies]
+lsp = [
+    "pygls >= 1.3",
+]
 ruff = [
     "ruff-api>=0.0.5",
 ]
@@ -40,6 +43,7 @@ dev = [
     "flit==3.9.0",
     "flake8==7.0.0",
     "mypy==1.10.0",
+    "pygls==1.3.1",
     "ruff-api==0.0.5",
     "usort==1.0.8.post1",
 ]

ufmt/cli.py

@@ -176,3 +176,23 @@ def format(ctx: click.Context, names: List[str]) -> None:
     _, error = echo_results(results, quiet=options.quiet)
     if error:
         ctx.exit(1)
+
+
+@main.command()
+@click.pass_context
+@click.option("--tcp", is_flag=True)
+@click.option("--ws", is_flag=True)
+@click.option("--port", type=int, default=8971)
+def lsp(ctx: click.Context, tcp: bool, ws: bool, port: int) -> None:
+    """Experimental: start an LSP formatting server"""
+    from .lsp import ufmt_lsp
+
+    options: Options = ctx.obj
+    server = ufmt_lsp(root=options.root)
+
+    if tcp:
+        server.start_tcp("localhost", port)
+    elif ws:
+        server.start_ws("localhost", port)
+    else:
+        server.start_io()

ufmt/lsp.py

@@ -0,0 +1,144 @@
+# Copyright 2022 Amethyst Reese, Tim Hatch
+# Licensed under the MIT license
+
+import logging
+from pathlib import Path
+from typing import List, Literal, Optional
+
+from lsprotocol.types import (
+    DocumentFormattingParams,
+    MessageType,
+    Position,
+    Range,
+    TEXT_DOCUMENT_FORMATTING,
+    TextEdit,
+)
+
+from pygls.server import LanguageServer
+from pygls.workspace import TextDocument
+
+from .__version__ import __version__
+from .config import load_config
+from .core import ufmt_bytes
+from .types import (
+    BlackConfigFactory,
+    Encoding,
+    FileContent,
+    Processor,
+    Result,
+    SkipFormatting,
+    UfmtConfigFactory,
+    UsortConfig,
+    UsortConfigFactory,
+)
+from .util import make_black_config
+
+ServerType = Literal["stdin", "tcp", "ws"]
+
+LOG = logging.getLogger(__name__)
+
+
+def _wrap_ufmt_bytes(  # pragma: nocover
+    path: Path,
+    content: FileContent,
+    *,
+    encoding: Encoding,
+    ufmt_config_factory: Optional[UfmtConfigFactory] = None,
+    black_config_factory: Optional[BlackConfigFactory] = None,
+    usort_config_factory: Optional[UsortConfigFactory] = None,
+    pre_processor: Optional[Processor] = None,
+    post_processor: Optional[Processor] = None,
+    root: Optional[Path] = None,
+) -> Result:
+    try:
+        ufmt_config = (ufmt_config_factory or load_config)(path, root)
+        black_config = (black_config_factory or make_black_config)(path)
+        usort_config = (usort_config_factory or UsortConfig.find)(path)
+
+        result = Result(path)
+
+        dst_contents = ufmt_bytes(
+            path,
+            content,
+            encoding=encoding,
+            ufmt_config=ufmt_config,
+            black_config=black_config,
+            usort_config=usort_config,
+            pre_processor=pre_processor,
+            post_processor=post_processor,
+        )
+        result.after = dst_contents
+
+    except SkipFormatting as e:
+        result.after = content
+        result.skipped = str(e) or True
+        return result
+
+    except Exception as e:
+        result.error = e
+        return result
+
+    return result
+
+
+def ufmt_lsp(  # pragma: nocover
+    *,
+    ufmt_config_factory: Optional[UfmtConfigFactory] = None,
+    black_config_factory: Optional[BlackConfigFactory] = None,
+    usort_config_factory: Optional[UsortConfigFactory] = None,
+    pre_processor: Optional[Processor] = None,
+    post_processor: Optional[Processor] = None,
+    root: Optional[Path] = None,
+) -> LanguageServer:
+    """
+    Prepare an LSP server instance.
+
+    Keyword arguments have the same semantics as :func:`ufmt_paths`.
+
+    Returns a LanguageServer object. User must call the ``start_io()`` or
+    ``start_tcp()`` method with appropriate arguments to actually start the LSP server.
+    """
+    server = LanguageServer("ufmt-lsp", __version__)
+
+    @server.feature(TEXT_DOCUMENT_FORMATTING)
+    def lsp_format_document(
+        ls: LanguageServer, params: DocumentFormattingParams
+    ) -> Optional[List[TextEdit]]:
+        document: TextDocument = ls.workspace.get_text_document(
+            params.text_document.uri
+        )
+        path = Path(document.path).resolve()
+
+        # XXX: we're assuming everything is UTF-8 because LSP doesn't track encodings...
+        encoding: Encoding = "utf-8"
+        content = document.source.encode(encoding)
+
+        result = _wrap_ufmt_bytes(
+            path,
+            content,
+            encoding=encoding,
+            ufmt_config_factory=ufmt_config_factory,
+            black_config_factory=black_config_factory,
+            usort_config_factory=usort_config_factory,
+            pre_processor=pre_processor,
+            post_processor=post_processor,
+            root=root,
+        )
+
+        if result.error:
+            ls.show_message(
+                f"Formatting failed: {str(result.error)}", MessageType.Error
+            )
+            return []
+
+        return [
+            TextEdit(
+                Range(
+                    Position(0, 0),
+                    Position(len(document.lines), 0),
+                ),
+                result.after.decode(encoding),
+            ),
+        ]
+
+    return server

ufmt/tests/cli.py

@@ -487,3 +487,26 @@ def reset() -> None:
             self.assertIn("1 file formatted, 1 file already formatted", result.stderr)
             self.assertEqual(POORLY_FORMATTED_CODE, beta.read_text())
             self.assertEqual(CORRECTLY_FORMATTED_CODE, kappa.read_text())
+
+    @patch("ufmt.lsp.ufmt_lsp")  # dynamic import, patch at definition
+    def test_lsp(self, lsp_mock: Mock) -> None:
+        with self.subTest("default"):
+            lsp_mock.reset_mock()
+            self.runner.invoke(main, ["lsp"])
+
+            lsp_mock.assert_called_with(root=None)
+            lsp_mock.return_value.start_io.assert_called_with()
+
+        with self.subTest("tcp"):
+            lsp_mock.reset_mock()
+            self.runner.invoke(main, ["lsp", "--tcp", "--port", "4567"])
+
+            lsp_mock.assert_called_with(root=None)
+            lsp_mock.return_value.start_tcp.assert_called_with("localhost", 4567)
+
+        with self.subTest("ws"):
+            lsp_mock.reset_mock()
+            self.runner.invoke(main, ["lsp", "--ws", "--port", "8765"])
+
+            lsp_mock.assert_called_with(root=None)
+            lsp_mock.return_value.start_ws.assert_called_with("localhost", 8765)