"""Nice output for Black. The double calls are for patching purposes in tests. """ import json import re import tempfile from typing import Any, List, Optional from click import echo, style from mypy_extensions import mypyc_attr @mypyc_attr(patchable=True) def _out(message: Optional[str] = None, nl: bool = True, **styles: Any) -> None: if message is not None: if "bold" not in styles: styles["bold"] = True message = style(message, **styles) echo(message, nl=nl, err=True) @mypyc_attr(patchable=True) def _err(message: Optional[str] = None, nl: bool = True, **styles: Any) -> None: if message is not None: if "fg" not in styles: styles["fg"] = "red" message = style(message, **styles) echo(message, nl=nl, err=True) @mypyc_attr(patchable=True) def out(message: Optional[str] = None, nl: bool = True, **styles: Any) -> None: _out(message, nl=nl, **styles) def err(message: Optional[str] = None, nl: bool = True, **styles: Any) -> None: _err(message, nl=nl, **styles) def ipynb_diff(a: str, b: str, a_name: str, b_name: str) -> str: """Return a unified diff string between each cell in notebooks `a` and `b`.""" a_nb = json.loads(a) b_nb = json.loads(b) diff_lines = [ diff( "".join(a_nb["cells"][cell_number]["source"]) + "\n", "".join(b_nb["cells"][cell_number]["source"]) + "\n", f"{a_name}:cell_{cell_number}", f"{b_name}:cell_{cell_number}", ) for cell_number, cell in enumerate(a_nb["cells"]) if cell["cell_type"] == "code" ] return "".join(diff_lines) _line_pattern = re.compile(r"(.*?(?:\r\n|\n|\r|$))") def _splitlines_no_ff(source: str) -> List[str]: """Split a string into lines ignoring form feed and other chars. This mimics how the Python parser splits source code. A simplified version of the function with the same name in Lib/ast.py """ result = [match[0] for match in _line_pattern.finditer(source)] if result[-1] == "": result.pop(-1) return result def diff(a: str, b: str, a_name: str, b_name: str) -> str: """Return a unified diff string between strings `a` and `b`.""" import difflib a_lines = _splitlines_no_ff(a) b_lines = _splitlines_no_ff(b) diff_lines = [] for line in difflib.unified_diff( a_lines, b_lines, fromfile=a_name, tofile=b_name, n=5 ): # Work around https://bugs.python.org/issue2142 # See: # https://www.gnu.org/software/diffutils/manual/html_node/Incomplete-Lines.html if line[-1] == "\n": diff_lines.append(line) else: diff_lines.append(line + "\n") diff_lines.append("\\ No newline at end of file\n") return "".join(diff_lines) def color_diff(contents: str) -> str: """Inject the ANSI color codes to the diff.""" lines = contents.split("\n") for i, line in enumerate(lines): if line.startswith("+++") or line.startswith("---"): line = "\033[1m" + line + "\033[0m" # bold, reset elif line.startswith("@@"): line = "\033[36m" + line + "\033[0m" # cyan, reset elif line.startswith("+"): line = "\033[32m" + line + "\033[0m" # green, reset elif line.startswith("-"): line = "\033[31m" + line + "\033[0m" # red, reset lines[i] = line return "\n".join(lines) @mypyc_attr(patchable=True) def dump_to_file(*output: str, ensure_final_newline: bool = True) -> str: """Dump `output` to a temporary file. Return path to the file.""" with tempfile.NamedTemporaryFile( mode="w", prefix="blk_", suffix=".log", delete=False, encoding="utf8" ) as f: for lines in output: f.write(lines) if ensure_final_newline and lines and lines[-1] != "\n": f.write("\n") return f.name