| """The dependency policy, enforced. |
| |
| Every module under align/ and scripts/ is AST-scanned. Any import whose top-level |
| package is not (a) stdlib, (b) the four allowed runtime deps, or (c) the package |
| itself, fails this test. Lazy imports inside functions count too -- a lazy |
| `import transformers` is still a dependency, just a hidden one. |
| |
| Whitelist is per-FILE and per-PACKAGE, so an optional extra can exist in exactly |
| one place and nowhere else. |
| """ |
| from __future__ import annotations |
|
|
| import ast |
| import sys |
| from pathlib import Path |
|
|
| ROOT = Path(__file__).resolve().parents[1] |
| ALLOWED = {"torch", "triton", "numpy", "safetensors", "align"} |
| |
| FILE_WHITELIST: dict[str, set[str]] = { |
| "align/data/ingest_parquet.py": {"pyarrow"}, |
| "scripts/teachers/download_teacher.py": {"huggingface_hub"}, |
| } |
| STDLIB = set(sys.stdlib_module_names) |
|
|
|
|
| def _imports(path: Path): |
| tree = ast.parse(path.read_text(), filename=str(path)) |
| for node in ast.walk(tree): |
| if isinstance(node, ast.Import): |
| for a in node.names: |
| yield node.lineno, a.name.split(".")[0] |
| elif isinstance(node, ast.ImportFrom): |
| if node.level: |
| continue |
| if node.module: |
| yield node.lineno, node.module.split(".")[0] |
|
|
|
|
| def test_no_forbidden_imports(): |
| bad = [] |
| for sub in ("align", "scripts"): |
| for py in sorted((ROOT / sub).rglob("*.py")): |
| rel = str(py.relative_to(ROOT)) |
| extra = FILE_WHITELIST.get(rel, set()) |
| for lineno, top in _imports(py): |
| if top in STDLIB or top in ALLOWED or top in extra: |
| continue |
| bad.append(f"{rel}:{lineno}: {top}") |
| assert not bad, "forbidden third-party imports:\n" + "\n".join(bad) |
|
|
|
|
| def test_extras_are_not_reachable_from_core(): |
| """extras/ may import anything, but nothing under align/ may import extras/.""" |
| for py in (ROOT / "align").rglob("*.py"): |
| for _, top in _imports(py): |
| assert top != "extras", f"{py} imports extras/" |
|
|
