|
| 1 | +#!/usr/bin/env python3 |
| 2 | +""" |
| 3 | +CI gate enforcing PyHealth's PR contribution rules. |
| 4 | +
|
| 5 | +Rules enforced whenever a PR touches pyhealth/**/*.py: |
| 6 | +
|
| 7 | + 1. Docs/examples: the PR must also modify at least one file under |
| 8 | + docs/** and one file under examples/**. |
| 9 | + 2. Lint: lines added or modified in touched pyhealth/**/*.py files must |
| 10 | + be free of ruff violations. Pre-existing violations elsewhere in a |
| 11 | + touched file are not flagged. |
| 12 | + 3. Docstring examples: new or modified top-level public classes/ |
| 13 | + functions in pyhealth/**/*.py must include a '>>>' usage example in |
| 14 | + their docstring. |
| 15 | +
|
| 16 | +Usage: |
| 17 | + python tools/check_pr_rules.py --base <base_sha> --head <head_sha> |
| 18 | +""" |
| 19 | +import argparse |
| 20 | +import ast |
| 21 | +import json |
| 22 | +import subprocess |
| 23 | +import sys |
| 24 | +from pathlib import Path |
| 25 | + |
| 26 | +REPO_ROOT = Path(__file__).resolve().parents[1] |
| 27 | + |
| 28 | + |
| 29 | +def sh(*args): |
| 30 | + return subprocess.run( |
| 31 | + args, cwd=REPO_ROOT, capture_output=True, text=True, check=True |
| 32 | + ).stdout |
| 33 | + |
| 34 | + |
| 35 | +def changed_files(base, head): |
| 36 | + out = sh("git", "diff", "--name-only", "--diff-filter=ACMR", f"{base}..{head}") |
| 37 | + return [line.strip() for line in out.splitlines() if line.strip()] |
| 38 | + |
| 39 | + |
| 40 | +def added_lines(base, head, path): |
| 41 | + """Line numbers in `path` at `head` that were added or modified vs `base`.""" |
| 42 | + out = sh("git", "diff", "--unified=0", f"{base}..{head}", "--", path) |
| 43 | + lines = set() |
| 44 | + for line in out.splitlines(): |
| 45 | + if not line.startswith("@@"): |
| 46 | + continue |
| 47 | + plus = line.split("+")[1].split(" ")[0] |
| 48 | + if "," in plus: |
| 49 | + start, count = (int(x) for x in plus.split(",")) |
| 50 | + else: |
| 51 | + start, count = int(plus), 1 |
| 52 | + lines |= set(range(start, start + count)) |
| 53 | + return lines |
| 54 | + |
| 55 | + |
| 56 | +def check_docs_examples(files): |
| 57 | + if not any(f.startswith("pyhealth/") and f.endswith(".py") for f in files): |
| 58 | + return [] |
| 59 | + problems = [] |
| 60 | + if not any(f.startswith("docs/") for f in files): |
| 61 | + problems.append( |
| 62 | + "PR modifies pyhealth/ source files but no file under docs/ " |
| 63 | + "was updated." |
| 64 | + ) |
| 65 | + if not any(f.startswith("examples/") for f in files): |
| 66 | + problems.append( |
| 67 | + "PR modifies pyhealth/ source files but no file under " |
| 68 | + "examples/ was updated." |
| 69 | + ) |
| 70 | + return problems |
| 71 | + |
| 72 | + |
| 73 | +def check_lint(files, base, head): |
| 74 | + py_files = [ |
| 75 | + f |
| 76 | + for f in files |
| 77 | + if f.startswith("pyhealth/") and f.endswith(".py") and (REPO_ROOT / f).exists() |
| 78 | + ] |
| 79 | + if not py_files: |
| 80 | + return [] |
| 81 | + result = subprocess.run( |
| 82 | + ["ruff", "check", "--output-format=json", *py_files], |
| 83 | + cwd=REPO_ROOT, |
| 84 | + capture_output=True, |
| 85 | + text=True, |
| 86 | + ) |
| 87 | + if not result.stdout.strip(): |
| 88 | + return [] |
| 89 | + problems = [] |
| 90 | + for v in json.loads(result.stdout): |
| 91 | + path = Path(v["filename"]).resolve().relative_to(REPO_ROOT).as_posix() |
| 92 | + line = v["location"]["row"] |
| 93 | + if line in added_lines(base, head, path): |
| 94 | + problems.append(f"{path}:{line}: {v['code']} {v['message']}") |
| 95 | + return problems |
| 96 | + |
| 97 | + |
| 98 | +def check_docstring_examples(files, base, head): |
| 99 | + problems = [] |
| 100 | + for path in files: |
| 101 | + if not (path.startswith("pyhealth/") and path.endswith(".py")): |
| 102 | + continue |
| 103 | + full = REPO_ROOT / path |
| 104 | + if not full.exists(): |
| 105 | + continue |
| 106 | + added = added_lines(base, head, path) |
| 107 | + if not added: |
| 108 | + continue |
| 109 | + try: |
| 110 | + tree = ast.parse(full.read_text()) |
| 111 | + except SyntaxError: |
| 112 | + continue |
| 113 | + for node in ast.iter_child_nodes(tree): |
| 114 | + if not isinstance(node, (ast.ClassDef, ast.FunctionDef, ast.AsyncFunctionDef)): |
| 115 | + continue |
| 116 | + if node.name.startswith("_"): |
| 117 | + continue |
| 118 | + span = set(range(node.lineno, node.end_lineno + 1)) |
| 119 | + if not span & added: |
| 120 | + continue |
| 121 | + doc = ast.get_docstring(node) |
| 122 | + if not doc or ">>>" not in doc: |
| 123 | + kind = "class" if isinstance(node, ast.ClassDef) else "function" |
| 124 | + problems.append( |
| 125 | + f"{path}:{node.lineno}: public {kind} '{node.name}' is " |
| 126 | + "new/modified but its docstring has no '>>>' usage example." |
| 127 | + ) |
| 128 | + return problems |
| 129 | + |
| 130 | + |
| 131 | +def main(): |
| 132 | + parser = argparse.ArgumentParser(description=__doc__) |
| 133 | + parser.add_argument("--base", required=True, help="base commit SHA") |
| 134 | + parser.add_argument("--head", required=True, help="head commit SHA") |
| 135 | + args = parser.parse_args() |
| 136 | + |
| 137 | + files = changed_files(args.base, args.head) |
| 138 | + problems = ( |
| 139 | + check_docs_examples(files) |
| 140 | + + check_lint(files, args.base, args.head) |
| 141 | + + check_docstring_examples(files, args.base, args.head) |
| 142 | + ) |
| 143 | + |
| 144 | + if problems: |
| 145 | + print("PR contribution rules failed:\n") |
| 146 | + for p in problems: |
| 147 | + print(f" - {p}") |
| 148 | + print(f"\n{len(problems)} issue(s) found. See CONTRIBUTING.md for details.") |
| 149 | + sys.exit(1) |
| 150 | + |
| 151 | + print("All PR contribution rules passed.") |
| 152 | + |
| 153 | + |
| 154 | +if __name__ == "__main__": |
| 155 | + main() |
0 commit comments