Merge pull request #1409 from PyCQA/issue/1405/add-coloring-to-diff-output

Adds colors to --diff output.

3 files changed, 85 insertions, 8 deletions

diff --git a/isort/api.py b/isort/api.py
index 059bbf9e..8b0ca237 100644
--- a/isort/api.py
+++ b/isort/api.py
@@ -131,6 +131,7 @@ def sort_stream(
             file_output=_output_stream.read(),
             file_path=file_path,
             output=output_stream if show_diff is True else cast(TextIO, show_diff),
+            color_output=config.color_output,
         )
         return changed
 
@@ -233,6 +234,7 @@ def check_stream(
                 file_output=output_stream.read(),
                 file_path=file_path,
                 output=None if show_diff is True else cast(TextIO, show_diff),
+                color_output=config.color_output,
             )
         return False
 
@@ -335,6 +337,7 @@ def sort_file(
                                     file_output=tmp_out.read(),
                                     file_path=file_path or source_file.path,
                                     output=None if show_diff is True else cast(TextIO, show_diff),
+                                    color_output=config.color_output,
                                 )
                                 if show_diff or (
                                     ask_to_apply
diff --git a/isort/format.py b/isort/format.py
index 3dbb1957..67c8c5b1 100644
--- a/isort/format.py
+++ b/isort/format.py
@@ -1,3 +1,4 @@
+import re
 import sys
 from datetime import datetime
 from difflib import unified_diff
@@ -13,6 +14,10 @@ else:
     colorama.init()
 
 
+ADDED_LINE_PATTERN = re.compile(r"\+[^+]")
+REMOVED_LINE_PATTERN = re.compile(r"-[^-]")
+
+
 def format_simplified(import_line: str) -> str:
     import_line = import_line.strip()
     if import_line.startswith("from "):
@@ -37,7 +42,12 @@ def format_natural(import_line: str) -> str:
 
 
 def show_unified_diff(
-    *, file_input: str, file_output: str, file_path: Optional[Path], output: Optional[TextIO] = None
+    *,
+    file_input: str,
+    file_output: str,
+    file_path: Optional[Path],
+    output: Optional[TextIO] = None,
+    color_output: bool = False,
 ):
     """Shows a unified_diff for the provided input and output against the provided file path.
 
@@ -45,8 +55,9 @@ def show_unified_diff(
     - **file_output**: A string that represents the contents of a file after changes.
     - **file_path**: A Path object that represents the file path of the file being changed.
     - **output**: A stream to output the diff to. If non is provided uses sys.stdout.
+    - **color_output**: Use color in output if True.
     """
-    output = sys.stdout if output is None else output
+    printer = create_terminal_printer(color_output, output)
     file_name = "" if file_path is None else str(file_path)
     file_mtime = str(
         datetime.now() if file_path is None else datetime.fromtimestamp(file_path.stat().st_mtime)
@@ -60,7 +71,7 @@ def show_unified_diff(
         tofiledate=str(datetime.now()),
     )
     for line in unified_diff_lines:
-        output.write(line)
+        printer.diff_line(line)
 
 
 def ask_whether_to_apply_changes_to_file(file_path: str) -> bool:
@@ -84,28 +95,49 @@ class BasicPrinter:
     ERROR = "ERROR"
     SUCCESS = "SUCCESS"
 
+    def __init__(self, output: Optional[TextIO] = None):
+        self.output = output or sys.stdout
+
     def success(self, message: str) -> None:
-        print(f"{self.SUCCESS}: {message}")
+        print(f"{self.SUCCESS}: {message}", file=self.output)
 
     def error(self, message: str) -> None:
         print(
             f"{self.ERROR}: {message}",
+            file=self.output,
             # TODO this should print to stderr, but don't want to make it backward incompatible now
             # file=sys.stderr
         )
 
+    def diff_line(self, line: str) -> None:
+        self.output.write(line)
+
 
 class ColoramaPrinter(BasicPrinter):
-    def __init__(self):
+    ADDED_LINE = colorama.Fore.GREEN
+    REMOVED_LINE = colorama.Fore.RED
+
+    def __init__(self, output: Optional[TextIO] = None):
+        self.output = output or sys.stdout
         self.ERROR = self.style_text("ERROR", colorama.Fore.RED)
         self.SUCCESS = self.style_text("SUCCESS", colorama.Fore.GREEN)
 
     @staticmethod
-    def style_text(text: str, style: str) -> str:
+    def style_text(text: str, style: Optional[str] = None) -> str:
+        if style is None:
+            return text
         return style + text + colorama.Style.RESET_ALL
 
+    def diff_line(self, line: str) -> None:
+        style = None
+        if re.match(ADDED_LINE_PATTERN, line):
+            style = self.ADDED_LINE
+        elif re.match(REMOVED_LINE_PATTERN, line):
+            style = self.REMOVED_LINE
+        self.output.write(self.style_text(line, style))
+
 
-def create_terminal_printer(color: bool):
+def create_terminal_printer(color: bool, output: Optional[TextIO] = None):
     if color and colorama_unavailable:
         no_colorama_message = (
             "\n"
@@ -118,4 +150,4 @@ def create_terminal_printer(color: bool):
         print(no_colorama_message, file=sys.stderr)
         sys.exit(1)
 
-    return ColoramaPrinter() if color else BasicPrinter()
+    return ColoramaPrinter(output) if color else BasicPrinter(output)
diff --git a/tests/unit/test_format.py b/tests/unit/test_format.py
index 3a8f7236..a2658367 100644
--- a/tests/unit/test_format.py
+++ b/tests/unit/test_format.py
@@ -1,3 +1,4 @@
+from io import StringIO
 from unittest.mock import MagicMock, patch
 
 import colorama
@@ -29,6 +30,15 @@ def test_basic_printer(capsys):
     assert out == "ERROR: Some error\n"
 
 
+def test_basic_printer_diff(capsys):
+    printer = isort.format.create_terminal_printer(color=False)
+    printer.diff_line("+ added line\n")
+    printer.diff_line("- removed line\n")
+
+    out, _ = capsys.readouterr()
+    assert out == "+ added line\n- removed line\n"
+
+
 def test_colored_printer_success(capsys):
     printer = isort.format.create_terminal_printer(color=True)
     printer.success("All good!")
@@ -47,6 +57,38 @@ def test_colored_printer_error(capsys):
     assert colorama.Fore.RED in out
 
 
+def test_colored_printer_diff(capsys):
+    printer = isort.format.create_terminal_printer(color=True)
+    printer.diff_line("+++ file1\n")
+    printer.diff_line("--- file2\n")
+    printer.diff_line("+ added line\n")
+    printer.diff_line("normal line\n")
+    printer.diff_line("- removed line\n")
+    printer.diff_line("normal line\n")
+
+    out, _ = capsys.readouterr()
+    # No color added to lines with multiple + and -'s
+    assert out.startswith("+++ file1\n--- file2\n")
+    # Added lines are green
+    assert colorama.Fore.GREEN + "+ added line" in out
+    # Removed lines are red
+    assert colorama.Fore.RED + "- removed line" in out
+    # Normal lines are resetted back
+    assert colorama.Style.RESET_ALL + "normal line" in out
+
+
+def test_colored_printer_diff_output(capsys):
+    output = StringIO()
+    printer = isort.format.create_terminal_printer(color=True, output=output)
+    printer.diff_line("a line\n")
+
+    out, _ = capsys.readouterr()
+    assert out == ""
+
+    output.seek(0)
+    assert output.read().startswith("a line\n")
+
+
 @patch("isort.format.colorama_unavailable", True)
 def test_colorama_not_available_handled_gracefully(capsys):
     with pytest.raises(SystemExit) as system_exit: