File size: 2,216 Bytes
201dc77 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 | """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 (relative to ROOT) -> extra packages it alone may import lazily
FILE_WHITELIST: dict[str, set[str]] = {
"align/data/ingest_parquet.py": {"pyarrow"},
"scripts/teachers/download_teacher.py": {"huggingface_hub"}, # offline weight download, never imported at runtime
}
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: # relative import -> the package itself
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/"
|