Add a dedicated 'Ruff: Format document' action

astral-sh · Sep 27, 2023 · 794fc1f · 794fc1f
1 parent e270ae5
commit 794fc1f
Show file tree

Hide file tree

Showing 2 changed files with 17 additions and 40 deletions.
diff --git a/ruff_lsp/server.py b/ruff_lsp/server.py
@@ -682,20 +682,29 @@ async def apply_organize_imports(arguments: tuple[TextDocument]):
 
 if RUFF_EXPERIMENTAL_FORMATTER:
 
+    @LSP_SERVER.command("ruff.applyFormat")
+    async def apply_format(arguments: tuple[TextDocument]):
+        uri = arguments[0]["uri"]
+        text_document = LSP_SERVER.workspace.get_document(uri)
+        results = await _format_document_impl(text_document)
+        LSP_SERVER.apply_edit(
+            _create_workspace_edits(text_document, results),
+            "Ruff: Format document",
+        )
+
     @LSP_SERVER.feature(TEXT_DOCUMENT_FORMATTING)
     async def format_document(
         ls: server.LanguageServer,
         params: DocumentFormattingParams,
     ) -> list[TextEdit] | None:
-        return await _format_document_impl(ls, params)
+        uri = params.text_document.uri
+        document = ls.workspace.get_document(uri)
+        return await _format_document_impl(document)
 
 
 async def _format_document_impl(
-    language_server: server.LanguageServer,
-    params: DocumentFormattingParams,
+    document: workspace.Document,
 ) -> list[TextEdit]:
-    uri = params.text_document.uri
-    document = language_server.workspace.get_document(uri)
     result = await _run_format_on_document(document)
     return _result_to_edits(document, result)
 

diff --git a/tests/test_format.py b/tests/test_format.py
@@ -1,16 +1,6 @@
 from __future__ import annotations
 
-from pathlib import Path
-from typing import cast
-
 import pytest
-from lsprotocol.types import (
-    DocumentFormattingParams,
-    FormattingOptions,
-    TextDocumentIdentifier,
-    WorkspaceEdit,
-)
-from pygls import server
 from pygls.workspace import Workspace
 
 from ruff_lsp.server import _format_document_impl
@@ -24,37 +14,15 @@
 """
 
 
-class MockLanguageServer:
-    root: Path
-    applied_edits: list[WorkspaceEdit] = []
-
-    def __init__(self, root):
-        self.root = root
-
-    @property
-    def workspace(self) -> Workspace:
-        return Workspace(str(self.root))
-
-    def apply_edit(self, edit: WorkspaceEdit, _label: str | None = None) -> None:
-        """Currently unused, but we keep it around for future tests."""
-        self.applied_edits.append(edit)
-
-
 @pytest.mark.asyncio
 async def test_format(tmp_path):
     test_file = tmp_path.joinpath("main.py")
     test_file.write_text(original)
     uri = utils.as_uri(str(test_file))
 
-    mock_language_server = MockLanguageServer(tmp_path)
-    dummy_params = DocumentFormattingParams(
-        text_document=TextDocumentIdentifier(uri=uri),
-        options=FormattingOptions(tab_size=4, insert_spaces=True),
-        work_done_token=None,
-    )
+    workspace = Workspace(str(tmp_path))
+    document = workspace.get_document(uri)
 
-    result = await _format_document_impl(
-        cast(server.LanguageServer, mock_language_server), dummy_params
-    )
+    result = await _format_document_impl(document)
     [edit] = result
     assert edit.new_text == expected