Repository: apache/airflow
File: scripts/ci/prek/check_core_imports_in_sdk.py (Python)

Code before:
                and not node.module.startswith("airflow.sdk")
            ):
                import_names = ", ".join(alias.name for alias in node.names)
                statement = f"from {node.module} import {import_names}"
                mismatches.append((node.lineno, statement))
        # for `import airflow.x` statements
        elif isinstance(node, ast.Import):
            for alias in node.names:
                if alias.name.startswith("airflow.") and not alias.name.startswith("airflow.sdk"):
                    statement = f"import {alias.name}"
                    if alias.asname:
                        statement += f" as {alias.asname}"
                    mismatches.append((node.lineno, statement))

    return mismatches


def main():
    parser = argparse.ArgumentParser(description="Check for core imports in task-sdk files")
    parser.add_argument("files", nargs="*", help="Files to check")
    args = parser.parse_args()

    if not args.files:
        return

    files_to_check = [Path(f) for f in args.files if f.endswith(".py")]
    total_violations = 0

    for file_path in files_to_check:
        mismatches = check_file_for_core_imports(file_path)
        if mismatches:
            console.print(f"[red]{file_path}[/red]:")
            for line_num, statement in mismatches:
                console.print(f"  [yellow]Line {line_num}[/yellow]: {statement}")
            total_violations += len(mismatches)

    if total_violations:
        console.print()
        console.print(f"[red]Found {total_violations} core import(s) in task-sdk files[/red]")
        sys.exit(1)


if __name__ == "__main__":
    main()
    sys.exit(0)


Review comment (suggestion):
Is it okay to check only Python files in the pre‑commit hook, as it already filters them similar to other hooks? like -files :` ^src/airflow/.*\.py$`

Code after:
            ):
                import_names = ", ".join(alias.name for alias in node.names)
                statement = f"from {node.module} import {import_names}"
                mismatches.append((node.lineno, statement))
        # for `import airflow.x` statements
        elif isinstance(node, ast.Import):
            for alias in node.names:
                if alias.name.startswith("airflow.") and not alias.name.startswith("airflow.sdk"):
                    statement = f"import {alias.name}"
                    if alias.asname:
                        statement += f" as {alias.asname}"
                    mismatches.append((node.lineno, statement))

    return mismatches


def main():
    parser = argparse.ArgumentParser(description="Check for core imports in task-sdk files")
    parser.add_argument("files", nargs="*", help="Files to check")
    args = parser.parse_args()

    if not args.files:
        return

    total_violations = 0

    for file_path in [Path(f) for f in args.files]:
        mismatches = check_file_for_core_imports(file_path)
        if mismatches:
            console.print(f"[red]{file_path}[/red]:")
            for line_num, statement in mismatches:
                console.print(f"  [yellow]Line {line_num}[/yellow]: {statement}")
            total_violations += len(mismatches)

    if total_violations:
        console.print()
        console.print(f"[red]Found {total_violations} core import(s) in task-sdk files[/red]")
        sys.exit(1)


if __name__ == "__main__":
    main()
    sys.exit(0)