Add flag for printing reports in colour

CHANGELOG.md

@@ -3,6 +3,7 @@
 ## Unreleased
 
 - Drop support for Python 3.7
+- Add `parse --color` flag for printing out change reports in color. (Aliases `-c`, `--colour`.)
 
 ## v1.0.4 [2023-05-09]
 

mypy_json_report.py

@@ -31,6 +31,7 @@
     Iterator,
     List,
     Optional,
+    Protocol,
     Union,
     cast,
 )
@@ -87,6 +88,13 @@ def main() -> None:
             """
         ),
     )
+    parse_parser.add_argument(
+        "-c",
+        "--color",
+        "--colour",
+        action="store_true",
+        help="Whether to colorize the diff-report output. Defaults to False.",
+    )
 
     parse_parser.set_defaults(func=_parse_command)
 
@@ -116,7 +124,12 @@ def _parse_command(args: argparse.Namespace) -> None:
     tracker = None
     if args.diff_old_report is not None:
         old_report = cast(ErrorSummary, _load_json_file(args.diff_old_report))
-        tracker = ChangeTracker(old_report)
+        change_report_writer: _ChangeReportWriter
+        if args.color:
+            change_report_writer = ColorChangeReportWriter()
+        else:
+            change_report_writer = DefaultChangeReportWriter()
+        tracker = ChangeTracker(old_report, report_writer=change_report_writer)
         processors.append(tracker)
 
     messages = MypyMessage.from_lines(sys.stdin)
@@ -229,6 +242,102 @@ class DiffReport:
     num_fixed_errors: int
 
 
+class _ChangeReportWriter(Protocol):
+    def write_report(self, diff: DiffReport) -> None:
+        ...
+
+
+class DefaultChangeReportWriter:
+    """Writes an error summary without color."""
+
+    def __init__(self, _write: Callable[[str], Any] = sys.stdout.write) -> None:
+        self.write = _write
+
+    def write_report(self, diff: DiffReport) -> None:
+        new_errors = "\n".join(diff.error_lines)
+        if new_errors:
+            self.write(new_errors + "\n\n")
+        self.write(f"Fixed errors: {diff.num_fixed_errors}\n")
+        self.write(f"New errors: {diff.num_new_errors}\n")
+        self.write(f"Total errors: {diff.total_errors}\n")
+
+
+class ColorChangeReportWriter:
+    """
+    Writes an error summary in color.
+
+    Inspired by the FancyFormatter in mypy.util.
+
+    Ref: https://github.com/python/mypy/blob/f9e8e0bda5cfbb54d6a8f9e482aa25da28a1a635/mypy/util.py#L761
+    """
+
+    _RESET = "\033[0m"
+    _BOLD = "\033[1m"
+    _BOLD_RED = "\033[31;1m"
+    _BOLD_YELLOW = "\033[33;1m"
+    _GREEN = "\033[32m"
+    _YELLOW = "\033[33m"
+    _BLUE = "\033[34m"
+
+    def __init__(self, _write: Callable[[str], Any] = sys.stdout.write) -> None:
+        self.write = _write
+
+    def write_report(self, diff: DiffReport) -> None:
+        new_errors = "\n".join([self._format_line(line) for line in diff.error_lines])
+        if new_errors:
+            self.write(new_errors + "\n\n")
+
+        fixed_color = self._BOLD_YELLOW if diff.num_fixed_errors else self._GREEN
+        error_color = self._BOLD_RED if diff.num_new_errors else self._GREEN
+
+        self.write(self._style(fixed_color, f"Fixed errors: {diff.num_fixed_errors}\n"))
+        self.write(self._style(error_color, f"New errors: {diff.num_new_errors}\n"))
+        self.write(self._style(self._BOLD, f"Total errors: {diff.total_errors}\n"))
+
+    def _style(self, style: str, message: str) -> str:
+        return f"{style}{message}{self._RESET}"
+
+    def _highlight_quotes(self, msg: str) -> str:
+        if msg.count('"') % 2:
+            return msg
+        parts = msg.split('"')
+        out = ""
+        for i, part in enumerate(parts):
+            if i % 2 == 0:
+                out += part
+            else:
+                out += self._style(self._BOLD, f'"{part}"')
+        return out
+
+    def _format_line(self, line: str) -> str:
+        if ": error: " in line:
+            # Separate the location from the message.
+            location, _, message = line.partition(" error: ")
+
+            # Extract the error code from the end of the message if it's there.
+            if message.endswith("]") and "  [" in message:
+                error_msg, _, code = message.rpartition("  [")
+                code = self._style(self._YELLOW, f"  [{code}")
+            else:
+                error_msg = message
+                code = ""
+
+            return (
+                location
+                + self._style(self._BOLD_RED, " error: ")
+                + self._highlight_quotes(error_msg)
+                + code
+            )
+        if ": note: " in line:
+            location, _, message = line.partition(" note: ")
+            return (
+                location
+                + self._style(self._BLUE, " note: ")
+                + self._highlight_quotes(message)
+            )
+        return line
+
+
 class ChangeTracker:
     """
     Compares the current Mypy report against a previous summary.
@@ -239,8 +348,11 @@ class ChangeTracker:
     that are cached in memory.
     """
 
-    def __init__(self, summary: ErrorSummary) -> None:
+    def __init__(
+        self, summary: ErrorSummary, report_writer: _ChangeReportWriter
+    ) -> None:
         self.old_report = summary
+        self.report_writer = report_writer
         self.error_lines: List[str] = []
         self.num_errors = 0
         self.num_new_errors = 0
@@ -282,13 +394,8 @@ def diff_report(self) -> DiffReport:
         )
 
     def write_report(self) -> Optional[ErrorCodes]:
-        write = sys.stderr.write
         diff = self.diff_report()
-        new_errors = "\n".join(diff.error_lines)
-        write(new_errors + "\n")
-        write(f"Fixed errors: {diff.num_fixed_errors}\n")
-        write(f"New errors: {diff.num_new_errors}\n")
-        write(f"Total errors: {diff.total_errors}\n")
+        self.report_writer.write_report(diff)
 
         if diff.num_new_errors or diff.num_fixed_errors:
             return ErrorCodes.ERROR_DIFF