Convert check-sdk-imports to capture occurences than files

airflow-core/.pre-commit-config.yaml

@@ -308,113 +308,12 @@ repos:
           ^tests/unit/core/test_configuration\.py$|
           ^tests/unit/models/test_renderedtifields\.py$|
           ^tests/unit/models/test_variable\.py$
-      - id: check-sdk-imports
+      - id: check-sdk-imports-in-core
         name: Check for SDK imports in core files
         entry: ../scripts/ci/prek/check_sdk_imports_in_core.py
         language: python
         types: [python]
         files: ^src/airflow/
-        exclude: |
-          (?x)
-          # Allow SDK imports in these legitimate locations
-          ^src/airflow/example_dags/.*\.py$|
-
-          # TODO: These files need to be refactored to remove SDK coupling
-          ^src/airflow/__init__\.py$|
-          ^src/airflow/api/common/mark_tasks\.py$|
-          ^src/airflow/api_fastapi/core_api/datamodels/assets\.py$|
-          ^src/airflow/api_fastapi/core_api/datamodels/connections\.py$|
-          ^src/airflow/api_fastapi/core_api/datamodels/hitl\.py$|
-          ^src/airflow/api_fastapi/core_api/datamodels/variables\.py$|
-          ^src/airflow/api_fastapi/core_api/routes/ui/grid\.py$|
-          ^src/airflow/api_fastapi/core_api/routes/ui/structure\.py$|
-          ^src/airflow/api_fastapi/core_api/services/public/connections\.py$|
-          ^src/airflow/api_fastapi/core_api/services/ui/connections\.py$|
-          ^src/airflow/api_fastapi/core_api/services/ui/grid\.py$|
-          ^src/airflow/api_fastapi/core_api/services/ui/task_group.py$|
-          ^src/airflow/api_fastapi/execution_api/routes/hitl\.py$|
-          ^src/airflow/api_fastapi/execution_api/routes/task_instances\.py$|
-          ^src/airflow/api_fastapi/execution_api/versions/v2026_04_06\.py$|
-          ^src/airflow/api_fastapi/logging/decorators\.py$|
-          ^src/airflow/assets/evaluation\.py$|
-          ^src/airflow/assets/manager\.py$|
-          ^src/airflow/cli/commands/connection_command\.py$|
-          ^src/airflow/cli/commands/task_command\.py$|
-          ^src/airflow/cli/commands/triggerer_command.py$|
-          ^src/airflow/configuration\.py$|
-          ^src/airflow/dag_processing/collection\.py$|
-          ^src/airflow/dag_processing/dagbag\.py$|
-          ^src/airflow/dag_processing/importers/.*\.py$|
-          ^src/airflow/dag_processing/manager\.py$|
-          ^src/airflow/dag_processing/processor\.py$|
-          ^src/airflow/datasets/metadata\.py$|
-          ^src/airflow/exceptions\.py$|
-          ^src/airflow/executors/base_executor\.py$|
-          ^src/airflow/jobs/triggerer_job_runner\.py$|
-          ^src/airflow/lineage/hook\.py$|
-          ^src/airflow/listeners/spec/asset\.py$|
-          ^src/airflow/listeners/spec/taskinstance\.py$|
-          ^src/airflow/logging/remote\.py$|
-          ^src/airflow/models/__init__\.py$|
-          ^src/airflow/models/asset\.py$|
-          ^src/airflow/models/baseoperator\.py$|
-          ^src/airflow/models/callback\.py$|
-          ^src/airflow/models/connection\.py$|
-          ^src/airflow/models/dag\.py$|
-          ^src/airflow/models/dagrun\.py$|
-          ^src/airflow/models/deadline\.py$|
-          ^src/airflow/models/expandinput\.py$|
-          ^src/airflow/models/mappedoperator\.py$|
-          ^src/airflow/models/operator\.py$|
-          ^src/airflow/models/param\.py$|
-          ^src/airflow/models/renderedtifields\.py$|
-          ^src/airflow/models/serialized_dag\.py$|
-          ^src/airflow/models/taskinstance\.py$|
-          ^src/airflow/models/taskinstancekey\.py$|
-          ^src/airflow/models/taskmap\.py$|
-          ^src/airflow/models/taskmixin\.py$|
-          ^src/airflow/models/taskreschedule\.py$|
-          ^src/airflow/models/trigger\.py$|
-          ^src/airflow/stats\.py$|
-          ^src/airflow/models/variable\.py$|
-          ^src/airflow/models/xcom\.py$|
-          ^src/airflow/models/xcom_arg\.py$|
-          ^src/airflow/operators/subdag\.py$|
-          ^src/airflow/plugins_manager\.py$|
-          ^src/airflow/providers_manager\.py$|
-          ^src/airflow/secrets/__init__.py$|
-          ^src/airflow/serialization/decoders\.py$|
-          ^src/airflow/serialization/definitions/[_/a-z]+\.py$|
-          ^src/airflow/serialization/encoders\.py$|
-          ^src/airflow/serialization/enums\.py$|
-          ^src/airflow/serialization/helpers\.py$|
-          ^src/airflow/serialization/serialized_objects\.py$|
-          ^src/airflow/settings\.py$|
-          ^src/airflow/task/task_runner/bash_task_runner\.py$|
-          ^src/airflow/task/task_runner/standard_task_runner\.py$|
-          ^src/airflow/ti_deps/deps/mapped_task_upstream_dep\.py$|
-          ^src/airflow/ti_deps/deps/prev_dagrun_dep\.py$|
-          ^src/airflow/ti_deps/deps/trigger_rule_dep\.py$|
-          ^src/airflow/timetables/assets\.py$|
-          ^src/airflow/timetables/base\.py$|
-          ^src/airflow/timetables/simple\.py$|
-          ^src/airflow/triggers/base\.py$|
-          ^src/airflow/utils/cli\.py$|
-          ^src/airflow/utils/context\.py$|
-          ^src/airflow/utils/dag_cycle_tester\.py$|
-          ^src/airflow/utils/dag_edges\.py$|
-          ^src/airflow/utils/dag_parsing_context\.py$|
-          ^src/airflow/utils/decorators\.py$|
-          ^src/airflow/utils/dot_renderer\.py$|
-          ^src/airflow/utils/edgemodifier\.py$|
-          ^src/airflow/utils/email\.py$|
-          ^src/airflow/utils/helpers\.py$|
-          ^src/airflow/utils/operator_helpers\.py$|
-          ^src/airflow/utils/session\.py$|
-          ^src/airflow/utils/task_group\.py$|
-          ^src/airflow/utils/trigger_rule\.py$|
-          ^src/airflow/utils/json\.py$|
-          ^src/airflow/utils/types\.py$
         ## ONLY ADD PREK HOOKS HERE THAT REQUIRE CI IMAGE
       - id: check-schema-defaults
         name: Check schema defaults match server-side defaults

generated/known_sdk_imports_in_core.txt

@@ -0,0 +1,44 @@
+airflow-core/src/airflow/__init__.py::1
+airflow-core/src/airflow/api_fastapi/execution_api/versions/v2026_04_06.py::1
+airflow-core/src/airflow/cli/commands/task_command.py::7
+airflow-core/src/airflow/cli/commands/triggerer_command.py::1
+airflow-core/src/airflow/configuration.py::1
+airflow-core/src/airflow/dag_processing/dagbag.py::1
+airflow-core/src/airflow/dag_processing/importers/base.py::1
+airflow-core/src/airflow/dag_processing/importers/python_importer.py::7
+airflow-core/src/airflow/dag_processing/manager.py::4
+airflow-core/src/airflow/dag_processing/processor.py::15
+airflow-core/src/airflow/exceptions.py::1
+airflow-core/src/airflow/executors/base_executor.py::3
+airflow-core/src/airflow/jobs/triggerer_job_runner.py::18
+airflow-core/src/airflow/models/__init__.py::5
+airflow-core/src/airflow/models/asset.py::2
+airflow-core/src/airflow/models/connection.py::5
+airflow-core/src/airflow/models/dag.py::1
+airflow-core/src/airflow/models/dagrun.py::2
+airflow-core/src/airflow/models/taskinstance.py::1
+airflow-core/src/airflow/models/trigger.py::4
+airflow-core/src/airflow/models/variable.py::8
+airflow-core/src/airflow/models/xcom_arg.py::1
+airflow-core/src/airflow/plugins_manager.py::1
+airflow-core/src/airflow/providers_manager.py::5
+airflow-core/src/airflow/secrets/__init__.py::1
+airflow-core/src/airflow/serialization/decoders.py::2
+airflow-core/src/airflow/serialization/definitions/baseoperator.py::1
+airflow-core/src/airflow/serialization/definitions/dag.py::2
+airflow-core/src/airflow/serialization/definitions/deadline.py::1
+airflow-core/src/airflow/serialization/definitions/mappedoperator.py::5
+airflow-core/src/airflow/serialization/encoders.py::11
+airflow-core/src/airflow/serialization/serialized_objects.py::16
+airflow-core/src/airflow/settings.py::1
+airflow-core/src/airflow/stats.py::1
+airflow-core/src/airflow/timetables/simple.py::1
+airflow-core/src/airflow/triggers/base.py::3
+airflow-core/src/airflow/utils/cli.py::2
+airflow-core/src/airflow/utils/context.py::2
+airflow-core/src/airflow/utils/dag_cycle_tester.py::1
+airflow-core/src/airflow/utils/dag_edges.py::1
+airflow-core/src/airflow/utils/dot_renderer.py::3
+airflow-core/src/airflow/utils/edgemodifier.py::1
+airflow-core/src/airflow/utils/json.py::4
+airflow-core/src/airflow/utils/operator_helpers.py::2

scripts/ci/prek/check_sdk_imports_in_core.py

@@ -22,16 +22,64 @@
 #   "rich>=13.6.0",
 # ]
 # ///
+"""Check that no new ``airflow.sdk`` imports are introduced in ``airflow-core``.
+
+All *existing* imports are recorded in ``generated/known_sdk_imports_in_core.txt``
+as ``relative/path::N`` entries (one per file), where ``N`` is the maximum
+number of ``airflow.sdk`` import statements allowed in that file. A file whose
+current count exceeds the recorded limit is treated as a violation -- core
+should not gain new runtime dependencies on the Task SDK. Two escape hatches
+exist: ``# noqa: SDK001`` on a single import line, which keeps that line out of
+the count and out of the allowlist entirely, or raising the file's recorded
+count via ``--generate FILE`` when a file genuinely needs several.
+
+``airflow-core/src/airflow/example_dags`` is skipped entirely: those files are
+Dags authored with the Task SDK, not core runtime code.
+
+Modes
+-----
+Default (files passed by prek/pre-commit):
+    Check only the supplied files; fail if any file's count exceeds the limit.
+    When a file's count has *decreased*, the allowlist entry is tightened
+    automatically and the hook exits with a non-zero code so that pre-commit
+    reports the modified allowlist -- just stage
+    ``generated/known_sdk_imports_in_core.txt`` and re-run.
+
+``--all-files``:
+    Walk all of ``airflow-core/src/airflow`` and check every ``.py`` file.
+
+``--cleanup``:
+    Remove entries for files that no longer exist. Safe to run at any time;
+    does not add new entries or raise limits.
+
+``--generate [FILE ...]``:
+    With files, record the current counts for just those files and leave every
+    other entry alone -- the usual way to approve a deliberate new import.
+    With no files, *rebuild* the allowlist from scratch; that also approves any
+    unrelated drift, so it is meant for the initial setup or a clean-up sprint.
+"""
+
 from __future__ import annotations
 
 import argparse
-import sys
+import ast
+from collections.abc import Iterable
 from pathlib import Path
 
-from common_prek_utils import find_import_violations, report_import_violations
+from common_prek_utils import AIRFLOW_CORE_ROOT_PATH, AllowlistManager, find_import_violations
+from rich.console import Console
+
+console = Console(color_system="standard", width=200)
+
+REPO_ROOT = AIRFLOW_CORE_ROOT_PATH.parent
+CORE_SRC_ROOT = AIRFLOW_CORE_ROOT_PATH / "src" / "airflow"
 
 NOCHECK_CODE = "SDK001"
 
+# Dags are authored *with* the task SDK, so example Dags shipped in core are not core code
+# for the purpose of this check and never take part in the ratchet.
+EXCLUDED_DIR_NAMES = frozenset({"example_dags"})
+
 
 def check_file_for_sdk_imports(file_path: Path) -> list[tuple[int, str]]:
     """Check file for airflow.sdk imports. Returns list of (line_num, import_statement)."""
@@ -42,23 +90,153 @@ def check_file_for_sdk_imports(file_path: Path) -> list[tuple[int, str]]:
     )
 
 
-def main():
-    parser = argparse.ArgumentParser(description="Check for SDK imports in airflow-core files")
-    parser.add_argument("files", nargs="*", help="Files to check")
-    args = parser.parse_args()
+def _can_parse(path: Path) -> bool:
+    """Whether *path* is valid enough Python for its import count to be trustworthy.
 
-    if not args.files:
-        return
+    ``find_import_violations`` swallows parse failures and reports zero violations for
+    them, which is indistinguishable from a file that genuinely has none. Callers must
+    filter those files out *before* counting, so a file mid-edit with a syntax error
+    never looks like a legitimate drop in its allowlist entry.
+    """
+    try:
+        ast.parse(path.read_text(encoding="utf-8"), filename=str(path))
+    except (OSError, UnicodeDecodeError, SyntaxError):
+        return False
+    return True
 
-    report_import_violations(
-        args.files,
-        check_func=check_file_for_sdk_imports,
-        violation_label="SDK import(s) in core files",
-        nocheck_code=NOCHECK_CODE,
-        only_python_files=True,
+
+class SdkImportsAllowlistManager(AllowlistManager):
+    def __init__(self, allowlist_file: Path) -> None:
+        super().__init__(allowlist_file, repo_root=REPO_ROOT)
+
+    def is_excluded(self, path: Path) -> bool:
+        return not EXCLUDED_DIR_NAMES.isdisjoint(path.parts)
+
+    def iter_files(self) -> Iterable[Path]:
+        return (
+            path for path in CORE_SRC_ROOT.rglob("*.py") if not self.is_excluded(path) and _can_parse(path)
+        )
+
+    def check(self, files: list[Path], allowlist: dict[str, int]) -> int:
+        checkable = []
+        for path in files:
+            if self.is_excluded(path):
+                continue
+            if path.exists() and path.suffix == ".py" and not _can_parse(path):
+                console.print(
+                    f"[yellow]Skipping unparsable file (not counted, allowlist entry left "
+                    f"untouched):[/yellow] {path}"
+                )
+                continue
+            checkable.append(path)
+        return super().check(checkable, allowlist)
+
+    def count_occurrences(self, path: Path) -> int:
+        return len(check_file_for_sdk_imports(path))
+
+    def generate_for(self, files: list[Path]) -> int:
+        """Rewrite the allowlist entries for *files* only, leaving every other entry untouched."""
+        allowlist = self.load()
+        updated: list[tuple[str, int, int]] = []
+
+        for path in files:
+            if self.is_excluded(path) or not path.exists() or path.suffix != ".py":
+                continue
+            if not _can_parse(path):
+                console.print(f"[yellow]Skipping unparsable file:[/yellow] {path}")
+                continue
+            try:
+                rel = str(path.resolve().relative_to(self.repo_root.resolve()))
+            except ValueError:
+                console.print(f"[yellow]Skipping file outside the repository:[/yellow] {path}")
+                continue
+            before = allowlist.get(rel, 0)
+            after = self.count_occurrences(path)
+            if before == after:
+                continue
+            if after:
+                allowlist[rel] = after
+            else:
+                del allowlist[rel]
+            updated.append((rel, before, after))
+
+        if not updated:
+            console.print("[green]Allowlist already up to date for the given files.[/green]")
+            return 0
+
+        self.save(allowlist)
+        console.print(
+            f"[green]Updated[/green] [cyan]{self.allowlist_file.relative_to(self.repo_root)}[/cyan]:"
+        )
+        for rel, before, after in updated:
+            console.print(f"  [cyan]{rel}[/cyan]  {before} -> {after}")
+        return 0
+
+    def violation_panel_text(self) -> str:
+        return (
+            "New [bold]airflow.sdk[/bold] import detected in airflow-core.\n"
+            "Core (scheduler/API server) should not gain new runtime dependencies "
+            "on the Task SDK.\n\n"
+            "The preferred fix is to drop the import. If it has to stay, pick one:\n\n"
+            "[bold]1. A single deliberate import line[/bold] -- append `# noqa: SDK001` to it.\n"
+            "   The line is then invisible to this check and never reaches the allowlist.\n\n"
+            "[bold]2. A file that genuinely needs several[/bold] -- record its new count:\n\n"
+            "  [cyan]uv run ./scripts/ci/prek/check_sdk_imports_in_core.py --generate FILE ...[/cyan]\n\n"
+            "   then commit the updated [cyan]generated/known_sdk_imports_in_core.txt[/cyan].\n"
+            "   Passing no FILE rebuilds the whole allowlist and approves any unrelated\n"
+            "   drift with it, so prefer the scoped form."
+        )
+
+
+def main(argv: list[str] | None = None) -> int:
+    parser = argparse.ArgumentParser(
+        description="Prevent new airflow.sdk imports in airflow-core.",
+        formatter_class=argparse.RawDescriptionHelpFormatter,
+        epilog=__doc__,
+    )
+    parser.add_argument("files", nargs="*", metavar="FILE", help="Files to check (provided by prek)")
+    parser.add_argument(
+        "--all-files",
+        action="store_true",
+        help="Check every Python file under airflow-core/src/airflow",
+    )
+    parser.add_argument(
+        "--cleanup",
+        action="store_true",
+        help="Remove stale entries from the allowlist and exit",
     )
+    parser.add_argument(
+        "--generate",
+        action="store_true",
+        help=(
+            "Record the current counts for the given FILEs and exit; "
+            "with no FILE, rebuild the whole allowlist from the current codebase"
+        ),
+    )
+    args = parser.parse_args(argv)
+
+    manager = SdkImportsAllowlistManager(REPO_ROOT / "generated" / "known_sdk_imports_in_core.txt")
+    files = [Path(f).resolve() for f in args.files]
+
+    if args.generate:
+        return manager.generate_for(files) if files else manager.generate()
+
+    if args.cleanup:
+        return manager.cleanup()
+
+    allowlist = manager.load()
+
+    if args.all_files:
+        return manager.check(list(manager.iter_files()), allowlist)
+
+    if not files:
+        console.print(
+            "[yellow]No files provided. Pass filenames or use --all-files to scan the whole repo.[/yellow]"
+        )
+        return 0
+
+    return manager.check(files, allowlist)
 
 
 if __name__ == "__main__":
-    main()
-    sys.exit(0)
+    raise SystemExit(main())