[CI] Prune redundant CPU test overhead (#34309)

Co-authored-by: hnyls2002 <lsyincs@gmail.com>
2026-08-13 19:51:17 -07:00
parent 6ad3f2d8fd
commit 85cdf1178d
78 changed files with 884 additions and 981 deletions
@@ -0,0 +1,154 @@
+#!/usr/bin/env python3
+
+import ast
+import pathlib
+import re
+import sys
+
+_PYTEST_MAIN = re.compile(r"pytest\s*\.\s*main")
+
+
+def is_main_guard(node: ast.expr) -> bool:
+    if not isinstance(node, ast.Compare) or len(node.ops) != 1:
+        return False
+    if not isinstance(node.ops[0], ast.Eq):
+        return False
+    sides = [node.left, *node.comparators]
+    has_name = any(
+        isinstance(side, ast.Name) and side.id == "__name__" for side in sides
+    )
+    has_main = any(
+        isinstance(side, ast.Constant) and side.value == "__main__" for side in sides
+    )
+    return has_name and has_main
+
+
+def is_pytest_main_call(node: ast.AST) -> bool:
+    if not isinstance(node, ast.Call):
+        return False
+    func = node.func
+    return (
+        isinstance(func, ast.Attribute)
+        and func.attr == "main"
+        and isinstance(func.value, ast.Name)
+        and func.value.id == "pytest"
+    )
+
+
+def is_exit_call(node: ast.AST, parents: dict[int, ast.AST]) -> bool:
+    """``sys.exit(...)``, or a ``SystemExit(...)`` that is actually raised."""
+    if not isinstance(node, ast.Call):
+        return False
+    func = node.func
+    if (
+        isinstance(func, ast.Attribute)
+        and func.attr == "exit"
+        and isinstance(func.value, ast.Name)
+        and func.value.id == "sys"
+    ):
+        return True
+    parent = parents.get(id(node))
+    return (
+        isinstance(func, ast.Name)
+        and func.id == "SystemExit"
+        and isinstance(parent, ast.Raise)
+        and parent.exc is node
+    )
+
+
+def assigned_names(node: ast.AST) -> list[str]:
+    if isinstance(node, ast.Assign):
+        return [t.id for t in node.targets if isinstance(t, ast.Name)]
+    if isinstance(node, (ast.AnnAssign, ast.NamedExpr)):
+        return [node.target.id] if isinstance(node.target, ast.Name) else []
+    return []
+
+
+def exited_names(nodes: list[ast.AST], parents: dict[int, ast.AST]) -> set[str]:
+    """Names handed to an exit call, so the two-step form still propagates."""
+    return {
+        arg.id
+        for node in nodes
+        if is_exit_call(node, parents)
+        for arg in node.args
+        if isinstance(arg, ast.Name)
+    }
+
+
+def propagates_exit_code(
+    node: ast.Call, parents: dict[int, ast.AST], exited: set[str]
+) -> bool:
+    parent = parents.get(id(node))
+    if (
+        isinstance(parent, ast.Call)
+        and node in parent.args
+        and is_exit_call(parent, parents)
+    ):
+        return True
+    return any(name in exited for name in assigned_names(parent))
+
+
+def runtime_nodes(node: ast.AST):
+    yield node
+    if isinstance(
+        node, (ast.FunctionDef, ast.AsyncFunctionDef, ast.ClassDef, ast.Lambda)
+    ):
+        return
+    for child in ast.iter_child_nodes(node):
+        yield from runtime_nodes(child)
+
+
+def find_bare_pytest_main(path: pathlib.Path) -> int | None:
+    try:
+        source = path.read_text(encoding="utf-8")
+    except (OSError, UnicodeDecodeError):
+        return None
+    if "__main__" not in source or _PYTEST_MAIN.search(source) is None:
+        return None
+    try:
+        tree = ast.parse(source, filename=str(path))
+    except SyntaxError:
+        return None
+
+    for node in ast.walk(tree):
+        if not isinstance(node, ast.If) or not is_main_guard(node.test):
+            continue
+        # Whole body at once: the call and the sys.exit() that propagates it
+        # are separate statements.
+        nodes = [n for statement in node.body for n in runtime_nodes(statement)]
+        parents = {
+            id(child): parent
+            for parent in nodes
+            for child in ast.iter_child_nodes(parent)
+        }
+        exited = exited_names(nodes, parents)
+        for candidate in nodes:
+            if is_pytest_main_call(candidate) and not propagates_exit_code(
+                candidate, parents, exited
+            ):
+                return candidate.lineno
+    return None
+
+
+def main(paths: list[str]) -> int:
+    offenders = []
+    for path_string in paths:
+        path = pathlib.Path(path_string)
+        line = find_bare_pytest_main(path)
+        if line is not None:
+            offenders.append(f"{path}:{line}")
+
+    if not offenders:
+        return 0
+
+    print(
+        "ERROR: pytest.main(...) in an __main__ block must propagate its exit "
+        "code with sys.exit(...) or raise SystemExit(...):"
+    )
+    for offender in offenders:
+        print(f"  {offender}")
+    return 1
+
+
+if __name__ == "__main__":
+    sys.exit(main(sys.argv[1:]))