scripts/ci/prek/common_prek_utils.py

# Licensed to the Apache Software Foundation (ASF) under one
# or more contributor license agreements.  See the NOTICE file
# distributed with this work for additional information
# regarding copyright ownership.  The ASF licenses this file
# to you under the Apache License, Version 2.0 (the
# "License"); you may not use this file except in compliance
# with the License.  You may obtain a copy of the License at
#
#   http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing,
# software distributed under the License is distributed on an
# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
# KIND, either express or implied.  See the License for the
# specific language governing permissions and limitations
# under the License.
from __future__ import annotations

import abc
import ast
import difflib
import os
import re
import shlex
import shutil
import subprocess
import sys
import textwrap
import time
from collections.abc import Callable, Generator, Iterable
from contextlib import contextmanager
from pathlib import Path
from tempfile import NamedTemporaryFile, _TemporaryFileWrapper
from typing import Any

AIRFLOW_ROOT_PATH = Path(__file__).parents[3].resolve()
AIRFLOW_CORE_ROOT_PATH = AIRFLOW_ROOT_PATH / "airflow-core"
AIRFLOW_CORE_SOURCES_PATH = AIRFLOW_CORE_ROOT_PATH / "src"
AIRFLOW_BREEZE_SOURCES_PATH = AIRFLOW_ROOT_PATH / "dev" / "breeze"
AIRFLOW_PROVIDERS_ROOT_PATH = AIRFLOW_ROOT_PATH / "providers"
AIRFLOW_TASK_SDK_ROOT_PATH = AIRFLOW_ROOT_PATH / "task-sdk"
AIRFLOW_TASK_SDK_SOURCES_PATH = AIRFLOW_TASK_SDK_ROOT_PATH / "src"

DEFAULT_PYTHON_MAJOR_MINOR_VERSION = "3.11"

# Maps a platform string (as declared in ``provider.yaml`` under ``excluded-platforms``)
# to the ``platform_machine`` values Python reports there. The two ARM spellings are not
# synonyms: ``platform.machine()`` returns ``aarch64`` on Linux and ``arm64`` on macOS, so
# each platform excludes exactly one of them. A provider whose native dependency is
# unavailable on every ARM target lists both platforms; one that only lacks a Linux ARM
# build lists ``linux/arm64`` alone and stays installable on Apple Silicon.
#
# Only the ``linux/*`` entries correspond to CI matrix platforms (see ``CI_PLATFORMS``);
# ``darwin/arm64`` never matches a CI run and exists purely to drive the install-time
# marker.
EXCLUDED_PLATFORM_MACHINES: dict[str, list[str]] = {
    "linux/arm64": ["aarch64"],
    "darwin/arm64": ["arm64"],
}

GITHUB_TOKEN_ENV_VARS = ("GH_TOKEN", "GITHUB_TOKEN")

try:
    from rich.console import Console
    from rich.panel import Panel

    console = Console(width=400, color_system="standard")
except ImportError:
    console = None  # type: ignore[assignment]
    Panel = None  # type: ignore[assignment,misc]


@contextmanager
def temporary_tsc_project(
    tsconfig_path: Path, files: list[str]
) -> Generator[_TemporaryFileWrapper, None, None]:
    """
    Create a temporary tsconfig.json file that extends the main tsconfig.json file.
    This is needed to run TypeScript compiler with specific files included only
    """
    if not tsconfig_path.exists():
        raise RuntimeError(f"Cannot find {tsconfig_path}")
    temp_tsconfig_path = NamedTemporaryFile(mode="wt", suffix=".json", dir=tsconfig_path.parent, delete=True)
    files_joined = ", ".join([f'"{file}"' for file in files])
    content = f'{{"extends": "./{tsconfig_path.name}", "include": [{files_joined}]}}'
    if console:
        console.print(f"[magenta]Creating temporary tsconfig.json at {temp_tsconfig_path.name}[/]")
        console.print(content)
    else:
        print(f"Creating temporary tsconfig.json at {temp_tsconfig_path.name}", file=sys.stderr)
        print(content, file=sys.stderr)
    temp_tsconfig_path.write(content)
    temp_tsconfig_path.flush()
    yield temp_tsconfig_path


def run_command(*args, **kwargs) -> None:
    """
    Run command with given arguments and return the result.
    """
    cmd = " ".join([shlex.quote(arg) for arg in args[0]])
    cwd = kwargs.get("cwd", os.getcwd())
    text = f"Running command: `{cmd}` in directory: `{cwd}`"
    if console:
        console.print(f"[magenta]{text}[/]")
    else:
        print("#" * min(len(text), 200), file=sys.stderr)
        print(text, file=sys.stderr)
        print("#" * min(len(text), 200), file=sys.stderr)
    time_start = time.monotonic()
    subprocess.check_call(*args, **kwargs)
    time_end = time.monotonic()
    if console:
        console.print(f"[green]After {text}[/]")
        console.print(f"[green]Command finished in {time_end - time_start:.2f} seconds[/]")
    else:
        print("#" * min(len(text), 200), file=sys.stderr)
        print(f"After {text}")
        print(f"Command finished in {time_end - time_start:.2f} seconds", file=sys.stderr)
        print("#" * min(len(text), 200), file=sys.stderr)


def read_airflow_version() -> str:
    ast_obj = ast.parse((AIRFLOW_CORE_SOURCES_PATH / "airflow" / "__init__.py").read_text(encoding="utf-8"))
    for node in ast_obj.body:
        if isinstance(node, ast.Assign):
            if node.targets[0].id == "__version__":  # type: ignore[attr-defined]
                return ast.literal_eval(node.value)

    raise RuntimeError("Couldn't find __version__ in AST")


GLOBAL_CONSTANTS_PATH = (
    AIRFLOW_ROOT_PATH / "dev" / "breeze" / "src" / "airflow_breeze" / "global_constants.py"
)


def _read_global_constants_assignment(name: str) -> Any:
    """Read a top-level assignment from global_constants.py.

    Handles both plain assignments (``NAME = ...``) and annotated assignments
    (``NAME: type = ...``). The value must be a literal so it can be safely
    evaluated with ``ast.literal_eval``.
    """
    tree = ast.parse(GLOBAL_CONSTANTS_PATH.read_text(encoding="utf-8"))
    for node in tree.body:
        if isinstance(node, ast.Assign):
            for target in node.targets:
                if isinstance(target, ast.Name) and target.id == name:
                    return ast.literal_eval(node.value)
        elif isinstance(node, ast.AnnAssign):
            if isinstance(node.target, ast.Name) and node.target.id == name and node.value is not None:
                return ast.literal_eval(node.value)
    raise RuntimeError(f"{name} not found in global_constants.py")


def read_allowed_kubernetes_versions() -> list[str]:
    """Parse ALLOWED_KUBERNETES_VERSIONS from global_constants.py (single source of truth).

    Returns versions without the ``v`` prefix, e.g. ``["1.30.13", "1.31.12", ...]``.
    """
    versions: list[str] = _read_global_constants_assignment("ALLOWED_KUBERNETES_VERSIONS")
    return [v.lstrip("v") for v in versions]


def read_allowed_python_major_minor_versions() -> list[str]:
    """Parse ALLOWED_PYTHON_MAJOR_MINOR_VERSIONS from global_constants.py (single source of truth)."""
    return list(_read_global_constants_assignment("ALLOWED_PYTHON_MAJOR_MINOR_VERSIONS"))


def read_current_postgres_versions() -> list[str]:
    """Parse CURRENT_POSTGRES_VERSIONS from global_constants.py (single source of truth)."""
    return list(_read_global_constants_assignment("CURRENT_POSTGRES_VERSIONS"))


def read_current_mysql_versions() -> list[str]:
    """The MySQL release versions Airflow currently tests with.

    Mirrors how ``CURRENT_MYSQL_VERSIONS`` is built in global_constants.py: the
    "old" releases plus the LTS releases, plus an innovation release when one is
    configured. Returns the numeric versions only (e.g. ``["8.0", "8.4"]``); the
    docs add the textual "Innovation" annotation on top of these.
    """
    versions: list[str] = list(_read_global_constants_assignment("MYSQL_OLD_RELEASES"))
    versions += list(_read_global_constants_assignment("MYSQL_LTS_RELEASES"))
    innovation = _read_global_constants_assignment("MYSQL_INNOVATION_RELEASE")
    if innovation:
        versions.append(innovation)
    return versions


def read_default_python_major_minor_version_for_images() -> str:
    """Parse DEFAULT_PYTHON_MAJOR_MINOR_VERSION_FOR_IMAGES from global_constants.py."""
    value = _read_global_constants_assignment("DEFAULT_PYTHON_MAJOR_MINOR_VERSION_FOR_IMAGES")
    if not isinstance(value, str):
        raise RuntimeError(
            "DEFAULT_PYTHON_MAJOR_MINOR_VERSION_FOR_IMAGES in global_constants.py "
            f"must be a string, got {type(value).__name__}"
        )
    return value


def pre_process_mypy_files(files: list[str]) -> list[str]:
    """Pre-process files passed to mypy.

    * Exclude conftest.py files and __init__.py files
    * When running build on non-main branch do not take providers into account.
    * When running "airflow-core" package, then we need to exclude providers.
    """

    files = [file for file in files if not file.endswith("conftest.py") and not file.endswith("__init__.py")]
    default_branch = os.environ.get("DEFAULT_BRANCH")
    if not default_branch or default_branch == "main":
        return files
    return [file for file in files if not file.startswith("providers")]


def is_hidden_within_root(path: Path, root: Path) -> bool:
    """Whether any path component below ``root`` is dot-prefixed."""
    return any(part.startswith(".") for part in path.relative_to(root).parts)


def insert_documentation(
    file_path: Path,
    content: list[str],
    header: str,
    footer: str,
    add_comment: bool = False,
    extra_information: str | None = None,
) -> bool:
    found = False
    old_content = file_path.read_text(encoding="utf-8")
    lines = old_content.splitlines(keepends=True)
    replacing = False
    result: list[str] = []
    for line in lines:
        if line.strip().startswith(header.strip()):
            replacing = True
            found = True
            result.append(line)
            if add_comment:
                result.extend(["# " + line if line != "\n" else "#\n" for line in content])
            else:
                result.extend(content)
        if line.strip().startswith(footer.strip()):
            replacing = False
        if not replacing:
            result.append(line)
    new_content = "".join(result)
    if not found:
        print(f"Header {header} not found in {file_path}")
        sys.exit(1)
    if new_content != old_content:
        file_path.write_text(new_content, encoding="utf-8")
        console.print(f"Updated {file_path} with {extra_information or 'generated documentation'}")
        return True
    return False


_UV_REQUIRED_VERSION_RE = re.compile(
    r"""^\s*required-version\s*=\s*["']\s*>=\s*(?P<ver>\d+(?:\.\d+){0,2})\s*["']""",
    re.MULTILINE,
)


def _parse_version(ver: str) -> tuple[int, ...]:
    """Turn "0.9.17" into (0, 9, 17). Extra pre-release/build suffixes are ignored."""
    match = re.match(r"^(\d+(?:\.\d+)*)", ver.strip())
    if not match:
        raise ValueError(f"Cannot parse version: {ver!r}")
    return tuple(int(part) for part in match.group(1).split("."))


def read_uv_required_min_version() -> tuple[str, tuple[int, ...]]:
    """Read the minimum uv version from the root ``pyproject.toml``.

    Parses ``[tool.uv] required-version = ">=X.Y.Z"`` and returns ``(raw, tuple)``.
    We parse by regex to avoid pulling a TOML dep into every prek script.
    """
    pyproject = (AIRFLOW_ROOT_PATH / "pyproject.toml").read_text(encoding="utf-8")
    # Narrow to the [tool.uv] section so we don't match a different required-version.
    match = re.search(r"^\[tool\.uv\]\s*$(?P<body>.*?)(?=^\[|\Z)", pyproject, re.MULTILINE | re.DOTALL)
    if not match:
        raise RuntimeError("`[tool.uv]` section not found in root pyproject.toml")
    ver_match = _UV_REQUIRED_VERSION_RE.search(match.group("body"))
    if not ver_match:
        raise RuntimeError('`required-version = ">=X.Y.Z"` not found under `[tool.uv]` in pyproject.toml')
    raw = ver_match.group("ver")
    return raw, _parse_version(raw)


def check_uv_version(uv_bin: str = "uv") -> None:
    """Fail the hook if ``uv_bin`` is older than ``[tool.uv] required-version``.

    Called manually by prek hooks that invoke ``uv`` (directly or via breeze) so a
    contributor with an outdated uv sees a clear error before the hook spends time
    running and emits a confusing downstream failure.
    """
    try:
        raw_min, min_tuple = read_uv_required_min_version()
    except Exception as exc:
        # Don't block hooks on parse bugs — warn and continue.
        message = f"Could not determine required uv version from pyproject.toml: {exc}"
        if console:
            console.print(f"[yellow]WARNING: {message}")
        else:
            print(f"WARNING: {message}")
        return

    try:
        output = subprocess.check_output([uv_bin, "--version"], text=True).strip()
    except (FileNotFoundError, subprocess.CalledProcessError) as exc:
        message = (
            f"Could not run `{uv_bin} --version` to verify uv version (required: >= {raw_min}). Error: {exc}"
        )
        if console:
            console.print(f"[red]{message}")
        else:
            print(message)
        sys.exit(1)

    match = re.search(r"\b(\d+\.\d+(?:\.\d+)?)", output)
    if not match:
        message = f"Unexpected `uv --version` output: {output!r}"
        if console:
            console.print(f"[yellow]WARNING: {message}")
        else:
            print(f"WARNING: {message}")
        return
    actual_raw = match.group(1)
    actual_tuple = _parse_version(actual_raw)
    if actual_tuple < min_tuple:
        message = (
            f"uv {actual_raw} at `{uv_bin}` is older than the project-required "
            f">= {raw_min} (see `[tool.uv] required-version` in pyproject.toml). "
            "Upgrade uv before running this hook, e.g.:\n"
            "  uv self update\n"
            "or, to refresh the project-pinned uv in the main venv (included in the "
            "`dev` dependency group via the `all` extras):\n"
            "  uv sync\n"
        )
        if console:
            console.print(f"[red]{message}")
        else:
            print(message)
        sys.exit(1)


BREEZE_SHIM_MARKER = "Apache Airflow breeze shim — managed by scripts/tools/setup_breeze"
BREEZE_SHIM_VERSION_PREFIX = "# breeze-shim-version:"
SETUP_BREEZE_SHIM_VERSION_PREFIX = "SHIM_VERSION="
SETUP_BREEZE_PATH = AIRFLOW_ROOT_PATH / "scripts" / "tools" / "setup_breeze"
BREEZE_LOCKED_VENV_PATH = AIRFLOW_BREEZE_SOURCES_PATH / ".venv"


def _read_shim_version(text: str, prefix: str) -> int | None:
    for line in text.splitlines():
        stripped = line.strip()
        if stripped.startswith(prefix):
            try:
                return int(stripped[len(prefix) :].strip().strip("\"'"))
            except ValueError:
                return None
    return None


def describe_breeze_not_running_from_lock() -> str | None:
    """Describe why the ``breeze`` on PATH does not run from ``dev/breeze/uv.lock``.

    Only the current shim (which dispatches to ``uv run --locked``) and the locked venv CI
    syncs run the versions the lock pins. An older shim and a legacy ``uv tool`` / ``pipx``
    install both resolve breeze's dependencies against the index instead, so anything derived
    from them — command hashes above all — reflects whatever the index served that day.

    :return: a description of the offending install, or None when breeze runs from the lock.
    """
    breeze_bin = shutil.which("breeze")
    if breeze_bin is None:
        return None
    resolved = Path(breeze_bin).resolve()
    if resolved.is_relative_to(BREEZE_LOCKED_VENV_PATH.resolve()):
        return None
    try:
        text = Path(breeze_bin).read_text(encoding="utf-8")
    except (OSError, UnicodeDecodeError):
        return None
    if BREEZE_SHIM_MARKER not in text:
        return f"`{breeze_bin}` is a legacy global install, which ignores the lock"
    expected_version = _read_shim_version(
        SETUP_BREEZE_PATH.read_text(encoding="utf-8"), SETUP_BREEZE_SHIM_VERSION_PREFIX
    )
    if expected_version is None:
        return None
    installed_version = _read_shim_version(text, BREEZE_SHIM_VERSION_PREFIX)
    if installed_version is not None and installed_version >= expected_version:
        return None
    installed_text = installed_version if installed_version is not None else "pre-versioning"
    return (
        f"the shim at `{breeze_bin}` needs to be upgraded "
        f"(installed: {installed_text}, current: {expected_version})"
    )


def initialize_breeze_prek(name: str, file: str):
    if name not in ("__main__", "__mp_main__"):
        raise SystemExit(
            "This file is intended to be executed as an executable program. You cannot use it as a module."
            f"To run this script, run the ./{file} command"
        )

    if os.environ.get("SKIP_BREEZE_PREK_HOOKS"):
        console.print("[yellow]Skipping breeze prek hooks as SKIP_BREEZE_PREK_HOOKS is set")
        sys.exit(0)
    # Breeze itself runs under uv, so enforce the project's minimum uv version up front.
    check_uv_version()
    if shutil.which("breeze") is None:
        console.print(
            "[red]The `breeze` command is not on path.[/]\n\n"
            "[yellow]Please install breeze. Recommended: run `./scripts/tools/setup_breeze` "
            "from the repo root — it installs a shim at `~/.local/bin/breeze` that runs breeze "
            "via `uv run --locked` from the current git worktree (see ADR 0017).\n"
            "Legacy global install (`uv tool install -e ./dev/breeze` or "
            "`pipx install -e ./dev/breeze`) still works but is no longer recommended.[/]\n\n"
            "[bright_blue]You can also set SKIP_BREEZE_PREK_HOOKS env variable to non-empty "
            "value to skip all breeze tests."
        )
        sys.exit(1)


def run_command_via_breeze_shell(
    cmd: list[str],
    python_version: str = DEFAULT_PYTHON_MAJOR_MINOR_VERSION,
    backend: str = "none",
    executor: str = "LocalExecutor",
    extra_env: dict[str, str] | None = None,
    project_name: str = "breeze-prek",
    skip_environment_initialization: bool = True,
    warn_image_upgrade_needed: bool = False,
    enable_pseudo_terminal: bool = False,
    **other_popen_kwargs,
) -> subprocess.CompletedProcess:
    extra_env = extra_env or {}
    subprocess_cmd: list[str] = [
        "breeze",
        "shell",
        "--python",
        python_version,
        "--backend",
        backend,
        "--executor",
        executor,
        "--quiet",
        "--restart",
        "--skip-image-upgrade-check",
        # Note: The terminal is disabled - because prek is run inside git without a pseudo-terminal
        "--tty",
        "enabled" if enable_pseudo_terminal else "disabled",
    ]
    if warn_image_upgrade_needed:
        subprocess_cmd.append("--warn-image-upgrade-needed")
    if skip_environment_initialization:
        subprocess_cmd.append("--skip-environment-initialization")
    if project_name:
        subprocess_cmd.extend(["--project-name", project_name])
    subprocess_cmd.append(" ".join([shlex.quote(arg) for arg in cmd]))
    new_env = {
        **os.environ,
        "SKIP_BREEZE_SELF_UPGRADE_CHECK": "true",
        "SKIP_GROUP_OUTPUT": "true",
        "SKIP_SAVING_CHOICES": "true",
        "ANSWER": "no",
        **extra_env,
    }

    if os.environ.get("VERBOSE_COMMANDS") or os.environ.get("CI") == "true":
        if console:
            console.print(
                f"[magenta]Running command: {' '.join([shlex.quote(item) for item in subprocess_cmd])}[/]"
            )
            console.print("[magenta]With environment:[/]")
            console.print(new_env)
        else:
            print(f"Running command: {' '.join([shlex.quote(item) for item in subprocess_cmd])}")
            print("With environment:")
            print(new_env)
    try:
        return subprocess.run(
            subprocess_cmd,
            check=False,
            text=True,
            **other_popen_kwargs,
            env=new_env,
        )
    finally:
        # Always clean up containers, networks, and volumes the breeze shell
        # invocation created — even if the subprocess raised (KeyboardInterrupt,
        # OSError, etc.). Without --volumes the next prek run inherits state
        # from the previous run, which is the bug this finally clause prevents.
        down_command = ["docker", "compose", "--progress", "quiet"]
        if project_name:
            down_command.extend(["--project-name", project_name])
        down_command.extend(["down", "--remove-orphans", "--volumes"])
        subprocess.run(down_command, check=False, stdout=subprocess.DEVNULL, stderr=subprocess.DEVNULL)


def run_command_via_breeze_run(
    cmd: list[str],
    python_version: str = DEFAULT_PYTHON_MAJOR_MINOR_VERSION,
    backend: str = "none",
    executor: str = "LocalExecutor",
    extra_env: dict[str, str] | None = None,
    project_name: str = "breeze-prek",
    skip_environment_initialization: bool = True,
    warn_image_upgrade_needed: bool = False,
    enable_pseudo_terminal: bool = False,
    **other_popen_kwargs,
) -> subprocess.CompletedProcess:
    extra_env = extra_env or {}
    # Kept for call-site compatibility. `breeze run` does not use `executor`.
    _ = executor
    subprocess_cmd: list[str] = [
        "breeze",
        "run",
        "--python",
        python_version,
        "--backend",
        backend,
        "--tty",
        "enabled" if enable_pseudo_terminal else "disabled",
    ]
    if not warn_image_upgrade_needed:
        subprocess_cmd.append("--skip-image-upgrade-check")
    if skip_environment_initialization:
        # `breeze run` always runs non-interactively; keep parameter for compatibility.
        pass
    if project_name:
        subprocess_cmd.extend(["--project-name", project_name])
    subprocess_cmd.extend(cmd)
    new_env = {
        **os.environ,
        "SKIP_BREEZE_SELF_UPGRADE_CHECK": "true",
        "SKIP_GROUP_OUTPUT": "true",
        "SKIP_SAVING_CHOICES": "true",
        "ANSWER": "no",
        **extra_env,
    }

    if os.environ.get("VERBOSE_COMMANDS") or os.environ.get("CI") == "true":
        if console:
            console.print(
                f"[magenta]Running command: {' '.join([shlex.quote(item) for item in subprocess_cmd])}[/]"
            )
            console.print("[magenta]With environment:[/]")
            console.print(new_env)
        else:
            print(f"Running command: {' '.join([shlex.quote(item) for item in subprocess_cmd])}")
            print("With environment:")
            print(new_env)
    return subprocess.run(
        subprocess_cmd,
        check=False,
        text=True,
        **other_popen_kwargs,
        env=new_env,
    )


class ConsoleDiff(difflib.Differ):
    def _dump(self, tag, x, lo, hi):
        """Generate comparison results for a same-tagged range."""
        for i in range(lo, hi):
            if tag == "+":
                yield f"[green]{tag} {x[i]}[/]"
            elif tag == "-":
                yield f"[red]{tag} {x[i]}[/]"
            else:
                yield f"{tag} {x[i]}"


def check_list_sorted(the_list: list[str], message: str, errors: list[str]) -> bool:
    sorted_list = sorted(set(the_list))
    if the_list == sorted_list:
        console.print(f"{message} is [green]ok[/]")
        console.print(the_list)
        console.print()
        return True
    console.print(f"{message} [red]NOK[/]")
    console.print(textwrap.indent("\n".join(ConsoleDiff().compare(the_list, sorted_list)), " " * 4))
    console.print()
    errors.append(f"ERROR in {message}. The elements are not sorted/unique.")
    return False


def validate_cmd_result(cmd_result, include_ci_env_check=False):
    if include_ci_env_check:
        if cmd_result.returncode != 0 and os.environ.get("CI") != "true":
            if console:
                console.print(
                    "\n[yellow]If you see strange stacktraces above, especially about missing imports "
                    "run this command:[/]\n"
                )
                console.print(
                    "[magenta]breeze ci-image build --python 3.11 --upgrade-to-newer-dependencies[/]\n"
                )
            else:
                print(
                    "\nIf you see strange stacktraces above, especially about missing imports "
                    "run this command:\nbreeze ci-image build --python 3.11 --upgrade-to-newer-dependencies\n"
                )

    elif cmd_result.returncode != 0:
        if console:
            console.print(
                "[warning]\nIf you see strange stacktraces above, "
                "run `breeze ci-image build --python 3.11` and try again."
            )
        else:
            print(
                "\nIf you see strange stacktraces above, "
                "run `breeze ci-image build --python 3.11` and try again."
            )
    sys.exit(cmd_result.returncode)


def get_provider_id_from_path(file_path: Path) -> str | None:
    """
    Get the provider id from the path of the file it belongs to.
    """
    for parent in file_path.parents:
        # This works fine for both new and old providers structure - because we moved provider.yaml to
        # the top-level of the provider and this code finding "providers"  will find the "providers" package
        # in old structure and "providers" directory in new structure - in both cases we can determine
        # the provider id from the relative folders
        if (parent / "provider.yaml").exists():
            for providers_root_candidate in parent.parents:
                if providers_root_candidate.name == "providers":
                    return parent.relative_to(providers_root_candidate).as_posix().replace("/", ".")
            return None
    return None


def get_provider_base_dir_from_path(file_path: Path) -> Path | None:
    """
    Get the provider base dir (where provider.yaml is) from the path of the file it belongs to.
    """
    for parent in file_path.parents:
        if (parent / "provider.yaml").exists():
            return parent
    return None


def get_provider_namespace_from_path(file_path: Path) -> str | None:
    """Get the namespace of the nested provider the file belongs to, None if it is not nested."""
    provider_id = get_provider_id_from_path(file_path)
    if not provider_id or "." not in provider_id:
        return None
    return provider_id.split(".")[0]


def is_duplicated_namespace_init(file_path: Path) -> bool:
    """Check whether the file is a namespace ``__init__.py`` repeated across the namespace."""
    if file_path.name != "__init__.py":
        return False
    namespace = get_provider_namespace_from_path(file_path)
    base_dir = get_provider_base_dir_from_path(file_path)
    if namespace is None or base_dir is None:
        return False
    return file_path.relative_to(base_dir).parts in {
        ("src", "airflow", "providers", namespace, "__init__.py"),
        ("tests", "unit", namespace, "__init__.py"),
        ("tests", "integration", namespace, "__init__.py"),
        ("tests", "system", namespace, "__init__.py"),
    }


def get_all_provider_ids(
    exclude_suspended_providers: bool = False, exclude_not_ready_providers: bool = False
) -> list[str]:
    """
    Get all providers from the new provider structure

    :param exclude_suspended_providers: skip providers whose state is ``suspended``
    :param exclude_not_ready_providers: skip providers whose state is ``not-ready`` - those have
        never been published, so anything describing what is installable must leave them out
    """
    all_provider_ids = []
    excluded_states = set()
    if exclude_suspended_providers:
        excluded_states.add("suspended")
    if exclude_not_ready_providers:
        excluded_states.add("not-ready")
    for provider_file in AIRFLOW_PROVIDERS_ROOT_PATH.rglob("provider.yaml"):
        if provider_file.is_relative_to(AIRFLOW_PROVIDERS_ROOT_PATH / "src"):
            continue
        if excluded_states:
            import yaml

            provider_info = yaml.safe_load(provider_file.read_text(encoding="utf-8"))
            if provider_info.get("state") in excluded_states:
                continue
        provider_id = get_provider_id_from_path(provider_file)
        if provider_id:
            all_provider_ids.append(provider_id)
    return all_provider_ids


def get_all_provider_yaml_files() -> list[Path]:
    """
    Get all providers from the new provider structure
    """
    all_provider_yaml_files = []
    for provider_file in AIRFLOW_PROVIDERS_ROOT_PATH.rglob("provider.yaml"):
        if provider_file.is_relative_to(AIRFLOW_PROVIDERS_ROOT_PATH / "src"):
            continue
        all_provider_yaml_files.append(provider_file)
    return all_provider_yaml_files


def get_all_provider_info_dicts() -> dict[str, dict]:
    """
    Get provider yaml info for all providers from the new provider structure
    """
    providers: dict[str, dict] = {}
    for provider_file in get_all_provider_yaml_files():
        provider_id = str(provider_file.parent.relative_to(AIRFLOW_PROVIDERS_ROOT_PATH)).replace(os.sep, ".")
        import yaml

        provider_info = yaml.safe_load(provider_file.read_text(encoding="utf-8"))
        if provider_info["state"] != "suspended":
            providers[provider_id] = provider_info
    return providers


_NOQA_RE = re.compile(r"#\s*noqa\s*:\s*([^\n]*)", re.IGNORECASE)
_NOQA_CODE_RE = re.compile(r"[A-Z]+\d+\b")


def _parse_noqa_codes(line: str) -> set[str]:
    """Extract codes from the leading comma-separated list in a ``# noqa: <codes>`` comment.

    Each code must be terminated by a word boundary, so tokens like ``SDK002x``
    or ``F401foo`` are not treated as the corresponding code.

    Anything after the first non-code token is treated as explanatory text and
    ignored, so ``# noqa: F401 - see SDK002 docs`` only yields ``{"F401"}``.
    """
    match = _NOQA_RE.search(line)
    if not match:
        return set()
    codes: set[str] = set()
    for raw in match.group(1).split(","):
        code_match = _NOQA_CODE_RE.match(raw.strip())
        if not code_match:
            break
        codes.add(code_match.group(0))
    return codes


def has_nocheck_marker(source_lines: list[str], node: ast.ImportFrom | ast.Import, nocheck_code: str) -> bool:
    """
    Check if the import statement has a ``# noqa: <codes>`` comment that lists
    ``nocheck_code`` on any of its lines. The code may appear anywhere in the
    comma-separated code list (e.g. ``# noqa: F401, SDK002``).
    """
    start = node.lineno
    end = node.end_lineno or start
    for lineno in range(start, end + 1):
        if lineno <= len(source_lines) and nocheck_code in _parse_noqa_codes(source_lines[lineno - 1]):
            return True
    return False


def find_import_violations(
    file_path: Path,
    *,
    is_violating_module: Callable[[str], bool],
    nocheck_code: str,
    check_plain_imports: bool = False,
) -> list[tuple[int, str]]:
    """
    Walk imports in ``file_path`` and return ``(lineno, statement)`` for each
    that matches ``is_violating_module`` and is not suppressed by a
    ``# noqa: <nocheck_code>`` comment.

    :param check_plain_imports: also check ``import x`` statements (in addition
        to ``from x import y``).
    """
    try:
        source = file_path.read_text(encoding="utf-8")
        tree = ast.parse(source, filename=str(file_path))
    except (OSError, UnicodeDecodeError, SyntaxError):
        return []

    source_lines = source.splitlines()
    violations: list[tuple[int, str]] = []

    for node in ast.walk(tree):
        if isinstance(node, ast.ImportFrom):
            if not node.module:
                continue
            if is_violating_module(node.module):
                violating_names = [alias.name for alias in node.names]
            else:
                # Catch ``from airflow import settings`` style imports where the
                # offending module is the dotted ``<module>.<name>`` path.
                violating_names = [
                    alias.name for alias in node.names if is_violating_module(f"{node.module}.{alias.name}")
                ]
            if not violating_names:
                continue
            if has_nocheck_marker(source_lines, node, nocheck_code):
                continue
            statement = f"from {node.module} import {', '.join(violating_names)}"
            violations.append((node.lineno, statement))
        elif check_plain_imports and isinstance(node, ast.Import):
            for alias in node.names:
                if is_violating_module(alias.name):
                    if has_nocheck_marker(source_lines, node, nocheck_code):
                        continue
                    statement = f"import {alias.name}"
                    if alias.asname:
                        statement += f" as {alias.asname}"
                    violations.append((node.lineno, statement))

    return violations


def report_import_violations(
    files: list[str],
    *,
    check_func: Callable[[Path], list[tuple[int, str]]],
    violation_label: str,
    nocheck_code: str | None = None,
    only_python_files: bool = False,
) -> None:
    """Run ``check_func`` on each file, print violations, and exit(1) if any are found.

    When ``nocheck_code`` is given, a hint pointing at the ``# noqa: <code>``
    escape hatch is printed alongside the failure summary.
    """
    file_paths = [Path(f) for f in files if not only_python_files or f.endswith(".py")]
    total_violations = 0

    for file_path in file_paths:
        mismatches = check_func(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} {violation_label}[/red]")
        if nocheck_code:
            console.print(
                f"[yellow]Hint:[/yellow] if an import above is intentional, append "
                f"`# noqa: {nocheck_code}` to the import line (single-line imports) "
                f"or to the opening/closing paren line (multi-line imports) to "
                f"suppress this check for that statement."
            )
        sys.exit(1)


def get_imports_from_file(file_path: Path, *, only_top_level: bool) -> list[str]:
    """
    Returns list of all imports in file.

    For following code:
    import os
    from collections import defaultdict
    import numpy as np
    from pandas import DataFrame as DF

    def inner():
        import json
        from pathlib import Path, PurePath
    from __future__ import annotations

    When only_top_level = False then returns
        ['os', 'collections.defaultdict', 'numpy', 'pandas.DataFrame']
    When only_top_level = False then returns
        ['os', 'collections.defaultdict', 'numpy', 'pandas.DataFrame', 'json', 'pathlib.Path', 'pathlib.PurePath']
    """
    root = ast.parse(file_path.read_text(encoding="utf-8"), file_path.name)
    imports: list[str] = []

    nodes = ast.iter_child_nodes(root) if only_top_level else ast.walk(root)
    for node in nodes:
        if isinstance(node, ast.Import):
            for alias in node.names:
                imports.append(alias.name)
        elif isinstance(node, ast.ImportFrom):
            if node.module == "__future__":
                continue
            for alias in node.names:
                name = alias.name
                fullname = f"{node.module}.{name}" if node.module else name
                imports.append(fullname)

    return imports


GIT_REPO_OVERRIDE_VARIABLES = (
    "GIT_INDEX_FILE",
    "GIT_DIR",
    "GIT_WORK_TREE",
    "GIT_PREFIX",
    "GIT_COMMON_DIR",
    "GIT_OBJECT_DIRECTORY",
    "GIT_ALTERNATE_OBJECT_DIRECTORIES",
)


def git_env_without_repo_overrides() -> dict[str, str]:
    """
    Return an environment for git commands that must not act on the repo being committed.

    ``git commit`` exports ``GIT_INDEX_FILE`` (and friends) to hooks; a child ``git worktree add``
    would otherwise check out into the committing repo's index instead of its own.
    """
    return {k: v for k, v in os.environ.items() if k not in GIT_REPO_OVERRIDE_VARIABLES}


def get_remote_for_main() -> str:
    """
    Return the remote name to use when fetching main.
    Prefers the remote that points to apache/airflow; otherwise uses origin.
    """
    result = subprocess.run(
        ["git", "remote", "-v"],
        capture_output=True,
        text=True,
        check=False,
    )
    if result.returncode != 0:
        return "origin"

    apache_remote = None
    origin_remote = None
    for line in result.stdout.splitlines():
        parts = line.split()
        if len(parts) >= 2:
            name, url = parts[0], parts[1]
            if "apache/airflow" in url:
                apache_remote = name
                break
            if name == "origin":
                origin_remote = name

    return apache_remote or origin_remote or "origin"


def env_without_github_tokens(env: dict[str, str] | None = None) -> dict[str, str]:
    cleaned_env = dict(os.environ if env is None else env)
    for token_env_var in GITHUB_TOKEN_ENV_VARS:
        cleaned_env.pop(token_env_var, None)
    return cleaned_env


def get_github_token_from_env(env: dict[str, str] | None = None) -> str | None:
    source_env = os.environ if env is None else env
    for token_env_var in GITHUB_TOKEN_ENV_VARS:
        token = source_env.get(token_env_var)
        if token:
            return token
    return None


def resolve_github_token(*, token: str | None = None, env: dict[str, str] | None = None) -> str | None:
    """Resolve a token while preventing ambient env tokens from shadowing ``gh auth login``."""
    if token:
        return token
    try:
        result = subprocess.run(
            ["gh", "auth", "token"],
            capture_output=True,
            text=True,
            check=False,
            env=env_without_github_tokens(env),
        )
    except FileNotFoundError:
        return get_github_token_from_env(env)
    if result.returncode == 0 and result.stdout.strip():
        return result.stdout.strip()
    return get_github_token_from_env(env)


def retrieve_gh_token(*, token: str | None = None, description: str, scopes: str) -> str:
    token = resolve_github_token(token=token)
    if not token:
        if not console:
            raise RuntimeError("Please add rich to your script dependencies and run it again")
        console.print(
            "[red]GITHUB_TOKEN environment variable is not set. "
            "This might lead to failures on rate limits.[/]\n"
            "You can fix that by installing `gh` and running `gh auth login` or "
            f"set it to a valid GitHub token with {scopes} scope. "
            f"You can create one by clicking the URL:\n\n"
            f"https://github.com/settings/tokens/new?scopes={scopes}&description={description}\n\n"
            "Once you have the token you can prepend prek command with GITHUB_TOKEN='<your token>' or"
            "set it in your environment with export GITHUB_TOKEN='<your token>'\n\n"
        )
        sys.exit(1)
    return token


def parse_operations(
    operations_file: Path, exclude_operation_classes: set, exclude_methods: set
) -> dict[str, list[str]]:
    """Parse airflowctl operations file and return a mapping of CLI group names to subcommands."""
    commands: dict[str, list[str]] = {}

    with open(operations_file) as f:
        tree = ast.parse(f.read(), filename=str(operations_file))

    for node in ast.walk(tree):
        if isinstance(node, ast.ClassDef) and node.name.endswith("Operations"):
            if node.name in exclude_operation_classes:
                continue

            group_name = node.name.replace("Operations", "").lower()
            commands[group_name] = []

            for child in node.body:
                if isinstance(child, ast.FunctionDef):
                    method_name = child.name
                    if method_name in exclude_methods or method_name.startswith("_"):
                        continue
                    subcommand = method_name.replace("_", "-")
                    commands[group_name].append(subcommand)

    return commands


def _is_safe_relative(rel: str, repo_root: Path) -> bool:
    """Whether ``rel`` is a plain relative path that stays inside ``repo_root``."""
    candidate = Path(rel)
    if candidate.is_absolute():
        return False
    try:
        (repo_root / candidate).resolve().relative_to(repo_root.resolve())
    except ValueError:
        return False
    return True


class AllowlistManager(abc.ABC):
    """Common base for prek hooks that track per-file occurrence counts in allowlist files.

    Subclasses implement :meth:`iter_files`, :meth:`count_occurrences`, and
    :meth:`violation_panel_text` to define what gets scanned, how violations
    are counted, and what help text to show.  Everything else — loading, saving,
    generating, cleaning up, and the check loop — is handled here.
    """

    def __init__(self, allowlist_file: Path, *, repo_root: Path = AIRFLOW_ROOT_PATH) -> None:
        self.allowlist_file = allowlist_file
        self.repo_root = repo_root

    def parse(self, text: str) -> dict[str, int]:
        """Parse allowlist *text* into a ``{rel_path: count}`` mapping.

        Entries that escape the repo root (absolute paths or ``..`` segments)
        are silently skipped.
        """
        result: dict[str, int] = {}
        for raw_line in text.splitlines():
            if not (stripped := raw_line.strip()):
                continue

            rel_str, _, count_str = stripped.rpartition("::")
            if not rel_str or not count_str:
                continue

            try:
                count = int(count_str)
            except ValueError:
                continue

            if not _is_safe_relative(rel_str, self.repo_root):
                if console:
                    console.print(
                        f"[yellow]Ignoring unsafe allowlist entry (escapes repo root):[/yellow] {rel_str}"
                    )
                continue

            result[rel_str] = count

        return result

    def load(self) -> dict[str, int]:
        """Return mapping of ``relative_path -> allowed_count``."""
        if not self.allowlist_file.exists():
            return {}
        return self.parse(self.allowlist_file.read_text(encoding="utf-8"))

    def save(self, counts: dict[str, int]) -> None:
        lines = [f"{rel}::{count}" for rel, count in sorted(counts.items())]
        self.allowlist_file.write_text("\n".join(lines) + "\n", encoding="utf-8")

    @abc.abstractmethod
    def iter_files(self) -> Iterable[Path]:
        """Return all files to scan during ``--generate`` or ``--all-files``."""

    @abc.abstractmethod
    def count_occurrences(self, path: Path) -> int:
        """Count the number of violations/occurrences in a single file."""

    @abc.abstractmethod
    def violation_panel_text(self) -> str:
        """Return the rich markup body for the violation help panel."""

    def format_violation_details(self, path: Path) -> list[str]:
        """Return extra detail lines for each violating file."""
        return []

    def check(self, files: list[Path], allowlist: dict[str, int]) -> int:
        """Run the check loop: compare counts, tighten entries, report violations."""
        violations: list[tuple[Path, int, int]] = []
        tightened: list[tuple[str, int, int]] = []

        for path in files:
            if not path.exists() or path.suffix != ".py":
                continue
            actual = self.count_occurrences(path)
            rel = str(path.relative_to(self.repo_root))
            allowed = allowlist.get(rel, 0)
            if actual > allowed:
                violations.append((path, actual, allowed))
            elif actual < allowed:
                if actual == 0:
                    del allowlist[rel]
                else:
                    allowlist[rel] = actual
                tightened.append((rel, allowed, actual))

        if tightened:
            self.save(allowlist)
            if console:
                console.print(
                    f"[green]Tightened {len(tightened)} entr{'y' if len(tightened) == 1 else 'ies'} "
                    f"in [cyan]{self.allowlist_file.relative_to(self.repo_root)}[/cyan][/green] "
                    "(stage the updated file):"
                )
                for rel, old, new in tightened:
                    console.print(f"  [cyan]{rel}[/cyan]  {old} → {new}")

        if violations:
            if console:
                console.print(
                    Panel.fit(
                        self.violation_panel_text(),
                        title="[red]Check failed[/red]",
                        border_style="red",
                    )
                )
                for path, actual, allowed in violations:
                    console.print(
                        f"  [cyan]{path.relative_to(self.repo_root)}[/cyan]  "
                        f"count={actual} (allowed={allowed})"
                    )
                    for detail in self.format_violation_details(path):
                        console.print(detail)
            return 1

        return 1 if tightened else 0

    def generate(self) -> int:
        if console:
            console.print(f"Scanning [cyan]{self.repo_root}[/cyan] …")
        counts: dict[str, int] = {}
        for path in self.iter_files():
            n = self.count_occurrences(path)
            if n > 0:
                counts[str(path.relative_to(self.repo_root))] = n

        self.save(counts)
        total = sum(counts.values())
        if console:
            console.print(
                f"[green]Generated[/green] [cyan]{self.allowlist_file.relative_to(self.repo_root)}[/cyan] "
                f"with [bold]{len(counts)}[/bold] files / [bold]{total}[/bold] occurrences."
            )
        return 0

    def cleanup(self) -> int:
        allowlist = self.load()
        if not allowlist:
            if console:
                console.print("[yellow]Allowlist is empty – nothing to clean up.[/yellow]")
            return 0

        stale: list[str] = [rel for rel in allowlist if not (self.repo_root / rel).exists()]
        if stale:
            if console:
                console.print(
                    f"[yellow]Removing {len(stale)} stale entr{'y' if len(stale) == 1 else 'ies'}:[/yellow]"
                )
                for s in sorted(stale):
                    console.print(f"  [dim]-[/dim] {s}")
            for s in stale:
                del allowlist[s]
            self.save(allowlist)
            if console:
                console.print(
                    f"\n[green]Updated[/green] [cyan]{self.allowlist_file.relative_to(self.repo_root)}[/cyan]"
                )
        else:
            if console:
                console.print("[green]No stale entries found.[/green]")
        return 0