tests/util.py

   1 import os
   2 import unittest
   3 from pathlib import Path
   4 from typing import Iterator, List, Tuple, Any
   5 from contextlib import contextmanager
   6 from functools import partial
   7
   8 import black
   9 from black.output import out, err
  10 from black.debug import DebugVisitor
  11
  12 THIS_DIR = Path(__file__).parent
  13 PROJECT_ROOT = THIS_DIR.parent
  14 EMPTY_LINE = "# EMPTY LINE WITH WHITESPACE" + " (this comment will be removed)"
  15 DETERMINISTIC_HEADER = "[Deterministic header]"
  16
  17
  18 DEFAULT_MODE = black.Mode()
  19 ff = partial(black.format_file_in_place, mode=DEFAULT_MODE, fast=True)
  20 fs = partial(black.format_str, mode=DEFAULT_MODE)
  21
  22
  23 def dump_to_stderr(*output: str) -> str:
  24     return "\n" + "\n".join(output) + "\n"
  25
  26
  27 class BlackBaseTestCase(unittest.TestCase):
  28     maxDiff = None
  29     _diffThreshold = 2 ** 20
  30
  31     def assertFormatEqual(self, expected: str, actual: str) -> None:
  32         if actual != expected and not os.environ.get("SKIP_AST_PRINT"):
  33             bdv: DebugVisitor[Any]
  34             out("Expected tree:", fg="green")
  35             try:
  36                 exp_node = black.lib2to3_parse(expected)
  37                 bdv = DebugVisitor()
  38                 list(bdv.visit(exp_node))
  39             except Exception as ve:
  40                 err(str(ve))
  41             out("Actual tree:", fg="red")
  42             try:
  43                 exp_node = black.lib2to3_parse(actual)
  44                 bdv = DebugVisitor()
  45                 list(bdv.visit(exp_node))
  46             except Exception as ve:
  47                 err(str(ve))
  48         self.assertMultiLineEqual(expected, actual)
  49
  50
  51 def read_data(name: str, data: bool = True) -> Tuple[str, str]:
  52     """read_data('test_name') -> 'input', 'output'"""
  53     if not name.endswith((".py", ".pyi", ".out", ".diff")):
  54         name += ".py"
  55     base_dir = THIS_DIR / "data" if data else PROJECT_ROOT
  56     return read_data_from_file(base_dir / name)
  57
  58
  59 def read_data_from_file(file_name: Path) -> Tuple[str, str]:
  60     with open(file_name, "r", encoding="utf8") as test:
  61         lines = test.readlines()
  62     _input: List[str] = []
  63     _output: List[str] = []
  64     result = _input
  65     for line in lines:
  66         line = line.replace(EMPTY_LINE, "")
  67         if line.rstrip() == "# output":
  68             result = _output
  69             continue
  70
  71         result.append(line)
  72     if _input and not _output:
  73         # If there's no output marker, treat the entire file as already pre-formatted.
  74         _output = _input[:]
  75     return "".join(_input).strip() + "\n", "".join(_output).strip() + "\n"
  76
  77
  78 @contextmanager
  79 def change_directory(path: Path) -> Iterator[None]:
  80     """Context manager to temporarily chdir to a different directory."""
  81     previous_dir = os.getcwd()
  82     try:
  83         os.chdir(path)
  84         yield
  85     finally:
  86         os.chdir(previous_dir)