X-Git-Url: https://git.madduck.net/etc/vim.git/blobdiff_plain/3500e1cda5bef73ddc7eaf79be6c67c918738936..d4a85643a465f5fae2113d07d22d021d4af4795a:/tests/util.py

diff --git a/tests/util.py b/tests/util.py
index 84e98bb..d65c2e6 100644
--- a/tests/util.py
+++ b/tests/util.py
@@ -9,7 +9,10 @@ from typing import Any, Iterator, List, Optional, Tuple
 import black
 from black.debug import DebugVisitor
 from black.mode import TargetVersion
-from black.output import err, out
+from black.output import diff, err, out
+
+PYTHON_SUFFIX = ".py"
+ALLOWED_SUFFIXES = (PYTHON_SUFFIX, ".pyi", ".out", ".diff", ".ipynb")
 
 THIS_DIR = Path(__file__).parent
 DATA_DIR = THIS_DIR / "data"
@@ -47,6 +50,9 @@ def _assert_format_equal(expected: str, actual: str) -> None:
         except Exception as ve:
             err(str(ve))
 
+    if actual != expected:
+        out(diff(expected, actual, "expected", "actual"))
+
     assert actual == expected
 
 
@@ -87,12 +93,30 @@ class BlackBaseTestCase(unittest.TestCase):
         _assert_format_equal(expected, actual)
 
 
-def read_data(name: str, data: bool = True) -> Tuple[str, str]:
+def get_base_dir(data: bool) -> Path:
+    return DATA_DIR if data else PROJECT_ROOT
+
+
+def all_data_cases(subdir_name: str, data: bool = True) -> List[str]:
+    cases_dir = get_base_dir(data) / subdir_name
+    assert cases_dir.is_dir()
+    return [case_path.stem for case_path in cases_dir.iterdir()]
+
+
+def get_case_path(
+    subdir_name: str, name: str, data: bool = True, suffix: str = PYTHON_SUFFIX
+) -> Path:
+    """Get case path from name"""
+    case_path = get_base_dir(data) / subdir_name / name
+    if not name.endswith(ALLOWED_SUFFIXES):
+        case_path = case_path.with_suffix(suffix)
+    assert case_path.is_file(), f"{case_path} is not a file."
+    return case_path
+
+
+def read_data(subdir_name: str, name: str, data: bool = True) -> Tuple[str, str]:
     """read_data('test_name') -> 'input', 'output'"""
-    if not name.endswith((".py", ".pyi", ".out", ".diff")):
-        name += ".py"
-    base_dir = DATA_DIR if data else PROJECT_ROOT
-    return read_data_from_file(base_dir / name)
+    return read_data_from_file(get_case_path(subdir_name, name, data))
 
 
 def read_data_from_file(file_name: Path) -> Tuple[str, str]:
@@ -114,6 +138,18 @@ def read_data_from_file(file_name: Path) -> Tuple[str, str]:
     return "".join(_input).strip() + "\n", "".join(_output).strip() + "\n"
 
 
+def read_jupyter_notebook(subdir_name: str, name: str, data: bool = True) -> str:
+    return read_jupyter_notebook_from_file(
+        get_case_path(subdir_name, name, data, suffix=".ipynb")
+    )
+
+
+def read_jupyter_notebook_from_file(file_name: Path) -> str:
+    with open(file_name, mode="rb") as fd:
+        content_bytes = fd.read()
+    return content_bytes.decode()
+
+
 @contextmanager
 def change_directory(path: Path) -> Iterator[None]:
     """Context manager to temporarily chdir to a different directory."""