src/black/output.py

   1 """Nice output for Black.
   2
   3 The double calls are for patching purposes in tests.
   4 """
   5
   6 import json
   7 from typing import Any, Optional
   8 from mypy_extensions import mypyc_attr
   9 import tempfile
  10
  11 from click import echo, style
  12
  13
  14 @mypyc_attr(patchable=True)
  15 def _out(message: Optional[str] = None, nl: bool = True, **styles: Any) -> None:
  16     if message is not None:
  17         if "bold" not in styles:
  18             styles["bold"] = True
  19         message = style(message, **styles)
  20     echo(message, nl=nl, err=True)
  21
  22
  23 @mypyc_attr(patchable=True)
  24 def _err(message: Optional[str] = None, nl: bool = True, **styles: Any) -> None:
  25     if message is not None:
  26         if "fg" not in styles:
  27             styles["fg"] = "red"
  28         message = style(message, **styles)
  29     echo(message, nl=nl, err=True)
  30
  31
  32 @mypyc_attr(patchable=True)
  33 def out(message: Optional[str] = None, nl: bool = True, **styles: Any) -> None:
  34     _out(message, nl=nl, **styles)
  35
  36
  37 def err(message: Optional[str] = None, nl: bool = True, **styles: Any) -> None:
  38     _err(message, nl=nl, **styles)
  39
  40
  41 def ipynb_diff(a: str, b: str, a_name: str, b_name: str) -> str:
  42     """Return a unified diff string between each cell in notebooks `a` and `b`."""
  43     a_nb = json.loads(a)
  44     b_nb = json.loads(b)
  45     diff_lines = [
  46         diff(
  47             "".join(a_nb["cells"][cell_number]["source"]) + "\n",
  48             "".join(b_nb["cells"][cell_number]["source"]) + "\n",
  49             f"{a_name}:cell_{cell_number}",
  50             f"{b_name}:cell_{cell_number}",
  51         )
  52         for cell_number, cell in enumerate(a_nb["cells"])
  53         if cell["cell_type"] == "code"
  54     ]
  55     return "".join(diff_lines)
  56
  57
  58 def diff(a: str, b: str, a_name: str, b_name: str) -> str:
  59     """Return a unified diff string between strings `a` and `b`."""
  60     import difflib
  61
  62     a_lines = [line for line in a.splitlines(keepends=True)]
  63     b_lines = [line for line in b.splitlines(keepends=True)]
  64     diff_lines = []
  65     for line in difflib.unified_diff(
  66         a_lines, b_lines, fromfile=a_name, tofile=b_name, n=5
  67     ):
  68         # Work around https://bugs.python.org/issue2142
  69         # See:
  70         # https://www.gnu.org/software/diffutils/manual/html_node/Incomplete-Lines.html
  71         if line[-1] == "\n":
  72             diff_lines.append(line)
  73         else:
  74             diff_lines.append(line + "\n")
  75             diff_lines.append("\\ No newline at end of file\n")
  76     return "".join(diff_lines)
  77
  78
  79 def color_diff(contents: str) -> str:
  80     """Inject the ANSI color codes to the diff."""
  81     lines = contents.split("\n")
  82     for i, line in enumerate(lines):
  83         if line.startswith("+++") or line.startswith("---"):
  84             line = "\033[1;37m" + line + "\033[0m"  # bold white, reset
  85         elif line.startswith("@@"):
  86             line = "\033[36m" + line + "\033[0m"  # cyan, reset
  87         elif line.startswith("+"):
  88             line = "\033[32m" + line + "\033[0m"  # green, reset
  89         elif line.startswith("-"):
  90             line = "\033[31m" + line + "\033[0m"  # red, reset
  91         lines[i] = line
  92     return "\n".join(lines)
  93
  94
  95 @mypyc_attr(patchable=True)
  96 def dump_to_file(*output: str, ensure_final_newline: bool = True) -> str:
  97     """Dump `output` to a temporary file. Return path to the file."""
  98     with tempfile.NamedTemporaryFile(
  99         mode="w", prefix="blk_", suffix=".log", delete=False, encoding="utf8"
 100     ) as f:
 101         for lines in output:
 102             f.write(lines)
 103             if ensure_final_newline and lines and lines[-1] != "\n":
 104                 f.write("\n")
 105     return f.name