Repository: asmeurer/removestar

Issue:
Coloured output
The output right now is bland, but most of the pr-commit tools offer [rich](https://github.com/Textualize/rich) output for ease of reading.

Discussion:

Patch:
diff --git a/removestar/__main__.py b/removestar/__main__.py
index c5ae9dc..a724953 100755
--- a/removestar/__main__.py
+++ b/removestar/__main__.py
@@ -21,6 +21,7 @@ import sys
 
 from . import __version__
 from .helper import get_diff_text
+from .output import get_colored_diff, red
 from .removestar import fix_code
 
 
@@ -30,7 +31,7 @@ class RawDescriptionHelpArgumentDefaultsHelpFormatter(
     pass
 
 
-def main():
+def main():  # noqa: PLR0912
     parser = argparse.ArgumentParser(
         description=__doc__,
         prog="removestar",
@@ -97,7 +98,7 @@ def main():
             continue
 
         if not os.path.isfile(file):
-            print(f"Error: {file}: no such file or directory", file=sys.stderr)
+            print(red(f"Error: {file}: no such file or directory"), file=sys.stderr)
             continue
 
         with open(file, encoding="utf-8") as f:
@@ -114,7 +115,7 @@ def main():
             )
         except (RuntimeError, NotImplementedError) as e:
             if not args.quiet:
-                print(f"Error with {file}: {e}", file=sys.stderr)
+                print(red(f"Error with {file}: {e}"), file=sys.stderr)
             continue
 
         if new_code != code:
@@ -122,12 +123,24 @@ def main():
             if args.in_place:
                 with open(file, "w", encoding="utf-8") as f:
                     f.write(new_code)
+                if not args.quiet:
+                    print(
+                        get_colored_diff(
+                            get_diff_text(
+                                io.StringIO(code).readlines(),
+                                io.StringIO(new_code).readlines(),
+                                file,
+                            )
+                        )
+                    )
             else:
                 print(
-                    get_diff_text(
-                        io.StringIO(code).readlines(),
-                        io.StringIO(new_code).readlines(),
-                        file,
+                    get_colored_diff(
+                        get_diff_text(
+                            io.StringIO(code).readlines(),
+                            io.StringIO(new_code).readlines(),
+                            file,
+                        )
                     )
                 )
 
diff --git a/removestar/output.py b/removestar/output.py
new file mode 100644
index 0000000..973fade
--- /dev/null
+++ b/removestar/output.py
@@ -0,0 +1,54 @@
+def bold(line):
+    return "\033[1m" + line + "\033[0m"  # bold, reset
+
+
+def red(line):
+    return "\033[31m" + line + "\033[0m"  # red, reset
+
+
+def yellow(line):
+    return "\033[33m" + line + "\033[0m"  # yellow, reset
+
+
+def cyan(line):
+    return "\033[36m" + line + "\033[0m"  # cyan, reset
+
+
+def green(line):
+    return "\033[32m" + line + "\033[0m"  # green, reset
+
+
+def get_colored_diff(contents):
+    """Inject the ANSI color codes to the diff."""
+    # taken from https://github.com/psf/black/blob/main/src/black/output.py
+    # Copyright (c) 2018 Łukasz Langa
+
+    # Permission is hereby granted, free of charge, to any person obtaining a copy
+    # of this software and associated documentation files (the "Software"), to deal
+    # in the Software without restriction, including without limitation the rights
+    # to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+    # copies of the Software, and to permit persons to whom the Software is
+    # furnished to do so, subject to the following conditions:
+
+    # The above copyright notice and this permission notice shall be included in all
+    # copies or substantial portions of the Software.
+
+    # THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+    # IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+    # FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+    # AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+    # LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+    # OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+    # SOFTWARE.
+    lines = contents.split("\n")
+    for i, line in enumerate(lines):
+        if line.startswith(("+++", "---")):
+            line = bold(line)  # bold, reset  # noqa: PLW2901
+        elif line.startswith("@@"):
+            line = cyan(line)  # cyan, reset  # noqa: PLW2901
+        elif line.startswith("+"):
+            line = green(line)  # green, reset  # noqa: PLW2901
+        elif line.startswith("-"):
+            line = red(line)  # red, reset  # noqa: PLW2901
+        lines[i] = line
+    return "\n".join(lines)
diff --git a/removestar/removestar.py b/removestar/removestar.py
index 9180be7..3c6742e 100644
--- a/removestar/removestar.py
+++ b/removestar/removestar.py
@@ -9,6 +9,8 @@ from pathlib import Path
 from pyflakes.checker import _MAGIC_GLOBALS, Checker, ModuleScope
 from pyflakes.messages import ImportStarUsage, ImportStarUsed
 
+from .output import green, yellow
+
 # quit and exit are not included in old versions of pyflakes
 MAGIC_GLOBALS = set(_MAGIC_GLOBALS).union({"quit", "exit"})
 
@@ -65,13 +67,15 @@ def fix_code(
         if not mods:
             if not quiet:
                 print(
-                    f"Warning: {file}: could not find import for '{name}'",
+                    yellow(f"Warning: {file}: could not find import for '{name}'"),
                     file=sys.stderr,
                 )
             continue
         if len(mods) > 1 and not quiet:
             print(
-                f"Warning: {file}: '{name}' comes from multiple modules: {', '.join(map(repr, mods))}. Using '{mods[-1]}'.",
+                yellow(
+                    f"Warning: {file}: '{name}' comes from multiple modules: {', '.join(map(repr, mods))}. Using '{mods[-1]}'."
+                ),
                 file=sys.stderr,
             )
 
@@ -157,22 +161,28 @@ def replace_imports(
             if comment and is_noqa_comment_allowing_star_import(comment):
                 if verbose:
                     print(
-                        f"{verbose_prefix}Retaining 'from {mod} import *' due to noqa comment",
+                        green(
+                            f"{verbose_prefix}Retaining 'from {mod} import *' due to noqa comment"
+                        ),
                         file=sys.stderr,
                     )
                 return original_import
 
             if verbose:
                 print(
-                    f"{verbose_prefix}Replacing 'from {mod} import *' with '{new_import.strip()}'",
+                    green(
+                        f"{verbose_prefix}Replacing 'from {mod} import *' with '{new_import.strip()}'"
+                    ),
                     file=sys.stderr,
                 )
 
             if not new_import and comment:
                 if not quiet:
                     print(
-                        f"{warning_prefix}The removed star import statement for '{mod}' "
-                        f"had an inline comment which may not make sense without the import",
+                        yellow(
+                            f"{warning_prefix}The removed star import statement for '{mod}' "
+                            f"had an inline comment which may not make sense without the import"
+                        ),
                         file=sys.stderr,
                     )
                 return f"{comment}\n"
@@ -184,7 +194,7 @@ def replace_imports(
         new_code, subs_made = star_import.subn(star_import_replacement, code)
         if subs_made == 0 and not quiet:
             print(
-                f"{warning_prefix}Could not find the star imports for '{mod}'",
+                yellow(f"{warning_prefix}Could not find the star imports for '{mod}'"),
                 file=sys.stderr,
             )
         code = new_code


Test patch:
diff --git a/tests/test_removestar.py b/tests/test_removestar.py
index 9c24839..a944ad3 100644
--- a/tests/test_removestar.py
+++ b/tests/test_removestar.py
@@ -7,6 +7,7 @@ from pathlib import Path
 
 import pytest
 from pyflakes.checker import Checker
+from removestar.output import get_colored_diff, green, red, yellow
 from removestar.removestar import (
     ExternalModuleError,
     fix_code,
@@ -1317,8 +1318,8 @@ def test_replace_imports():
             {"code": code_mod4, "repls": {".mod1": ["a"], ".mod2": ["b", "c"]}},
             code_mod4_fixed,
             [
-                "Replacing 'from .mod1 import *' with 'from .mod1 import a'",
-                "Replacing 'from .mod2 import *' with 'from .mod2 import b, c'",
+                green("Replacing 'from .mod1 import *' with 'from .mod1 import a'"),
+                green("Replacing 'from .mod2 import *' with 'from .mod2 import b, c'"),
             ],
         ),
         (
@@ -1329,8 +1330,12 @@ def test_replace_imports():
             },
             code_mod4_fixed,
             [
-                "directory/mod4.py: Replacing 'from .mod1 import *' with 'from .mod1 import a'",
-                "directory/mod4.py: Replacing 'from .mod2 import *' with 'from .mod2 import b, c'",
+                green(
+                    "directory/mod4.py: Replacing 'from .mod1 import *' with 'from .mod1 import a'"
+                ),
+                green(
+                    "directory/mod4.py: Replacing 'from .mod2 import *' with 'from .mod2 import b, c'"
+                ),
             ],
         ),
         (
@@ -1340,9 +1345,9 @@ def test_replace_imports():
             },
             code_mod_commented_star_fixed,
             [
-                "Replacing 'from .mod3 import *' with 'from .mod3 import name'",
-                "Retaining 'from .mod1 import *' due to noqa comment",
-                "Retaining 'from .mod2 import *' due to noqa comment",
+                green("Replacing 'from .mod3 import *' with 'from .mod3 import name'"),
+                green("Retaining 'from .mod1 import *' due to noqa comment"),
+                green("Retaining 'from .mod2 import *' due to noqa comment"),
             ],
         ),
         (
@@ -1353,9 +1358,15 @@ def test_replace_imports():
             },
             code_mod_commented_star_fixed,
             [
-                "directory/mod_commented_star.py: Replacing 'from .mod3 import *' with 'from .mod3 import name'",
-                "directory/mod_commented_star.py: Retaining 'from .mod1 import *' due to noqa comment",
-                "directory/mod_commented_star.py: Retaining 'from .mod2 import *' due to noqa comment",
+                green(
+                    "directory/mod_commented_star.py: Replacing 'from .mod3 import *' with 'from .mod3 import name'"
+                ),
+                green(
+                    "directory/mod_commented_star.py: Retaining 'from .mod1 import *' due to noqa comment"
+                ),
+                green(
+                    "directory/mod_commented_star.py: Retaining 'from .mod2 import *' due to noqa comment"
+                ),
             ],
         ),
     ],
@@ -1388,8 +1399,12 @@ def test_replace_imports_warnings(capsys):
     )
     out, err = capsys.readouterr()
     assert set(err.splitlines()) == {
-        "Warning: module/mod_unfixable.py: Could not find the star imports for '.mod1'",
-        "Warning: module/mod_unfixable.py: Could not find the star imports for '.mod2'",
+        yellow(
+            "Warning: module/mod_unfixable.py: Could not find the star imports for '.mod1'"
+        ),
+        yellow(
+            "Warning: module/mod_unfixable.py: Could not find the star imports for '.mod2'"
+        ),
     }
 
     assert (
@@ -1400,8 +1415,8 @@ def test_replace_imports_warnings(capsys):
     )
     out, err = capsys.readouterr()
     assert set(err.splitlines()) == {
-        "Warning: Could not find the star imports for '.mod1'",
-        "Warning: Could not find the star imports for '.mod2'",
+        yellow("Warning: Could not find the star imports for '.mod1'"),
+        yellow("Warning: Could not find the star imports for '.mod2'"),
     }
 
     assert (
@@ -1423,7 +1438,9 @@ def test_replace_imports_warnings(capsys):
     )
     out, err = capsys.readouterr()
     assert set(err.splitlines()) == {
-        "Warning: module/mod_commented_unused_star.py: The removed star import statement for '.mod1' had an inline comment which may not make sense without the import",
+        yellow(
+            "Warning: module/mod_commented_unused_star.py: The removed star import statement for '.mod1' had an inline comment which may not make sense without the import"
+        ),
     }
 
     assert (
@@ -1434,7 +1451,9 @@ def test_replace_imports_warnings(capsys):
     )
     out, err = capsys.readouterr()
     assert set(err.splitlines()) == {
-        "Warning: The removed star import statement for '.mod1' had an inline comment which may not make sense without the import",
+        yellow(
+            "Warning: The removed star import statement for '.mod1' had an inline comment which may not make sense without the import"
+        ),
     }
 
     assert (
@@ -1656,9 +1675,12 @@ Warning: {directory}/mod_unfixable.py: Could not find the star imports for '.mod
 Warning: {directory}/mod_commented_unused_star.py: The removed star import statement for '.mod1' had an inline comment which may not make sense without the import
 """.splitlines()
     )
+    colored_warnings = {yellow(warning) for warning in warnings}
 
-    error = f"Error with {directory}/mod_bad.py: SyntaxError: invalid syntax (mod_bad.py, line 1)"
-    assert set(p.stderr.splitlines()) == warnings.union({error})
+    error = red(
+        f"Error with {directory}/mod_bad.py: SyntaxError: invalid syntax (mod_bad.py, line 1)"
+    )
+    assert set(p.stderr.splitlines()) == colored_warnings.union({error})
 
     diffs = [
         f"""\
@@ -1780,7 +1802,7 @@ Warning: {directory}/mod_commented_unused_star.py: The removed star import state
     ]
     unchanged = ["__init__.py", "mod_bad.py", "mod_unfixable.py"]
     for d in diffs:
-        assert d in p.stdout, p.stdout
+        assert get_colored_diff(d) in p.stdout, p.stdout
     for mod_path in unchanged:
         assert f"--- original/{directory}/{mod_path}" not in p.stdout
     cmp = dircmp(directory, directory_orig)
@@ -1793,7 +1815,7 @@ Warning: {directory}/mod_commented_unused_star.py: The removed star import state
     )
     assert p.stderr == ""
     for d in diffs:
-        assert d in p.stdout
+        assert get_colored_diff(d) in p.stdout
     cmp = dircmp(directory, directory_orig)
     assert _dirs_equal(cmp)
 
@@ -1826,10 +1848,13 @@ Warning: {directory}/mod_commented_unused_star.py: The removed star import state
 {directory}/submod_recursive/submod2.py: Replacing 'from . import *' with 'from . import a'
 """.splitlines()
     )
+    colored_changes = {green(change) for change in changes}
 
-    assert set(p.stderr.splitlines()) == changes.union({error}).union(warnings)
+    assert set(p.stderr.splitlines()) == colored_changes.union({error}).union(
+        colored_warnings
+    )
     for d in diffs:
-        assert d in p.stdout, p.stdout
+        assert get_colored_diff(d) in p.stdout, p.stdout
     cmp = dircmp(directory, directory_orig)
     assert _dirs_equal(cmp)
 
@@ -1844,12 +1869,15 @@ Error with {directory}/mod6.py: Static determination of external module imports
 Error with {directory}/mod7.py: Static determination of external module imports is not supported.
 """.splitlines()
     )
-    assert set(p.stderr.splitlines()) == {error}.union(static_error).union(warnings)
+    colored_static_error = {red(err) for err in static_error}
+    assert set(p.stderr.splitlines()) == {error}.union(colored_static_error).union(
+        colored_warnings
+    )
     for d in diffs:
         if "mod6" in d:
-            assert d not in p.stdout
+            assert get_colored_diff(d) not in p.stdout
         else:
-            assert d in p.stdout, p.stdout
+            assert get_colored_diff(d) in p.stdout, p.stdout
     cmp = dircmp(directory, directory_orig)
     assert _dirs_equal(cmp)
 
@@ -1869,9 +1897,9 @@ Error with {directory}/mod7.py: Static determination of external module imports
     assert p.stderr == ""
     for d in diffs:
         if "mod6" in d:
-            assert d not in p.stdout
+            assert get_colored_diff(d) not in p.stdout
         else:
-            assert d in p.stdout, p.stdout
+            assert get_colored_diff(d) in p.stdout, p.stdout
     cmp = dircmp(directory, directory_orig)
     assert _dirs_equal(cmp)
 
@@ -1934,6 +1962,7 @@ Error with {directory}/mod7.py: Static determination of external module imports
         encoding="utf-8",
     )
     assert (
-        p.stderr == f"Error: {directory}/notarealfile.py: no such file or directory\n"
+        p.stderr
+        == red(f"Error: {directory}/notarealfile.py: no such file or directory") + "\n"
     )
     assert p.stdout == ""