tests/util.py

   1 import os
   2 import sys
   3 import unittest
   4 from contextlib import contextmanager
   5 from functools import partial
   6 from pathlib import Path
   7 from typing import Any, Iterator, List, Optional, Tuple
   8
   9 import black
  10 from black.debug import DebugVisitor
  11 from black.mode import TargetVersion
  12 from black.output import err, out
  13
  14 THIS_DIR = Path(__file__).parent
  15 DATA_DIR = THIS_DIR / "data"
  16 PROJECT_ROOT = THIS_DIR.parent
  17 EMPTY_LINE = "# EMPTY LINE WITH WHITESPACE" + " (this comment will be removed)"
  18 DETERMINISTIC_HEADER = "[Deterministic header]"
  19
  20 PY36_VERSIONS = {
  21     TargetVersion.PY36,
  22     TargetVersion.PY37,
  23     TargetVersion.PY38,
  24     TargetVersion.PY39,
  25 }
  26
  27 DEFAULT_MODE = black.Mode()
  28 ff = partial(black.format_file_in_place, mode=DEFAULT_MODE, fast=True)
  29 fs = partial(black.format_str, mode=DEFAULT_MODE)
  30
  31
  32 def _assert_format_equal(expected: str, actual: str) -> None:
  33     if actual != expected and not os.environ.get("SKIP_AST_PRINT"):
  34         bdv: DebugVisitor[Any]
  35         out("Expected tree:", fg="green")
  36         try:
  37             exp_node = black.lib2to3_parse(expected)
  38             bdv = DebugVisitor()
  39             list(bdv.visit(exp_node))
  40         except Exception as ve:
  41             err(str(ve))
  42         out("Actual tree:", fg="red")
  43         try:
  44             exp_node = black.lib2to3_parse(actual)
  45             bdv = DebugVisitor()
  46             list(bdv.visit(exp_node))
  47         except Exception as ve:
  48             err(str(ve))
  49
  50     assert actual == expected
  51
  52
  53 def assert_format(
  54     source: str,
  55     expected: str,
  56     mode: black.Mode = DEFAULT_MODE,
  57     *,
  58     fast: bool = False,
  59     minimum_version: Optional[Tuple[int, int]] = None,
  60 ) -> None:
  61     """Convenience function to check that Black formats as expected.
  62
  63     You can pass @minimum_version if you're passing code with newer syntax to guard
  64     safety guards so they don't just crash with a SyntaxError. Please note this is
  65     separate from TargetVerson Mode configuration.
  66     """
  67     actual = black.format_str(source, mode=mode)
  68     _assert_format_equal(expected, actual)
  69     # It's not useful to run safety checks if we're expecting no changes anyway. The
  70     # assertion right above will raise if reality does actually make changes. This just
  71     # avoids wasted CPU cycles.
  72     if not fast and source != expected:
  73         # Unfortunately the AST equivalence check relies on the built-in ast module
  74         # being able to parse the code being formatted. This doesn't always work out
  75         # when checking modern code on older versions.
  76         if minimum_version is None or sys.version_info >= minimum_version:
  77             black.assert_equivalent(source, actual)
  78         black.assert_stable(source, actual, mode=mode)
  79
  80
  81 def dump_to_stderr(*output: str) -> str:
  82     return "\n" + "\n".join(output) + "\n"
  83
  84
  85 class BlackBaseTestCase(unittest.TestCase):
  86     def assertFormatEqual(self, expected: str, actual: str) -> None:
  87         _assert_format_equal(expected, actual)
  88
  89
  90 def read_data(name: str, data: bool = True) -> Tuple[str, str]:
  91     """read_data('test_name') -> 'input', 'output'"""
  92     if not name.endswith((".py", ".pyi", ".out", ".diff")):
  93         name += ".py"
  94     base_dir = DATA_DIR if data else PROJECT_ROOT
  95     return read_data_from_file(base_dir / name)
  96
  97
  98 def read_data_from_file(file_name: Path) -> Tuple[str, str]:
  99     with open(file_name, "r", encoding="utf8") as test:
 100         lines = test.readlines()
 101     _input: List[str] = []
 102     _output: List[str] = []
 103     result = _input
 104     for line in lines:
 105         line = line.replace(EMPTY_LINE, "")
 106         if line.rstrip() == "# output":
 107             result = _output
 108             continue
 109
 110         result.append(line)
 111     if _input and not _output:
 112         # If there's no output marker, treat the entire file as already pre-formatted.
 113         _output = _input[:]
 114     return "".join(_input).strip() + "\n", "".join(_output).strip() + "\n"
 115
 116
 117 @contextmanager
 118 def change_directory(path: Path) -> Iterator[None]:
 119     """Context manager to temporarily chdir to a different directory."""
 120     previous_dir = os.getcwd()
 121     try:
 122         os.chdir(path)
 123         yield
 124     finally:
 125         os.chdir(previous_dir)