tests/util.py

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