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

Code before:
# ]
# ///
from __future__ import annotations

import argparse
import ast
import sys
from pathlib import Path

sys.path.insert(0, str(Path(__file__).parent.resolve()))
from common_prek_utils import console


def check_file_for_core_imports(file_path: Path) -> list[tuple[int, str]]:
    """Check file for airflow-core imports (anything except airflow.sdk). Returns list of (line_num, import_statement)."""
    try:
        with open(file_path, encoding="utf-8") as f:
            source = f.read()
            tree = ast.parse(source, filename=str(file_path))
    except (OSError, UnicodeDecodeError, SyntaxError):
        return []

    mismatches = []

    for node in ast.walk(tree):
        if isinstance(node, ast.ImportFrom):
            if (
                node.module
                and node.module.startswith("airflow.")
                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))

    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)

Review comment (suggestion):
Yep - look at our sphinx.ext code:

```python
        if isinstance(node, (ast.Import, ast.ImportFrom)):
            for alias in node.names:
                module_prefix = f"{node.module}." if hasattr(node, "module") and node.module else ""
                imports[alias.asname or alias.name] = f"{module_prefix}{alias.name}"
```

You also should handle ast.Import

Code after:
# ]
# ///
from __future__ import annotations

import argparse
import ast
import sys
from pathlib import Path

sys.path.insert(0, str(Path(__file__).parent.resolve()))
from common_prek_utils import console


def check_file_for_core_imports(file_path: Path) -> list[tuple[int, str]]:
    """Check file for airflow-core imports (anything except airflow.sdk). Returns list of (line_num, import_statement)."""
    try:
        source = file_path.read_text(encoding="utf-8")
        tree = ast.parse(source, filename=str(file_path))
    except (OSError, UnicodeDecodeError, SyntaxError):
        return []

    mismatches = []

    for node in ast.walk(tree):
        # for `from airflow.x import y` statements
        if isinstance(node, ast.ImportFrom):
            if (
                node.module
                and node.module.startswith("airflow.")
                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()