Adds syntax highlighting to python compilation errors.

Summary:
Adds syntax highlighting for error messages that occur during the compilation step in the Buck2 preludes.

It also drops all parts of the stack trace that are not useful. If you still want the full stack trace you can use the `--debug` flag (intended primarily for PLF or folks working on the `compile.py` file itself).

If you hate having colors, the env variable `NO_COLOR` (https://no-color.org/) will disable all colors.
If you hate the colors I've chosen, please let me know a better color scheme.

Reviewed By: zsol

Differential Revision: D67802134

fbshipit-source-id: d1ad506ee02383252d894053cf8eceec1482a0b3

prelude/python/tools/compile.py

@@ -8,21 +8,28 @@
 # pyre-strict
 
 """
-Example usage:
+Example usage (internal):
 $ cat inputs.manifest
 [["foo.py", "input/foo.py", "//my_rule:foo"]]
-$ compile.py --output=out-dir --bytecode-manifest=output.manifest --ignore-errors inputs.manifest
+$ buck build //fbcode//python/build/compile:compile --show-full-output
+$ python <fulloutput/__main__.py> --output=out-dir --bytecode-manifest=output.manifest inputs.manifest
 $ find out-dir -type f
 out-dir/foo.pyc
+
+Or (external):
+compile.py --output=out-dir --bytecode-manifest=output.manifest --ignore-errors inputs.manifest
 """
 
 import argparse
 import errno
 import json
 import os
+import re
 import sys
-from py_compile import compile, PycInvalidationMode
-from typing import List
+import traceback
+from py_compile import compile, PycInvalidationMode, PyCompileError
+from types import TracebackType
+from typing import List, Type
 
 
 if sys.version_info[0] == 3:
@@ -60,6 +67,57 @@ def _mkdirs(dirpath: str) -> None:
             raise
 
 
+def _stderr_print(msg: str) -> None:
+    print(msg, file=sys.stderr, end="")
+
+
+def pretty_exception(
+    typ: Type[BaseException], exc: BaseException, tb: TracebackType
+) -> None:
+    try:
+        from colorama import Fore, just_fix_windows_console, Style
+
+        just_fix_windows_console()
+
+        trace = traceback.format_exception(typ, exc, tb)
+        prev_line = ""
+        for line in trace:
+            if line.startswith(
+                "\nDuring handling of the above exception, another exception occurred:"
+            ):
+                # Drop everything beyond the original exception
+                return
+
+            if line.startswith("  File"):
+                prev_line = line
+            else:
+                if prev_line.startswith("  File"):
+                    # Magenta for last listed filename and line number
+                    s = re.sub(
+                        r'"(.*?)"',
+                        lambda match: Fore.MAGENTA + f'"{match.group(1)}"' + Fore.RESET,
+                        prev_line,
+                    )
+                    s = re.sub(
+                        r"\b\d+$",
+                        lambda match: Fore.MAGENTA + f"{match.group()}" + Fore.RESET,
+                        s,
+                    )
+                    _stderr_print(s)
+                    prev_line = ""
+            if line.startswith("    "):
+                line = Fore.CYAN + line + Fore.RESET  # Cyan for code lines
+                _stderr_print(line)
+            elif line.startswith("SyntaxError") or line.startswith("ValueError"):
+                # Bold magenta for exception
+                line = Style.BRIGHT + Fore.MAGENTA + line + Style.RESET_ALL
+                _stderr_print(line)
+    except Exception:
+        # If anything goes wrong, fall back to the default exception printing
+        sys.excepthook = sys.__excepthook__
+        traceback.print_exception(typ, exc, tb)
+
+
 def main(argv: List[str]) -> None:
     parser = argparse.ArgumentParser(fromfile_prefix_chars="@")
     parser.add_argument("-o", "--output", required=True)
@@ -73,6 +131,11 @@ def main(argv: List[str]) -> None:
         default=PycInvalidationMode.UNCHECKED_HASH.name,
         choices=[m.name for m in PycInvalidationMode],
     )
+    parser.add_argument(
+        "--debug",
+        action="store_true",
+        default=False,
+    )
     parser.add_argument("manifests", nargs="*")
     args = parser.parse_args(argv[1:])
     invalidation_mode = PycInvalidationMode.__members__[args.invalidation_mode]
@@ -95,13 +158,18 @@ def main(argv: List[str]) -> None:
             dest_pyc = get_pyc_path(module, args.format)
             pyc = os.path.join(args.output, dest_pyc)
             _mkdirs(os.path.dirname(pyc))
-            compile(
-                src,
-                cfile=pyc,
-                dfile=get_py_path(module),
-                doraise=True,
-                invalidation_mode=invalidation_mode,
-            )
+            try:
+                compile(
+                    src,
+                    cfile=pyc,
+                    dfile=get_py_path(module),
+                    doraise=True,
+                    invalidation_mode=invalidation_mode,
+                )
+            except PyCompileError:
+                if not args.debug:
+                    sys.excepthook = pretty_exception
+                raise
             bytecode_manifest.append((dest_pyc, pyc, src))
     json.dump(bytecode_manifest, args.bytecode_manifest, indent=2)