Support range formatting when using Black 23.11.0+. (#380)

Co-authored-by: Karthik Nadig <kanadig@microsoft.com>

bundled/tool/lsp_server.py

@@ -12,7 +12,7 @@
 import sys
 import sysconfig
 import traceback
-from typing import Any, Dict, List, Optional, Sequence
+from typing import Any, Dict, List, Optional, Sequence, Tuple
 
 
 # **********************************************************
@@ -89,6 +89,12 @@ def update_environ_path() -> None:
 # Minimum version of black supported.
 MIN_VERSION = "22.3.0"
 
+# Minimum version of black that supports the `--line-ranges` CLI option.
+LINE_RANGES_MIN_VERSION = (23, 11, 0)
+
+# Versions of black found by workspace
+VERSION_LOOKUP: Dict[str, Tuple[int, int, int]] = {}
+
 # **********************************************************
 # Formatting features start here
 # **********************************************************
@@ -102,13 +108,52 @@ def formatting(params: lsp.DocumentFormattingParams) -> list[lsp.TextEdit] | Non
     return _formatting_helper(document)
 
 
-@LSP_SERVER.feature(lsp.TEXT_DOCUMENT_RANGE_FORMATTING)
-def range_formatting(params: lsp.DocumentFormattingParams) -> list[lsp.TextEdit] | None:
-    """LSP handler for textDocument/formatting request."""
+@LSP_SERVER.feature(
+    lsp.TEXT_DOCUMENT_RANGE_FORMATTING,
+    lsp.DocumentRangeFormattingOptions(ranges_support=True),
+)
+def range_formatting(
+    params: lsp.DocumentRangeFormattingParams,
+) -> list[lsp.TextEdit] | None:
+    """LSP handler for textDocument/rangeFormatting request."""
+    document = LSP_SERVER.workspace.get_text_document(params.text_document.uri)
+    settings = _get_settings_by_document(document)
+    version = VERSION_LOOKUP[settings["workspaceFS"]]
+
+    if version >= LINE_RANGES_MIN_VERSION:
+        return _formatting_helper(
+            document,
+            args=[
+                "--line-ranges",
+                f"{params.range.start.line + 1}-{params.range.end.line + 1}",
+            ],
+        )
+    else:
+        log_warning(
+            "Black version earlier than 23.11.0 does not support range formatting. Formatting entire document."
+        )
+        return _formatting_helper(document)
+
 
-    log_warning("Black does not support range formatting. Formatting entire document.")
+@LSP_SERVER.feature(lsp.TEXT_DOCUMENT_RANGES_FORMATTING)
+def ranges_formatting(
+    params: lsp.DocumentRangesFormattingParams,
+) -> list[lsp.TextEdit] | None:
+    """LSP handler for textDocument/rangesFormatting request."""
     document = LSP_SERVER.workspace.get_text_document(params.text_document.uri)
-    return _formatting_helper(document)
+    settings = _get_settings_by_document(document)
+    version = VERSION_LOOKUP[settings["workspaceFS"]]
+
+    if version >= LINE_RANGES_MIN_VERSION:
+        args = []
+        for r in params.ranges:
+            args += ["--line-ranges", f"{r.start.line + 1}-{r.end.line + 1}"]
+        return _formatting_helper(document, args=args)
+    else:
+        log_warning(
+            "Black version earlier than 23.11.0 does not support range formatting. Formatting entire document."
+        )
+        return _formatting_helper(document)
 
 
 def is_python(code: str, file_path: str) -> bool:
@@ -121,8 +166,11 @@ def is_python(code: str, file_path: str) -> bool:
     return True
 
 
-def _formatting_helper(document: workspace.Document) -> list[lsp.TextEdit] | None:
-    extra_args = _get_args_by_file_extension(document)
+def _formatting_helper(
+    document: workspace.Document, args: Sequence[str] = None
+) -> list[lsp.TextEdit] | None:
+    args = [] if args is None else args
+    extra_args = args + _get_args_by_file_extension(document)
     extra_args += ["--stdin-filename", _get_filename_for_black(document)]
     result = _run_tool_on_document(document, use_stdin=True, extra_args=extra_args)
     if result and result.stdout:
@@ -226,8 +274,6 @@ def initialize(params: lsp.InitializeParams) -> None:
     paths = "\r\n   ".join(sys.path)
     log_to_output(f"sys.path used to run Server:\r\n   {paths}")
 
-    _log_version_info()
-
 
 @LSP_SERVER.feature(lsp.EXIT)
 def on_exit(_params: Optional[Any] = None) -> None:
@@ -241,12 +287,13 @@ def on_shutdown(_params: Optional[Any] = None) -> None:
     jsonrpc.shutdown_json_rpc()
 
 
-def _log_version_info() -> None:
-    for value in WORKSPACE_SETTINGS.values():
+def _update_workspace_settings_with_version_info(
+    workspace_settings: dict[str, Any]
+) -> None:
+    for settings in workspace_settings.values():
         try:
             from packaging.version import parse as parse_version
 
-            settings = copy.deepcopy(value)
             result = _run_tool(["--version"], settings)
             code_workspace = settings["workspaceFS"]
             log_to_output(
@@ -269,6 +316,11 @@ def _log_version_info() -> None:
 
             version = parse_version(actual_version)
             min_version = parse_version(MIN_VERSION)
+            VERSION_LOOKUP[code_workspace] = (
+                version.major,
+                version.minor,
+                version.micro,
+            )
 
             if version < min_version:
                 log_error(
@@ -281,6 +333,7 @@ def _log_version_info() -> None:
                     f"SUPPORTED {TOOL_MODULE}>={min_version}\r\n"
                     f"FOUND {TOOL_MODULE}=={actual_version}\r\n"
                 )
+
         except:  # pylint: disable=bare-except
             log_to_output(
                 f"Error while detecting black version:\r\n{traceback.format_exc()}"
@@ -318,6 +371,8 @@ def _update_workspace_settings(settings):
             "workspaceFS": key,
         }
 
+    _update_workspace_settings_with_version_info(WORKSPACE_SETTINGS)
+
 
 def _get_settings_by_path(file_path: pathlib.Path):
     workspaces = {s["workspaceFS"] for s in WORKSPACE_SETTINGS.values()}

noxfile.py

@@ -152,7 +152,13 @@ def lint(session: nox.Session) -> None:
     # check formatting using black
     session.install("black")
     session.run("black", "--check", "./bundled/tool")
-    session.run("black", "--check", "./src/test/python_tests")
+    session.run(
+        "black",
+        "--check",
+        "./src/test/python_tests",
+        "--exclude",
+        "test_data",
+    )
     session.run("black", "--check", "noxfile.py")
 
     # check typescript code

src/test/python_tests/lsp_test_client/session.py

@@ -143,6 +143,20 @@ def text_document_formatting(self, formatting_params):
         fut = self._send_request("textDocument/formatting", params=formatting_params)
         return fut.result()
 
+    def text_document_range_formatting(self, range_formatting_params):
+        """Sends text document references request to LSP server."""
+        fut = self._send_request(
+            "textDocument/rangeFormatting", params=range_formatting_params
+        )
+        return fut.result()
+
+    def text_document_ranges_formatting(self, ranges_formatting_params):
+        """Sends text document references request to LSP server."""
+        fut = self._send_request(
+            "textDocument/rangesFormatting", params=ranges_formatting_params
+        )
+        return fut.result()
+
     def set_notification_callback(self, notification_name, callback):
         """Set custom LS notification handler."""
         self._notification_callbacks[notification_name] = callback

src/test/python_tests/test_data/sample4/sample.py

@@ -0,0 +1,2 @@
+x = [1, 2, 3, 4, 5]
+y = [1,2,3,4,5]

src/test/python_tests/test_data/sample4/sample.unformatted

@@ -0,0 +1,2 @@
+x = [1,2,3,4,5]
+y = [1,2,3,4,5]

src/test/python_tests/test_data/sample5/sample.py

@@ -0,0 +1,4 @@
+w = [1, 2, 3, 4, 5]
+x = [1,2,3,4,5]
+y = [1, 2, 3, 4, 5]
+z = [1,2,3,4,5]

src/test/python_tests/test_data/sample5/sample.unformatted

@@ -0,0 +1,4 @@
+w = [1,2,3,4,5]
+x = [1,2,3,4,5]
+y = [1,2,3,4,5]
+z = [1,2,3,4,5]

src/test/python_tests/test_formatting.py

@@ -121,3 +121,67 @@ def test_skipping_site_packages_files():
 
     expected = None
     assert_that(actual, is_(expected))
+
+
+@pytest.mark.parametrize(
+    "sample, ranges", [("sample4", "single-range"), ("sample5", "multi-range")]
+)
+def test_range_formatting(sample: str, ranges: str):
+    """Test formatting a python file."""
+    FORMATTED_TEST_FILE_PATH = constants.TEST_DATA / sample / "sample.py"
+    UNFORMATTED_TEST_FILE_PATH = constants.TEST_DATA / sample / "sample.unformatted"
+
+    contents = UNFORMATTED_TEST_FILE_PATH.read_text(encoding="utf-8")
+    lines = contents.splitlines()
+
+    actual = []
+    with utils.python_file(contents, UNFORMATTED_TEST_FILE_PATH.parent) as pf:
+        uri = utils.as_uri(str(pf))
+
+        with session.LspSession() as ls_session:
+            ls_session.initialize()
+            ls_session.notify_did_open(
+                {
+                    "textDocument": {
+                        "uri": uri,
+                        "languageId": "python",
+                        "version": 1,
+                        "text": contents,
+                    }
+                }
+            )
+
+            if ranges == "single-range":
+                actual = ls_session.text_document_range_formatting(
+                    {
+                        "textDocument": {"uri": uri},
+                        # `options` is not used by black
+                        "options": {"tabSize": 4, "insertSpaces": True},
+                        "range": {
+                            "start": {"line": 0, "character": 0},
+                            "end": {"line": 0, "character": len(lines[0])},
+                        },
+                    }
+                )
+            else:
+                actual = ls_session.text_document_ranges_formatting(
+                    {
+                        "textDocument": {"uri": uri},
+                        # `options` is not used by black
+                        "options": {"tabSize": 4, "insertSpaces": True},
+                        "ranges": [
+                            {
+                                "start": {"line": 0, "character": 0},
+                                "end": {"line": 0, "character": len(lines[0])},
+                            },
+                            {
+                                "start": {"line": 2, "character": 0},
+                                "end": {"line": 2, "character": len(lines[2])},
+                            },
+                        ],
+                    }
+                )
+
+    expected_text = FORMATTED_TEST_FILE_PATH.read_text(encoding="utf-8")
+    actual_text = utils.apply_text_edits(contents, utils.destructure_text_edits(actual))
+    assert_that(actual_text, is_(expected_text))