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 diff, 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     if actual != expected:
  51         out(diff(expected, actual, "expected", "actual"))
  52
  53     assert actual == expected
  54
  55
  56 def assert_format(
  57     source: str,
  58     expected: str,
  59     mode: black.Mode = DEFAULT_MODE,
  60     *,
  61     fast: bool = False,
  62     minimum_version: Optional[Tuple[int, int]] = None,
  63 ) -> None:
  64     """Convenience function to check that Black formats as expected.
  65
  66     You can pass @minimum_version if you're passing code with newer syntax to guard
  67     safety guards so they don't just crash with a SyntaxError. Please note this is
  68     separate from TargetVerson Mode configuration.
  69     """
  70     actual = black.format_str(source, mode=mode)
  71     _assert_format_equal(expected, actual)
  72     # It's not useful to run safety checks if we're expecting no changes anyway. The
  73     # assertion right above will raise if reality does actually make changes. This just
  74     # avoids wasted CPU cycles.
  75     if not fast and source != expected:
  76         # Unfortunately the AST equivalence check relies on the built-in ast module
  77         # being able to parse the code being formatted. This doesn't always work out
  78         # when checking modern code on older versions.
  79         if minimum_version is None or sys.version_info >= minimum_version:
  80             black.assert_equivalent(source, actual)
  81         black.assert_stable(source, actual, mode=mode)
  82
  83
  84 def dump_to_stderr(*output: str) -> str:
  85     return "\n" + "\n".join(output) + "\n"
  86
  87
  88 class BlackBaseTestCase(unittest.TestCase):
  89     def assertFormatEqual(self, expected: str, actual: str) -> None:
  90         _assert_format_equal(expected, actual)
  91
  92
  93 def read_data(name: str, data: bool = True) -> Tuple[str, str]:
  94     """read_data('test_name') -> 'input', 'output'"""
  95     if not name.endswith((".py", ".pyi", ".out", ".diff")):
  96         name += ".py"
  97     base_dir = DATA_DIR if data else PROJECT_ROOT
  98     return read_data_from_file(base_dir / name)
  99
 100
 101 def read_data_from_file(file_name: Path) -> Tuple[str, str]:
 102     with open(file_name, "r", encoding="utf8") as test:
 103         lines = test.readlines()
 104     _input: List[str] = []
 105     _output: List[str] = []
 106     result = _input
 107     for line in lines:
 108         line = line.replace(EMPTY_LINE, "")
 109         if line.rstrip() == "# output":
 110             result = _output
 111             continue
 112
 113         result.append(line)
 114     if _input and not _output:
 115         # If there's no output marker, treat the entire file as already pre-formatted.
 116         _output = _input[:]
 117     return "".join(_input).strip() + "\n", "".join(_output).strip() + "\n"
 118
 119
 120 @contextmanager
 121 def change_directory(path: Path) -> Iterator[None]:
 122     """Context manager to temporarily chdir to a different directory."""
 123     previous_dir = os.getcwd()
 124     try:
 125         os.chdir(path)
 126         yield
 127     finally:
 128         os.chdir(previous_dir)