tests/util.py

   1 import unittest
   2 from contextlib import contextmanager
   3 from pathlib import Path
   4 from typing import List, Tuple, Iterator
   5
   6 THIS_DIR = Path(__file__).parent
   7 PROJECT_ROOT = THIS_DIR.parent
   8 EMPTY_LINE = "# EMPTY LINE WITH WHITESPACE" + " (this comment will be removed)"
   9 DETERMINISTIC_HEADER = "[Deterministic header]"
  10
  11
  12 @contextmanager
  13 def skip_if_exception(e: str) -> Iterator[None]:
  14     try:
  15         yield
  16     except Exception as exc:
  17         if exc.__class__.__name__ == e:
  18             unittest.skip(f"Encountered expected exception {exc}, skipping")
  19         else:
  20             raise
  21
  22
  23 def read_data(name: str, data: bool = True) -> Tuple[str, str]:
  24     """read_data('test_name') -> 'input', 'output'"""
  25     if not name.endswith((".py", ".pyi", ".out", ".diff")):
  26         name += ".py"
  27     _input: List[str] = []
  28     _output: List[str] = []
  29     base_dir = THIS_DIR / "data" if data else PROJECT_ROOT
  30     with open(base_dir / name, "r", encoding="utf8") as test:
  31         lines = test.readlines()
  32     result = _input
  33     for line in lines:
  34         line = line.replace(EMPTY_LINE, "")
  35         if line.rstrip() == "# output":
  36             result = _output
  37             continue
  38
  39         result.append(line)
  40     if _input and not _output:
  41         # If there's no output marker, treat the entire file as already pre-formatted.
  42         _output = _input[:]
  43     return "".join(_input).strip() + "\n", "".join(_output).strip() + "\n"