OneScience-Group/ESM

1from pathlib import Path2import re3import sys4 5 6PROJECT_ROOT = Path(__file__).resolve().parents[1]7MODEL_ROOT = PROJECT_ROOT / "model"8SCRIPT_ROOT = PROJECT_ROOT / "scripts"9 10 11def _python_files(root: Path):12    return [path for path in root.rglob("*.py") if "__pycache__" not in path.parts]13 14 15def _check_no_onescience_models_imports():16    offenders = []17    for path in _python_files(MODEL_ROOT):18        text = path.read_text(encoding="utf-8")19        if "onescience.models." in text:20            offenders.append(str(path.relative_to(PROJECT_ROOT)))21    return offenders22 23 24def _check_no_legacy_imports():25    pattern = re.compile(26        r"^\s*(from\s+(esm|openfold)(\.|\s)|import\s+(esm|openfold)(\.|\s|$))",27        re.MULTILINE,28    )29    offenders = []30    for root in (MODEL_ROOT, SCRIPT_ROOT):31        for path in _python_files(root):32            text = path.read_text(encoding="utf-8")33            if pattern.search(text):34                offenders.append(str(path.relative_to(PROJECT_ROOT)))35    return offenders36 37 38def _check_script_bootstrap():39    offenders = []40    for path in _python_files(SCRIPT_ROOT):41        text = path.read_text(encoding="utf-8")42        if "model.esm" in text and '_PROJECT_ROOT / "model"' not in text:43            offenders.append(str(path.relative_to(PROJECT_ROOT)))44    return offenders45 46 47def main():48    checks = {49        "onescience.models imports": _check_no_onescience_models_imports(),50        "legacy esm/openfold imports": _check_no_legacy_imports(),51        "script project-root bootstrap": _check_script_bootstrap(),52    }53    failed = {name: offenders for name, offenders in checks.items() if offenders}54    if failed:55        for name, offenders in failed.items():56            print(f"{name}:")57            for offender in offenders:58                print(f"  {offender}")59        return 160 61    print("Import boundary checks passed")62    return 063 64 65if __name__ == "__main__":66    sys.exit(main())67 68