From 43499c535df87e46c7f66e341303bb1b0dc76a14 Mon Sep 17 00:00:00 2001
From: Micha Reiser <micha@reiser.io>
Date: Fri, 2 Feb 2024 18:11:26 +0100
Subject: [PATCH] Use --range, extract `_format_document_impl`

---
 ruff_lsp/server.py   | 75 +++++++++++++++++++++-----------------------
 tests/test_format.py | 11 +++----
 2 files changed, 40 insertions(+), 46 deletions(-)

diff --git a/ruff_lsp/server.py b/ruff_lsp/server.py
index d2d33bb..a13afba 100755
--- a/ruff_lsp/server.py
+++ b/ruff_lsp/server.py
@@ -1215,14 +1215,41 @@ async def apply_format(arguments: tuple[TextDocument]):
 
 
 @LSP_SERVER.feature(TEXT_DOCUMENT_FORMATTING)
-async def format_document(
-    params: DocumentFormattingParams, range: Range | None = None
+async def format_document(params: DocumentFormattingParams) -> list[TextEdit] | None:
+    return await _format_document_impl(params, None)
+
+
+@LSP_SERVER.feature(
+    TEXT_DOCUMENT_RANGE_FORMATTING,
+    DocumentRangeFormattingRegistrationOptions(
+        document_selector=[
+            TextDocumentFilter_Type1(language="python", scheme="file"),
+            TextDocumentFilter_Type1(language="python", scheme="untitled"),
+        ],
+        ranges_support=False,
+        work_done_progress=False,
+    ),
+)
+async def format_document_range(
+    params: DocumentRangeFormattingParams,
+) -> list[TextEdit] | None:
+    return await _format_document_impl(
+        DocumentFormattingParams(
+            params.text_document, params.options, params.work_done_token
+        ),
+        params.range,
+    )
+
+
+async def _format_document_impl(
+    params: DocumentFormattingParams, range: Range | None
 ) -> list[TextEdit] | None:
     # For a Jupyter Notebook, this request can only format a single cell as the
     # request itself can only act on a text document. A cell in a Notebook is
     # represented as a text document. The "Notebook: Format notebook" action calls
     # this request for every cell.
     document = Document.from_cell_or_text_uri(params.text_document.uri)
+
     settings = _get_settings_by_document(document.path)
 
     # We don't support range formatting of notebooks yet but VS Code
@@ -1253,28 +1280,6 @@ async def format_document(
         )
 
 
-@LSP_SERVER.feature(
-    TEXT_DOCUMENT_RANGE_FORMATTING,
-    DocumentRangeFormattingRegistrationOptions(
-        document_selector=[
-            TextDocumentFilter_Type1(language="python", scheme="file"),
-            TextDocumentFilter_Type1(language="python", scheme="untitled"),
-        ],
-        ranges_support=False,
-        work_done_progress=False,
-    ),
-)
-async def format_document_range(
-    params: DocumentRangeFormattingParams,
-) -> list[TextEdit] | None:
-    return await format_document(
-        DocumentFormattingParams(
-            params.text_document, params.options, params.work_done_token
-        ),
-        params.range,
-    )
-
-
 async def _fix_document_impl(
     document: Document,
     settings: WorkspaceSettings,
@@ -1913,23 +1918,15 @@ async def _run_format_on_document(
     ]
 
     if format_range:
-        # Convert the start and end to character offsets instead of line:column
-        lines = document.source.splitlines(True)
         codec = PositionCodec(PositionEncodingKind.Utf16)
-        range = codec.range_from_client_units(lines, format_range)
-        start = format_range.start.character
-        end = format_range.end.character
-
-        for index, line in enumerate(lines):
-            if index < range.start.line:
-                start += len(line)
-
-            if index < range.end.line:
-                end += len(line)
-            else:
-                break
+        format_range = codec.range_from_client_units(document.source.splitlines(True), format_range)
 
-        argv.extend(["--range-start", str(start), "--range-end", str(end)])
+        argv.extend(
+            [
+                "--range",
+                f"{format_range.start.line + 1}:{format_range.start.character + 1}-{format_range.end.line + 1}:{format_range.end.character + 1}",
+            ]
+        )
 
     for arg in settings.get("format", {}).get("args", []):
         if arg in UNSUPPORTED_FORMAT_ARGS:
diff --git a/tests/test_format.py b/tests/test_format.py
index 89bf2e6..153e37e 100644
--- a/tests/test_format.py
+++ b/tests/test_format.py
@@ -91,13 +91,7 @@ async def test_format_range(tmp_path, ruff_version: Version):
 print ("Not formatted")
 """
 
-    expected = """x   = 1
-
-
-print("Formatted")
-
-print ("Not formatted")
-"""
+    expected = """print("Formatted")\n"""
 
     test_file = tmp_path.joinpath("main.py")
     test_file.write_text(original)
@@ -128,3 +122,6 @@ async def test_format_range(tmp_path, ruff_version: Version):
             original_source=document.source, fixed_source=result.stdout.decode("utf-8")
         )
         assert edit.new_text == expected
+        assert edit.range == Range(
+            start=Position(line=3, character=0), end=Position(line=5, character=0)
+        )