From 8d6d92aa5b5248b5ff70ebf7977f8af5cbcb10b9 Mon Sep 17 00:00:00 2001 From: Douglas Thor Date: Fri, 8 May 2020 06:30:10 -0700 Subject: [PATCH] Add option for printing a colored diff (#1266) --- CHANGES.md | 1 + black.py | 78 ++++++++++++++++++++++++++++++++++++++++++--- setup.py | 6 +++- tests/test_black.py | 41 ++++++++++++++++++++++++ 4 files changed, 120 insertions(+), 6 deletions(-) diff --git a/CHANGES.md b/CHANGES.md index b94d806..ebd1dd1 100644 --- a/CHANGES.md +++ b/CHANGES.md @@ -3,6 +3,7 @@ ### Unreleased - reindent docstrings when reindenting code around it (#1053) +- show colored diffs (#1266) ### 19.10b0 diff --git a/black.py b/black.py index 087683a..2a913fc 100644 --- a/black.py +++ b/black.py @@ -39,6 +39,7 @@ from typing import ( TypeVar, Union, cast, + TYPE_CHECKING, ) from typing_extensions import Final from mypy_extensions import mypyc_attr @@ -59,6 +60,9 @@ from blib2to3.pgen2.parse import ParseError from _black_version import version as __version__ +if TYPE_CHECKING: + import colorama # noqa: F401 + DEFAULT_LINE_LENGTH = 88 DEFAULT_EXCLUDES = r"/(\.eggs|\.git|\.hg|\.mypy_cache|\.nox|\.tox|\.venv|\.svn|_build|buck-out|build|dist)/" # noqa: B950 DEFAULT_INCLUDES = r"\.pyi?$" @@ -140,12 +144,18 @@ class WriteBack(Enum): YES = 1 DIFF = 2 CHECK = 3 + COLOR_DIFF = 4 @classmethod - def from_configuration(cls, *, check: bool, diff: bool) -> "WriteBack": + def from_configuration( + cls, *, check: bool, diff: bool, color: bool = False + ) -> "WriteBack": if check and not diff: return cls.CHECK + if diff and color: + return cls.COLOR_DIFF + return cls.DIFF if diff else cls.YES @@ -380,6 +390,11 @@ def target_version_option_callback( is_flag=True, help="Don't write the files back, just output a diff for each file on stdout.", ) +@click.option( + "--color/--no-color", + is_flag=True, + help="Show colored diff. Only applies when `--diff` is given.", +) @click.option( "--fast/--safe", is_flag=True, @@ -458,6 +473,7 @@ def main( target_version: List[TargetVersion], check: bool, diff: bool, + color: bool, fast: bool, pyi: bool, py36: bool, @@ -470,7 +486,7 @@ def main( config: Optional[str], ) -> None: """The uncompromising code formatter.""" - write_back = WriteBack.from_configuration(check=check, diff=diff) + write_back = WriteBack.from_configuration(check=check, diff=diff, color=color) if target_version: if py36: err("Cannot use both --target-version and --py36") @@ -718,12 +734,15 @@ def format_file_in_place( if write_back == WriteBack.YES: with open(src, "w", encoding=encoding, newline=newline) as f: f.write(dst_contents) - elif write_back == WriteBack.DIFF: + elif write_back in (WriteBack.DIFF, WriteBack.COLOR_DIFF): now = datetime.utcnow() src_name = f"{src}\t{then} +0000" dst_name = f"{src}\t{now} +0000" diff_contents = diff(src_contents, dst_contents, src_name, dst_name) + if write_back == write_back.COLOR_DIFF: + diff_contents = color_diff(diff_contents) + with lock or nullcontext(): f = io.TextIOWrapper( sys.stdout.buffer, @@ -731,12 +750,57 @@ def format_file_in_place( newline=newline, write_through=True, ) + f = wrap_stream_for_windows(f) f.write(diff_contents) f.detach() return True +def color_diff(contents: str) -> str: + """Inject the ANSI color codes to the diff.""" + lines = contents.split("\n") + for i, line in enumerate(lines): + if line.startswith("+++") or line.startswith("---"): + line = "\033[1;37m" + line + "\033[0m" # bold white, reset + if line.startswith("@@"): + line = "\033[36m" + line + "\033[0m" # cyan, reset + if line.startswith("+"): + line = "\033[32m" + line + "\033[0m" # green, reset + elif line.startswith("-"): + line = "\033[31m" + line + "\033[0m" # red, reset + lines[i] = line + return "\n".join(lines) + + +def wrap_stream_for_windows( + f: io.TextIOWrapper, +) -> Union[io.TextIOWrapper, "colorama.AnsiToWin32.AnsiToWin32"]: + """ + Wrap the stream in colorama's wrap_stream so colors are shown on Windows. + + If `colorama` is not found, then no change is made. If `colorama` does + exist, then it handles the logic to determine whether or not to change + things. + """ + try: + from colorama import initialise + + # We set `strip=False` so that we can don't have to modify + # test_express_diff_with_color. + f = initialise.wrap_stream( + f, convert=None, strip=False, autoreset=False, wrap=True + ) + + # wrap_stream returns a `colorama.AnsiToWin32.AnsiToWin32` object + # which does not have a `detach()` method. So we fake one. + f.detach = lambda *args, **kwargs: None # type: ignore + except ImportError: + pass + + return f + + def format_stdin_to_stdout( fast: bool, *, write_back: WriteBack = WriteBack.NO, mode: Mode ) -> bool: @@ -762,11 +826,15 @@ def format_stdin_to_stdout( ) if write_back == WriteBack.YES: f.write(dst) - elif write_back == WriteBack.DIFF: + elif write_back in (WriteBack.DIFF, WriteBack.COLOR_DIFF): now = datetime.utcnow() src_name = f"STDIN\t{then} +0000" dst_name = f"STDOUT\t{now} +0000" - f.write(diff(src, dst, src_name, dst_name)) + d = diff(src, dst, src_name, dst_name) + if write_back == WriteBack.COLOR_DIFF: + d = color_diff(d) + f = wrap_stream_for_windows(f) + f.write(d) f.detach() diff --git a/setup.py b/setup.py index 6192131..c9db684 100644 --- a/setup.py +++ b/setup.py @@ -77,7 +77,11 @@ setup( "typing_extensions>=3.7.4", "mypy_extensions>=0.4.3", ], - extras_require={"d": ["aiohttp>=3.3.2", "aiohttp-cors"]}, + extras_require={ + "d": ["aiohttp>=3.3.2", "aiohttp-cors"], + "colorama": ["colorama>=0.4.3"], + }, + test_suite="tests.test_black", classifiers=[ "Development Status :: 4 - Beta", "Environment :: Console", diff --git a/tests/test_black.py b/tests/test_black.py index fdd1994..11bba1f 100644 --- a/tests/test_black.py +++ b/tests/test_black.py @@ -264,6 +264,28 @@ class BlackTestCase(unittest.TestCase): actual = actual.rstrip() + "\n" # the diff output has a trailing space self.assertEqual(expected, actual) + def test_piping_diff_with_color(self) -> None: + source, _ = read_data("expression.py") + config = THIS_DIR / "data" / "empty_pyproject.toml" + args = [ + "-", + "--fast", + f"--line-length={black.DEFAULT_LINE_LENGTH}", + "--diff", + "--color", + f"--config={config}", + ] + result = BlackRunner().invoke( + black.main, args, input=BytesIO(source.encode("utf8")) + ) + actual = result.output + # Again, the contents are checked in a different test, so only look for colors. + self.assertIn("\033[1;37m", actual) + self.assertIn("\033[36m", actual) + self.assertIn("\033[32m", actual) + self.assertIn("\033[31m", actual) + self.assertIn("\033[0m", actual) + @patch("black.dump_to_file", dump_to_stderr) def test_function(self) -> None: source, expected = read_data("function") @@ -352,6 +374,25 @@ class BlackTestCase(unittest.TestCase): ) self.assertEqual(expected, actual, msg) + def test_expression_diff_with_color(self) -> None: + source, _ = read_data("expression.py") + expected, _ = read_data("expression.diff") + tmp_file = Path(black.dump_to_file(source)) + try: + result = BlackRunner().invoke( + black.main, ["--diff", "--color", str(tmp_file)] + ) + finally: + os.unlink(tmp_file) + actual = result.output + # We check the contents of the diff in `test_expression_diff`. All + # we need to check here is that color codes exist in the result. + self.assertIn("\033[1;37m", actual) + self.assertIn("\033[36m", actual) + self.assertIn("\033[32m", actual) + self.assertIn("\033[31m", actual) + self.assertIn("\033[0m", actual) + @patch("black.dump_to_file", dump_to_stderr) def test_fstring(self) -> None: source, expected = read_data("fstring") -- 2.39.2