Reimplement black.format_str as a line generator

This is an experiment. Still concatenates the lines afterwards without making any
use of the extra information we could get.

Author: akaihola

src/darker/black_diff.py

@@ -43,13 +43,13 @@
 from black import (
     TargetVersion,
     find_pyproject_toml,
-    format_str,
     parse_pyproject_toml,
     re_compile_maybe_verbose,
 )
 from black.const import DEFAULT_EXCLUDES, DEFAULT_INCLUDES
 from black.files import gen_python_files
 from black.report import Report
+from darker.linewise_black import format_str_to_lines
 
 from darker.utils import TextDocument
 
@@ -181,7 +181,8 @@ def run_black(src_contents: TextDocument, black_config: BlackConfig) -> TextDocu
     # https://github.com/psf/black/pull/2484 lands in Black.
     contents_for_black = src_contents.string_with_newline("\n")
     if contents_for_black.strip():
-        dst_contents = format_str(contents_for_black, mode=Mode(**mode))
+        dst_lines = format_str_to_lines(contents_for_black, mode=Mode(**mode))
+        dst_contents = "".join(dst_lines)
     else:
         dst_contents = "\n" if "\n" in src_contents.string else ""
     return TextDocument.from_str(

src/darker/linewise_black.py

@@ -0,0 +1,61 @@
+"""Re-implementation of :func:`black.format_str` as a line generator"""
+
+from typing import Generator
+from black import get_future_imports, detect_target_versions, decode_bytes
+from black.lines import Line, EmptyLineTracker
+from black.linegen import transform_line, LineGenerator
+from black.comments import normalize_fmt_off
+from black.mode import Mode
+from black.mode import Feature, supports_feature
+from black.parsing import lib2to3_parse
+
+
+def format_str_to_lines(
+    src_contents: str, *, mode: Mode
+) -> Generator[str, None, None]:  # pylint: disable=too-many-locals
+    """Reformat a string and yield each line of new contents
+
+    This is a re-implementation of :func:`black.format_str` modified to be a generator
+    which yields each resulting line instead of concatenating them into a single string.
+
+    """
+    src_node = lib2to3_parse(src_contents.lstrip(), mode.target_versions)
+    future_imports = get_future_imports(src_node)
+    if mode.target_versions:
+        versions = mode.target_versions
+    else:
+        versions = detect_target_versions(src_node)
+    normalize_fmt_off(src_node)
+    lines = LineGenerator(
+        mode=mode,
+        remove_u_prefix="unicode_literals" in future_imports
+        or supports_feature(versions, Feature.UNICODE_LITERALS),
+    )
+    elt = EmptyLineTracker(is_pyi=mode.is_pyi)
+    empty_line = str(Line(mode=mode))
+    empty_line_len = len(empty_line)
+    after = 0
+    split_line_features = {
+        feature
+        for feature in {Feature.TRAILING_COMMA_IN_CALL, Feature.TRAILING_COMMA_IN_DEF}
+        if supports_feature(versions, feature)
+    }
+    num_chars = 0
+    for current_line in lines.visit(src_node):
+        for _ in range(after):
+            yield empty_line
+        num_chars += after * empty_line_len
+        before, after = elt.maybe_empty_lines(current_line)
+        for _ in range(before):
+            yield empty_line
+        num_chars += before * empty_line_len
+        for line in transform_line(
+            current_line, mode=mode, features=split_line_features
+        ):
+            line_str = str(line)
+            yield line_str
+            num_chars += len(line_str)
+    if not num_chars:
+        normalized_content, _, newline = decode_bytes(src_contents.encode("utf-8"))
+        if "\n" in normalized_content:
+            yield newline

src/darker/tests/test_black_diff.py

@@ -174,12 +174,12 @@ def test_run_black(encoding, newline):
 def test_run_black_always_uses_unix_newlines(newline):
     """Content is always passed to Black with Unix newlines"""
     src = TextDocument.from_str(f"print ( 'touché' ){newline}")
-    with patch.object(black_diff, "format_str") as format_str:
-        format_str.return_value = 'print("touché")\n'
+    with patch.object(black_diff, "format_str_to_lines") as format_str_to_lines:
+        format_str_to_lines.return_value = ['print("touché")\n']
 
         _ = run_black(src, BlackConfig())
 
-    format_str.assert_called_once_with("print ( 'touché' )\n", mode=ANY)
+    format_str_to_lines.assert_called_once_with("print ( 'touché' )\n", mode=ANY)
 
 
 def test_run_black_ignores_excludes():

src/darker/tests/test_command_line.py

@@ -470,8 +470,10 @@ def test_black_options_skip_string_normalization(git_repo, config, options, expe
     added_files["main.py"].write_bytes(b"bar")
     mode_class_mock = Mock(wraps=black_diff.Mode)
     # Speed up tests by mocking `format_str` to skip running Black
-    format_str = Mock(return_value="bar")
-    with patch.multiple(black_diff, Mode=mode_class_mock, format_str=format_str):
+    format_str_to_lines = Mock(return_value=["bar"])
+    with patch.multiple(
+        black_diff, Mode=mode_class_mock, format_str_to_lines=format_str_to_lines
+    ):
 
         main(options + [str(path) for path in added_files.values()])
 
@@ -496,8 +498,10 @@ def test_black_options_skip_magic_trailing_comma(git_repo, config, options, expe
     added_files["main.py"].write_bytes(b"a = [1, 2,]")
     mode_class_mock = Mock(wraps=black_diff.Mode)
     # Speed up tests by mocking `format_str` to skip running Black
-    format_str = Mock(return_value="a = [1, 2,]")
-    with patch.multiple(black_diff, Mode=mode_class_mock, format_str=format_str):
+    format_str_to_lines = Mock(return_value=["a = [1, 2,]"])
+    with patch.multiple(
+        black_diff, Mode=mode_class_mock, format_str_to_lines=format_str_to_lines
+    ):
 
         main(options + [str(path) for path in added_files.values()])