feat: support formatting

Author: betaboon

pylsp_ruff/plugin.py

@@ -4,7 +4,7 @@
 import sys
 from pathlib import PurePath
 from subprocess import PIPE, Popen
-from typing import Dict, List
+from typing import Dict, Generator, List, Optional
 
 from lsprotocol.types import (
     CodeAction,
@@ -64,6 +64,34 @@ def pylsp_settings():
     return converter.unstructure(settings)
 
 
+@hookimpl(hookwrapper=True)
+def pylsp_format_document(workspace: Workspace, document: Document) -> Generator:
+    """
+    Provide formatting through ruff.
+
+    Parameters
+    ----------
+    workspace : pylsp.workspace.Workspace
+        Current workspace.
+    document : pylsp.workspace.Document
+        Document to apply ruff on.
+    """
+    log.debug(f"textDocument/formatting: {document}")
+    outcome = yield
+    results = outcome.get_result()
+    if results:
+        document.source = results[0]["new_text"]
+
+    new_text = run_ruff_format(workspace, document)
+    range = Range(
+        start=Position(line=0, character=0),
+        end=Position(line=len(document.lines), character=0),
+    )
+    text_edit = TextEdit(range=range, new_text=new_text)
+
+    outcome.force_result(converter.unstructure([text_edit]))
+
+
 @hookimpl
 def pylsp_lint(workspace: Workspace, document: Document) -> List[Dict]:
     """
@@ -315,7 +343,23 @@ def run_ruff_fix(workspace: Workspace, document: Document) -> str:
     return result
 
 
-def run_ruff(workspace: Workspace, document: Document, fix: bool = False) -> str:
+def run_ruff_format(workspace: Workspace, document: Document) -> str:
+    settings = load_settings(workspace, document)
+    extra_arguments = []
+    if settings.format:
+        extra_arguments.append(f"--fixable={','.join(settings.format)}")
+    else:
+        extra_arguments.append("--unfixable=ALL")
+    result = run_ruff(workspace, document, fix=True, extra_arguments=extra_arguments)
+    return result
+
+
+def run_ruff(
+    workspace: Workspace,
+    document: Document,
+    fix: bool = False,
+    extra_arguments: Optional[List[str]] = None,
+) -> str:
     """
     Run ruff on the given document and the given arguments.
 
@@ -327,14 +371,16 @@ def run_ruff(workspace: Workspace, document: Document, fix: bool = False) -> str
         File to run ruff on.
     fix : bool
         Whether to run fix or no-fix.
+    extra_arguments : List[str]
+        Extra arguments to pass to ruff.
 
     Returns
     -------
     String containing the result in json format.
     """
     settings = load_settings(workspace, document)
     executable = settings.executable
-    arguments = build_arguments(document, settings, fix)
+    arguments = build_arguments(document, settings, fix, extra_arguments)
 
     log.debug(f"Calling {executable} with args: {arguments} on '{document.path}'")
     try:
@@ -358,6 +404,7 @@ def build_arguments(
     document: Document,
     settings: PluginSettings,
     fix: bool = False,
+    extra_arguments: Optional[List[str]] = None,
 ) -> List[str]:
     """
     Build arguments for ruff.
@@ -368,6 +415,10 @@ def build_arguments(
         Document to apply ruff on.
     settings : PluginSettings
         Settings to use for arguments to pass to ruff.
+    fix : bool
+        Whether to execute with --fix.
+    extra_arguments : List[str]
+        Extra arguments to pass to ruff.
 
     Returns
     -------
@@ -416,6 +467,9 @@ def build_arguments(
                 continue
             args.append(f"--ignore={','.join(errors)}")
 
+    if extra_arguments:
+        args.extend(extra_arguments)
+
     args.extend(["--", "-"])
 
     return args
@@ -456,6 +510,7 @@ def load_settings(workspace: Workspace, document: Document) -> PluginSettings:
             executable=plugin_settings.executable,
             extend_ignore=plugin_settings.extend_ignore,
             extend_select=plugin_settings.extend_select,
+            format=plugin_settings.format,
         )
 
     return plugin_settings

pylsp_ruff/settings.py

@@ -22,6 +22,8 @@ class PluginSettings:
     extend_ignore: Optional[List[str]] = None
     per_file_ignores: Optional[Dict[str, List[str]]] = None
 
+    format: Optional[List[str]] = None
+
 
 def to_camel_case(snake_str: str) -> str:
     components = snake_str.split("_")

tests/test_code_actions.py

@@ -2,6 +2,7 @@
 # Copyright 2021- Python Language Server Contributors.
 
 import tempfile
+from textwrap import dedent
 from typing import List
 from unittest.mock import Mock
 
@@ -26,16 +27,20 @@ def workspace(tmp_path):
     return ws
 
 
-codeaction_str = r"""
-import os
-def f():
-    a = 2
-"""
+codeaction_str = dedent(
+    """
+    import os
+    def f():
+        a = 2
+    """
+)
 
-import_str = r"""
-import pathlib
-import os
-"""
+import_str = dedent(
+    """
+    import pathlib
+    import os
+    """
+)
 
 codeactions = [
     "Ruff (F401): Remove unused import: `os`",
@@ -102,13 +107,40 @@ def test_import_action(workspace):
 
 
 def test_fix_all(workspace):
+    expected_str = dedent(
+        """
+        def f():
+            pass
+        """
+    )
     _, doc = temp_document(codeaction_str, workspace)
+    fixed_str = ruff_lint.run_ruff_fix(workspace, doc)
+    assert fixed_str == expected_str
+
+
+def test_format_document_default_settings(workspace):
+    _, doc = temp_document(import_str, workspace)
+    formatted_str = ruff_lint.run_ruff_format(workspace, doc)
+    assert formatted_str == import_str
 
-    fix_all = ruff_lint.run_ruff_fix(workspace, doc)
-    assert (
-        fix_all
-        == r"""
-def f():
-    pass
-"""
+
+def test_format_document_settings(workspace):
+    expected_str = dedent(
+        """
+        import os
+        import pathlib
+        """
+    )
+    workspace._config.update(
+        {
+            "plugins": {
+                "ruff": {
+                    "select": ["I"],
+                    "format": ["I001"],
+                }
+            }
+        }
     )
+    _, doc = temp_document(import_str, workspace)
+    formatted_str = ruff_lint.run_ruff_format(workspace, doc)
+    assert formatted_str == expected_str