This is an automated email from the ASF dual-hosted git repository.
potiuk pushed a commit to branch v3-3-test
in repository https://gitbox.apache.org/repos/asf/airflow.git
The following commit(s) were added to refs/heads/v3-3-test by this push:
new a75b7cdf239 [v3-3-test] Convert check-sdk-imports to capture
occurences than files (#71212) (#71398)
a75b7cdf239 is described below
commit a75b7cdf2396558c9d88d15efeafb67e0a469ec9
Author: github-actions[bot]
<41898282+github-actions[bot]@users.noreply.github.com>
AuthorDate: Mon Aug 17 17:57:52 2026 +0200
[v3-3-test] Convert check-sdk-imports to capture occurences than files
(#71212) (#71398)
* Consolidate AllowlistManager for prek hooks (#69327)
* Extract common AllowListManager base class for prek hooks
The allowlist load/save/generate/cleanup logic was duplicated across
check_new_airflow_exception_usage.py and check_provide_session_kwargs.py.
This consolidates the shared pattern into an abstract AllowListManager
in common_prek_utils.py, with subclasses providing only iter_files()
and count_occurrences().
* Move check loop and parse into AllowlistManager base class
The initial extraction left the scan/tighten/report loop duplicated
in both hook scripts and exposed `parse` as a staticmethod that
silently ignored the instance's `repo_root`.
- Convert `parse` from `@staticmethod` to instance method so it
always uses `self.repo_root` (fixes the latent footgun where
`_parse_tracked_allowlist` called parse unbound).
- Hoist the check loop into `AllowlistManager.check()` with
`violation_panel_text()` (abstract) and `format_violation_details()`
(overridable) hooks — both `_check_*` wrappers are now one-liners.
- Fix naming: `AllowListManager` → `AllowlistManager` to match the
subclass convention used throughout.
- Add test coverage for `check_new_airflow_exception_usage.py`.
(cherry picked from commit 7e47a528bca17f7ad5ba41c4a94e642727d33b9b)
* [v3-3-test] Convert check-sdk-imports to capture occurences than files
(#71212)
(cherry picked from commit d8cfc951f7edcb15a0473b25a26b9495ff53d3c1)
Co-authored-by: Amogh Desai <[email protected]>
---------
Co-authored-by: Jason(Zhe-You) Liu
<[email protected]>
Co-authored-by: Amogh Desai <[email protected]>
---
airflow-core/.pre-commit-config.yaml | 103 +--------
generated/known_sdk_imports_in_core.txt | 44 ++++
.../ci/prek/check_new_airflow_exception_usage.py | 141 ++-----------
scripts/ci/prek/check_provide_session_kwargs.py | 190 +++--------------
scripts/ci/prek/check_sdk_imports_in_core.py | 210 +++++++++++++++++--
scripts/ci/prek/common_prek_utils.py | 184 +++++++++++++++-
.../prek/test_check_new_airflow_exception_usage.py | 194 +++++++++++++++++
.../ci/prek/test_check_provide_session_kwargs.py | 44 ++--
.../ci/prek/test_check_sdk_imports_in_core.py | 232 ++++++++++++++++++++-
9 files changed, 925 insertions(+), 417 deletions(-)
diff --git a/airflow-core/.pre-commit-config.yaml
b/airflow-core/.pre-commit-config.yaml
index 11e8861647a..0d2fa10c1d8 100644
--- a/airflow-core/.pre-commit-config.yaml
+++ b/airflow-core/.pre-commit-config.yaml
@@ -307,113 +307,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
diff --git a/generated/known_sdk_imports_in_core.txt
b/generated/known_sdk_imports_in_core.txt
new file mode 100644
index 00000000000..b4b0877197f
--- /dev/null
+++ b/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::14
+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::17
+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
diff --git a/scripts/ci/prek/check_new_airflow_exception_usage.py
b/scripts/ci/prek/check_new_airflow_exception_usage.py
index 6ba97db865e..5c096021920 100755
--- a/scripts/ci/prek/check_new_airflow_exception_usage.py
+++ b/scripts/ci/prek/check_new_airflow_exception_usage.py
@@ -54,87 +54,40 @@ from __future__ import annotations
import argparse
import re
+from collections.abc import Iterable
from pathlib import Path
+from common_prek_utils import AIRFLOW_ROOT_PATH, AllowlistManager
from rich.console import Console
-from rich.panel import Panel
console = Console(color_system="standard", width=200)
-REPO_ROOT = Path(__file__).parents[3]
+REPO_ROOT = AIRFLOW_ROOT_PATH
# Match lines that actually raise AirflowException. Comment filtering is done
# in _raise_lines() by skipping lines whose stripped form starts with "#".
_RAISE_RE = re.compile(r"raise\s+AirflowException\b")
-class AllowlistManager:
+class AirflowExceptionAllowlistManager(AllowlistManager):
def __init__(self, allowlist_file: Path) -> None:
- self.allowlist_file = allowlist_file
-
- def load(self) -> dict[str, int]:
- """Return mapping of ``relative_path -> allowed_count``."""
- if not self.allowlist_file.exists():
- return {}
-
- result: dict[str, int] = {}
- for raw_line in self.allowlist_file.read_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:
- result[rel_str] = int(count_str)
- except ValueError:
- continue
-
- return result
-
- 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")
-
- def generate(self) -> int:
- console.print(f"Scanning [cyan]{REPO_ROOT}[/cyan] for raise
AirflowException …")
- counts: dict[str, int] = {}
- for path in _iter_python_files():
- n = len(_raise_lines(path))
- if n > 0:
- counts[str(path.relative_to(REPO_ROOT))] = n
-
- self.save(counts)
- total = sum(counts.values())
- console.print(
- f"[green]✓ Generated[/green]
[cyan]{self.allowlist_file.relative_to(REPO_ROOT)}[/cyan] "
- f"with [bold]{len(counts)}[/bold] files / [bold]{total}[/bold]
occurrences."
+ super().__init__(allowlist_file, repo_root=REPO_ROOT)
+
+ def iter_files(self) -> Iterable[Path]:
+ return _iter_python_files()
+
+ def count_occurrences(self, path: Path) -> int:
+ return len(_raise_lines(path))
+
+ def violation_panel_text(self) -> str:
+ return (
+ "New [bold]raise AirflowException[/bold] usage detected.\n"
+ "Define a dedicated exception class or use an existing specific
exception.\n"
+ "If this usage is intentional and pre-existing, run:\n\n"
+ " [cyan]uv run
./scripts/ci/prek/check_new_airflow_exception_usage.py --generate[/cyan]\n\n"
+ "to regenerate the allowlist, then commit the updated\n"
+ "[cyan]generated/known_airflow_exceptions.txt[/cyan]."
)
- return 0
-
- def cleanup(self) -> int:
- allowlist = self.load()
- if not allowlist:
- console.print("[yellow]Allowlist is empty – nothing to clean
up.[/yellow]")
- return 0
-
- stale: list[str] = [rel for rel in allowlist if not (REPO_ROOT /
rel).exists()]
- if stale:
- 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)
- console.print(
- f"\n[green]Updated[/green]
[cyan]{self.allowlist_file.relative_to(REPO_ROOT)}[/cyan]"
- )
- else:
- console.print("[green]✓ No stale entries found.[/green]")
- return 0
def _raise_lines(path: Path) -> list[str]:
@@ -162,57 +115,9 @@ def _iter_python_files() -> list[Path]:
def _check_airflow_exception_usage(
- files: list[Path], allowlist: dict[str, int], manager: AllowlistManager
+ files: list[Path], allowlist: dict[str, int], manager:
AirflowExceptionAllowlistManager
) -> int:
- violations: list[tuple[Path, int, int]] = []
- tightened: list[tuple[str, int, int]] = [] # (rel, old_count, new_count)
-
- for path in files:
- if not path.exists() or path.suffix != ".py":
- continue
- actual = len(_raise_lines(path))
- rel = str(path.relative_to(REPO_ROOT))
- allowed = allowlist.get(rel, 0)
- if actual > allowed:
- violations.append((path, actual, allowed))
- elif actual < allowed:
- # Usage was reduced — tighten the allowlist entry so it can't
creep back up.
- if actual == 0:
- del allowlist[rel]
- else:
- allowlist[rel] = actual
- tightened.append((rel, allowed, actual))
-
- if tightened:
- manager.save(allowlist)
- console.print(
- f"[green]✓ Tightened {len(tightened)} entr{'y' if len(tightened)
== 1 else 'ies'} "
- f"in
[cyan]{manager.allowlist_file.relative_to(REPO_ROOT)}[/cyan][/green] "
- "(stage the updated file):"
- )
- for rel, old, new in tightened:
- console.print(f" [cyan]{rel}[/cyan] {old} → {new}")
-
- if violations:
- console.print(
- Panel.fit(
- "New [bold]raise AirflowException[/bold] usage detected.\n"
- "Define a dedicated exception class or use an existing
specific exception.\n"
- "If this usage is intentional and pre-existing, run:\n\n"
- " [cyan]uv run
./scripts/ci/prek/check_new_airflow_exception_usage.py --generate[/cyan]\n\n"
- "to regenerate the allowlist, then commit the updated\n"
- "[cyan]generated/known_airflow_exceptions.txt[/cyan].",
- title="[red]❌ Check failed[/red]",
- border_style="red",
- )
- )
- for path, actual, allowed in violations:
- console.print(f" [cyan]{path.relative_to(REPO_ROOT)}[/cyan]
count={actual} (allowed={allowed})")
- return 1
-
- # Return 1 when the allowlist was tightened so pre-commit reports the file
as modified
- # and prompts the user to stage the updated allowlist.
- return 1 if tightened else 0
+ return manager.check(files, allowlist)
def main(argv: list[str] | None = None) -> int:
@@ -239,7 +144,7 @@ def main(argv: list[str] | None = None) -> int:
)
args = parser.parse_args(argv)
- manager = AllowlistManager(REPO_ROOT / "generated" /
"known_airflow_exceptions.txt")
+ manager = AirflowExceptionAllowlistManager(REPO_ROOT / "generated" /
"known_airflow_exceptions.txt")
if args.generate:
return manager.generate()
diff --git a/scripts/ci/prek/check_provide_session_kwargs.py
b/scripts/ci/prek/check_provide_session_kwargs.py
index 29152e13a4d..0025126d52f 100755
--- a/scripts/ci/prek/check_provide_session_kwargs.py
+++ b/scripts/ci/prek/check_provide_session_kwargs.py
@@ -65,14 +65,16 @@ import argparse
import ast
import subprocess
import typing
+from collections.abc import Iterable
from pathlib import Path
+from common_prek_utils import AIRFLOW_ROOT_PATH, AllowlistManager
from rich.console import Console
from rich.panel import Panel
console = Console(color_system="standard", width=200)
-REPO_ROOT = Path(__file__).parents[3]
+REPO_ROOT = AIRFLOW_ROOT_PATH
_PROVIDE_SESSION_DECORATOR = "provide_session"
@@ -133,109 +135,33 @@ def _count_violations(path: Path) -> int:
return sum(1 for _ in _iter_positional_session_in_provide_session(path))
-def _is_safe_relative(rel: str) -> bool:
- """Whether ``rel`` is a plain relative path that stays inside
``REPO_ROOT``.
-
- Rejects absolute paths and any entry that resolves outside the repo root so
- callers can ``relative_to(REPO_ROOT)`` without fear of a ``ValueError``.
- """
- 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:
+class ProvideSessionAllowlistManager(AllowlistManager):
def __init__(self, allowlist_file: Path) -> None:
- self.allowlist_file = allowlist_file
-
- @staticmethod
- def parse(text: str) -> dict[str, int]:
- """Parse allowlist *text* into a ``{rel_path: count}`` mapping.
-
- Same validation rules as :meth:`load` so we can reuse parsing for the
- on-disk allowlist *and* for the git-tracked version fetched from
- ``HEAD`` when guarding against entry-removal bypasses.
- """
- 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):
- 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]:
- if not self.allowlist_file.exists():
- return {}
- return self.parse(self.allowlist_file.read_text())
-
- 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")
-
- def generate(self) -> int:
- roots = ", ".join(_PROJECT_SOURCE_ROOTS)
- console.print(
- f"Scanning project source roots ([cyan]{roots}[/cyan]) under
[cyan]{REPO_ROOT}[/cyan] "
- "for @provide_session functions with positional session …"
+ super().__init__(allowlist_file, repo_root=REPO_ROOT)
+
+ def iter_files(self) -> Iterable[Path]:
+ return _iter_python_files()
+
+ def count_occurrences(self, path: Path) -> int:
+ return _count_violations(path)
+
+ def violation_panel_text(self) -> str:
+ return (
+ "New [bold]@provide_session[/bold] function with positional
``session`` detected.\n"
+ "Move ``session`` after a bare ``*`` in the signature so callers
must pass it by keyword:\n\n"
+ " [cyan]@provide_session\n"
+ " def foo(arg, *, session: Session = NEW_SESSION) -> None:
...[/cyan]\n\n"
+ "If this usage is intentional and pre-existing, run:\n\n"
+ " [cyan]uv run ./scripts/ci/prek/check_provide_session_kwargs.py
--generate[/cyan]\n\n"
+ "to regenerate the allowlist, then commit the updated\n"
+
"[cyan]scripts/ci/prek/known_provide_session_positional.txt[/cyan]."
)
- counts: dict[str, int] = {}
- for path in _iter_python_files():
- n = _count_violations(path)
- if n > 0:
- counts[str(path.relative_to(REPO_ROOT))] = n
-
- self.save(counts)
- total = sum(counts.values())
- console.print(
- f"[green]Generated[/green]
[cyan]{self.allowlist_file.relative_to(REPO_ROOT)}[/cyan] "
- f"with [bold]{len(counts)}[/bold] files / [bold]{total}[/bold]
offenders."
- )
- return 0
-
- def cleanup(self) -> int:
- allowlist = self.load()
- if not allowlist:
- console.print("[yellow]Allowlist is empty - nothing to clean
up.[/yellow]")
- return 0
- stale: list[str] = [rel for rel in allowlist if not (REPO_ROOT /
rel).exists()]
- if stale:
- 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)
- console.print(
- f"\n[green]Updated[/green]
[cyan]{self.allowlist_file.relative_to(REPO_ROOT)}[/cyan]"
- )
- else:
- console.print("[green]No stale entries found.[/green]")
- return 0
+ def format_violation_details(self, path: Path) -> list[str]:
+ return [
+ f" [dim]L{argument.lineno}[/dim] def {func.name}(...)"
+ for func, argument in
_iter_positional_session_in_provide_session(path)
+ ]
def _iter_python_files() -> list[Path]:
@@ -250,7 +176,7 @@ def _iter_python_files() -> list[Path]:
def _check_provide_session_kwargs(
- files: list[Path], allowlist: dict[str, int], manager: AllowlistManager
+ files: list[Path], allowlist: dict[str, int], manager:
ProvideSessionAllowlistManager
) -> int:
allowlist_file = manager.allowlist_file.resolve()
if any(p.resolve() == allowlist_file for p in files) and not
allowlist_file.exists():
@@ -265,57 +191,7 @@ def _check_provide_session_kwargs(
)
)
return 1
-
- 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 = _count_violations(path)
- rel = str(path.relative_to(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:
- manager.save(allowlist)
- console.print(
- f"[green]Tightened {len(tightened)} entr{'y' if len(tightened) ==
1 else 'ies'} "
- f"in
[cyan]{manager.allowlist_file.relative_to(REPO_ROOT)}[/cyan][/green] "
- "(stage the updated file):"
- )
- for rel, old, new in tightened:
- console.print(f" [cyan]{rel}[/cyan] {old} -> {new}")
-
- if violations:
- console.print(
- Panel.fit(
- "New [bold]@provide_session[/bold] function with positional
``session`` detected.\n"
- "Move ``session`` after a bare ``*`` in the signature so
callers must pass it by keyword:\n\n"
- " [cyan]@provide_session\n"
- " def foo(arg, *, session: Session = NEW_SESSION) -> None:
...[/cyan]\n\n"
- "If this usage is intentional and pre-existing, run:\n\n"
- " [cyan]uv run
./scripts/ci/prek/check_provide_session_kwargs.py --generate[/cyan]\n\n"
- "to regenerate the allowlist, then commit the updated\n"
-
"[cyan]scripts/ci/prek/known_provide_session_positional.txt[/cyan].",
- title="[red]Check failed[/red]",
- border_style="red",
- )
- )
- for path, actual, allowed in violations:
- console.print(f" [cyan]{path.relative_to(REPO_ROOT)}[/cyan]
count={actual} (allowed={allowed})")
- for func, argument in
_iter_positional_session_in_provide_session(path):
- console.print(f" [dim]L{argument.lineno}[/dim] def
{func.name}(...)")
- return 1
-
- return 1 if tightened else 0
+ return manager.check(files, allowlist)
def main(argv: list[str] | None = None) -> int:
@@ -342,7 +218,7 @@ def main(argv: list[str] | None = None) -> int:
)
args = parser.parse_args(argv)
- manager = AllowlistManager(Path(__file__).parent /
"known_provide_session_positional.txt")
+ manager = ProvideSessionAllowlistManager(Path(__file__).parent /
"known_provide_session_positional.txt")
if args.generate:
return manager.generate()
@@ -366,7 +242,7 @@ def main(argv: list[str] | None = None) -> int:
return _check_provide_session_kwargs(paths, allowlist, manager)
-def _parse_tracked_allowlist(manager: AllowlistManager) -> dict[str, int]:
+def _parse_tracked_allowlist(manager: ProvideSessionAllowlistManager) ->
dict[str, int]:
"""Return the allowlist as recorded at ``HEAD`` (the git-tracked version).
Used by :func:`_expand_for_allowlist_edits` so that *removing* an entry
@@ -390,11 +266,11 @@ def _parse_tracked_allowlist(manager: AllowlistManager)
-> dict[str, int]:
return {}
if completed.returncode != 0:
return {}
- return AllowlistManager.parse(completed.stdout)
+ return manager.parse(completed.stdout)
def _expand_for_allowlist_edits(
- paths: list[Path], manager: AllowlistManager, allowlist: dict[str, int]
+ paths: list[Path], manager: ProvideSessionAllowlistManager, allowlist:
dict[str, int]
) -> list[Path]:
"""Add allowlisted files when the allowlist itself is being changed.
diff --git a/scripts/ci/prek/check_sdk_imports_in_core.py
b/scripts/ci/prek/check_sdk_imports_in_core.py
index 03192728930..aea9446305b 100755
--- a/scripts/ci/prek/check_sdk_imports_in_core.py
+++ b/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())
diff --git a/scripts/ci/prek/common_prek_utils.py
b/scripts/ci/prek/common_prek_utils.py
index 61b914c891f..559c2f0a047 100644
--- a/scripts/ci/prek/common_prek_utils.py
+++ b/scripts/ci/prek/common_prek_utils.py
@@ -16,6 +16,7 @@
# under the License.
from __future__ import annotations
+import abc
import ast
import difflib
import os
@@ -26,7 +27,7 @@ import subprocess
import sys
import textwrap
import time
-from collections.abc import Callable, Generator
+from collections.abc import Callable, Generator, Iterable
from contextlib import contextmanager
from pathlib import Path
from tempfile import NamedTemporaryFile, _TemporaryFileWrapper
@@ -58,10 +59,12 @@ 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
@@ -925,3 +928,182 @@ def parse_operations(
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())
+
+ 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")
+
+ @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
diff --git a/scripts/tests/ci/prek/test_check_new_airflow_exception_usage.py
b/scripts/tests/ci/prek/test_check_new_airflow_exception_usage.py
new file mode 100644
index 00000000000..f1abb7d6308
--- /dev/null
+++ b/scripts/tests/ci/prek/test_check_new_airflow_exception_usage.py
@@ -0,0 +1,194 @@
+# 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 textwrap
+from pathlib import Path
+
+import pytest
+from ci.prek import check_new_airflow_exception_usage as hook
+from ci.prek.check_new_airflow_exception_usage import (
+ AirflowExceptionAllowlistManager,
+ _check_airflow_exception_usage,
+ _raise_lines,
+)
+
+
[email protected]
+def create_fake_repo(tmp_path, monkeypatch):
+ monkeypatch.setattr(hook, "REPO_ROOT", tmp_path)
+
+ def _write(rel: str, code: str) -> Path:
+ path = tmp_path / rel
+ path.parent.mkdir(parents=True, exist_ok=True)
+ path.write_text(textwrap.dedent(code))
+ return path
+
+ return _write
+
+
+class TestRaiseLines:
+ def test_counts_raise_airflow_exception(self, write_python_file):
+ path = write_python_file(
+ """\
+ from airflow.exceptions import AirflowException
+ raise AirflowException("boom")
+ raise AirflowException("bang")
+ """
+ )
+ assert len(_raise_lines(path)) == 2
+
+ def test_ignores_commented_lines(self, write_python_file):
+ path = write_python_file(
+ """\
+ # raise AirflowException("commented out")
+ raise AirflowException("real")
+ """
+ )
+ assert len(_raise_lines(path)) == 1
+
+ def test_ignores_other_raises(self, write_python_file):
+ path = write_python_file(
+ """\
+ raise ValueError("not this")
+ raise TypeError("nor this")
+ """
+ )
+ assert len(_raise_lines(path)) == 0
+
+ def test_missing_file_returns_empty(self, tmp_path):
+ assert _raise_lines(tmp_path / "nonexistent.py") == []
+
+
+class TestAirflowExceptionAllowlistManager:
+ def test_load_missing_file_returns_empty(self, tmp_path):
+ manager = AirflowExceptionAllowlistManager(tmp_path / "missing.txt")
+ assert manager.load() == {}
+
+ def test_save_and_load_round_trip(self, tmp_path):
+ manager = AirflowExceptionAllowlistManager(tmp_path / "allowlist.txt")
+ manager.save({"b/file.py": 2, "a/file.py": 1})
+ text = (tmp_path / "allowlist.txt").read_text()
+ assert text.splitlines() == ["a/file.py::1", "b/file.py::2"]
+ assert manager.load() == {"a/file.py": 1, "b/file.py": 2}
+
+ def test_load_skips_blank_and_malformed_lines(self, tmp_path):
+ path = tmp_path / "allowlist.txt"
+ path.write_text("\nvalid/file.py::3\nnocount\n::5\nbad::notanumber\n")
+ assert AirflowExceptionAllowlistManager(path).load() ==
{"valid/file.py": 3}
+
+ @pytest.mark.usefixtures("create_fake_repo")
+ def test_load_skips_unsafe_entries(self, tmp_path):
+ path = tmp_path / "allowlist.txt"
+
path.write_text("airflow-core/src/airflow/safe.py::1\n../escape.py::1\n/etc/passwd::1\n")
+ assert AirflowExceptionAllowlistManager(path).load() ==
{"airflow-core/src/airflow/safe.py": 1}
+
+
+class TestCheckAirflowExceptionUsage:
+ def test_no_violations_passes(self, create_fake_repo, tmp_path):
+ path = create_fake_repo(
+ "airflow-core/src/airflow/clean.py",
+ """\
+ raise ValueError("specific exception")
+ """,
+ )
+ manager = AirflowExceptionAllowlistManager(tmp_path / "allowlist.txt")
+ assert _check_airflow_exception_usage([path], {}, manager) == 0
+
+ def test_new_violation_fails(self, create_fake_repo, tmp_path):
+ path = create_fake_repo(
+ "airflow-core/src/airflow/bad.py",
+ """\
+ raise AirflowException("boom")
+ """,
+ )
+ manager = AirflowExceptionAllowlistManager(tmp_path / "allowlist.txt")
+ assert _check_airflow_exception_usage([path], {}, manager) == 1
+
+ def test_violation_within_allowlist_passes(self, create_fake_repo,
tmp_path):
+ path = create_fake_repo(
+ "airflow-core/src/airflow/grandfathered.py",
+ """\
+ raise AirflowException("old")
+ """,
+ )
+ manager = AirflowExceptionAllowlistManager(tmp_path / "allowlist.txt")
+ allowlist = {"airflow-core/src/airflow/grandfathered.py": 1}
+ assert _check_airflow_exception_usage([path], allowlist, manager) == 0
+
+ def test_exceeding_allowlist_fails(self, create_fake_repo, tmp_path):
+ path = create_fake_repo(
+ "airflow-core/src/airflow/grew.py",
+ """\
+ raise AirflowException("one")
+ raise AirflowException("two")
+ """,
+ )
+ manager = AirflowExceptionAllowlistManager(tmp_path / "allowlist.txt")
+ allowlist = {"airflow-core/src/airflow/grew.py": 1}
+ assert _check_airflow_exception_usage([path], allowlist, manager) == 1
+
+ def test_reducing_violations_tightens_allowlist(self, create_fake_repo,
tmp_path):
+ path = create_fake_repo(
+ "airflow-core/src/airflow/improved.py",
+ """\
+ raise AirflowException("one remains")
+ """,
+ )
+ manager = AirflowExceptionAllowlistManager(tmp_path / "allowlist.txt")
+ allowlist = {"airflow-core/src/airflow/improved.py": 2}
+ assert _check_airflow_exception_usage([path], allowlist, manager) == 1
+ assert manager.load() == {"airflow-core/src/airflow/improved.py": 1}
+
+ def test_fixing_all_violations_removes_entry(self, create_fake_repo,
tmp_path):
+ path = create_fake_repo(
+ "airflow-core/src/airflow/fixed.py",
+ """\
+ raise ValueError("migrated away from AirflowException")
+ """,
+ )
+ manager = AirflowExceptionAllowlistManager(tmp_path / "allowlist.txt")
+ allowlist = {"airflow-core/src/airflow/fixed.py": 1}
+ assert _check_airflow_exception_usage([path], allowlist, manager) == 1
+ assert manager.load() == {}
+
+ def test_non_python_file_is_skipped(self, create_fake_repo, tmp_path):
+ path = create_fake_repo(
+ "airflow-core/src/airflow/not_python.txt",
+ "raise AirflowException('in a text file')\n",
+ )
+ manager = AirflowExceptionAllowlistManager(tmp_path / "allowlist.txt")
+ assert _check_airflow_exception_usage([path], {}, manager) == 0
+
+
+class TestCleanup:
+ def test_cleanup_removes_stale_entries(self, create_fake_repo, tmp_path):
+ create_fake_repo("airflow-core/src/airflow/keeper.py", "pass")
+ allowlist_path = tmp_path / "allowlist.txt"
+ manager = AirflowExceptionAllowlistManager(allowlist_path)
+ manager.save(
+ {
+ "airflow-core/src/airflow/keeper.py": 1,
+ "airflow-core/src/airflow/gone.py": 1,
+ }
+ )
+ assert manager.cleanup() == 0
+ assert manager.load() == {"airflow-core/src/airflow/keeper.py": 1}
+
+ def test_cleanup_empty_allowlist(self, tmp_path):
+ manager = AirflowExceptionAllowlistManager(tmp_path / "allowlist.txt")
+ assert manager.cleanup() == 0
diff --git a/scripts/tests/ci/prek/test_check_provide_session_kwargs.py
b/scripts/tests/ci/prek/test_check_provide_session_kwargs.py
index 78b85cd270b..4e96f1a413e 100644
--- a/scripts/tests/ci/prek/test_check_provide_session_kwargs.py
+++ b/scripts/tests/ci/prek/test_check_provide_session_kwargs.py
@@ -25,7 +25,7 @@ from pathlib import Path
import pytest
from ci.prek import check_provide_session_kwargs as hook
from ci.prek.check_provide_session_kwargs import (
- AllowlistManager,
+ ProvideSessionAllowlistManager,
_check_provide_session_kwargs,
_count_violations,
_expand_for_allowlist_edits,
@@ -228,13 +228,13 @@ class TestIterPositionalSessionInProvideSession:
assert _count_violations(path) == 1
-class TestAllowlistManager:
+class TestProvideSessionAllowlistManager:
def test_load_missing_file_returns_empty(self, tmp_path):
- manager = AllowlistManager(tmp_path / "missing.txt")
+ manager = ProvideSessionAllowlistManager(tmp_path / "missing.txt")
assert manager.load() == {}
def test_save_and_load_round_trip(self, tmp_path):
- manager = AllowlistManager(tmp_path / "allowlist.txt")
+ manager = ProvideSessionAllowlistManager(tmp_path / "allowlist.txt")
manager.save({"b/file.py": 2, "a/file.py": 1})
# Sorted by key in the file
text = (tmp_path / "allowlist.txt").read_text()
@@ -244,7 +244,7 @@ class TestAllowlistManager:
def test_load_skips_blank_and_malformed_lines(self, tmp_path):
path = tmp_path / "allowlist.txt"
path.write_text("\nvalid/file.py::3\nnocount\n::5\nbad::notanumber\n")
- assert AllowlistManager(path).load() == {"valid/file.py": 3}
+ assert ProvideSessionAllowlistManager(path).load() ==
{"valid/file.py": 3}
@pytest.mark.usefixtures("create_fake_repo")
def test_load_skips_unsafe_entries(self, tmp_path):
@@ -252,7 +252,7 @@ class TestAllowlistManager:
path = tmp_path / "allowlist.txt"
path.write_text("airflow-core/src/airflow/safe.py::1\n../escape.py::1\n/etc/passwd::1\n")
# `create_fake_repo` patches REPO_ROOT to tmp_path so the safety check
is meaningful.
- assert AllowlistManager(path).load() ==
{"airflow-core/src/airflow/safe.py": 1}
+ assert ProvideSessionAllowlistManager(path).load() ==
{"airflow-core/src/airflow/safe.py": 1}
class TestCheckProvideSessionKwargs:
@@ -265,7 +265,7 @@ class TestCheckProvideSessionKwargs:
pass
""",
)
- manager = AllowlistManager(tmp_path / "allowlist.txt")
+ manager = ProvideSessionAllowlistManager(tmp_path / "allowlist.txt")
assert _check_provide_session_kwargs([path], {}, manager) == 0
def test_new_violation_fails(self, create_fake_repo, tmp_path):
@@ -277,7 +277,7 @@ class TestCheckProvideSessionKwargs:
pass
""",
)
- manager = AllowlistManager(tmp_path / "allowlist.txt")
+ manager = ProvideSessionAllowlistManager(tmp_path / "allowlist.txt")
assert _check_provide_session_kwargs([path], {}, manager) == 1
def test_violation_within_allowlist_passes(self, create_fake_repo,
tmp_path):
@@ -289,7 +289,7 @@ class TestCheckProvideSessionKwargs:
pass
""",
)
- manager = AllowlistManager(tmp_path / "allowlist.txt")
+ manager = ProvideSessionAllowlistManager(tmp_path / "allowlist.txt")
allowlist = {"airflow-core/src/airflow/grandfathered.py": 1}
assert _check_provide_session_kwargs([path], allowlist, manager) == 0
@@ -306,7 +306,7 @@ class TestCheckProvideSessionKwargs:
pass
""",
)
- manager = AllowlistManager(tmp_path / "allowlist.txt")
+ manager = ProvideSessionAllowlistManager(tmp_path / "allowlist.txt")
allowlist = {"airflow-core/src/airflow/grew.py": 1}
assert _check_provide_session_kwargs([path], allowlist, manager) == 1
@@ -323,7 +323,7 @@ class TestCheckProvideSessionKwargs:
pass
""",
)
- manager = AllowlistManager(tmp_path / "allowlist.txt")
+ manager = ProvideSessionAllowlistManager(tmp_path / "allowlist.txt")
allowlist = {"airflow-core/src/airflow/improved.py": 2}
# Exit non-zero so pre-commit reports the modified allowlist
assert _check_provide_session_kwargs([path], allowlist, manager) == 1
@@ -338,7 +338,7 @@ class TestCheckProvideSessionKwargs:
pass
""",
)
- manager = AllowlistManager(tmp_path / "allowlist.txt")
+ manager = ProvideSessionAllowlistManager(tmp_path / "allowlist.txt")
allowlist = {"airflow-core/src/airflow/fixed.py": 1}
assert _check_provide_session_kwargs([path], allowlist, manager) == 1
assert manager.load() == {}
@@ -347,14 +347,14 @@ class TestCheckProvideSessionKwargs:
path = create_fake_repo(
"airflow-core/src/airflow/not_python.txt", "@provide_session\ndef
foo(session=N): pass\n"
)
- manager = AllowlistManager(tmp_path / "allowlist.txt")
+ manager = ProvideSessionAllowlistManager(tmp_path / "allowlist.txt")
assert _check_provide_session_kwargs([path], {}, manager) == 0
@pytest.mark.usefixtures("create_fake_repo")
def test_missing_allowlist_file_fails_loudly(self, tmp_path):
"""Passing the allowlist path when the file is missing must fail, not
silently pass."""
allowlist_path = tmp_path / "allowlist.txt"
- manager = AllowlistManager(allowlist_path)
+ manager = ProvideSessionAllowlistManager(allowlist_path)
assert not allowlist_path.exists()
assert _check_provide_session_kwargs([allowlist_path.resolve()], {},
manager) == 1
@@ -362,12 +362,12 @@ class TestCheckProvideSessionKwargs:
class TestExpandForAllowlistEdits:
def test_unchanged_when_allowlist_not_in_paths(self, create_fake_repo,
tmp_path):
py = create_fake_repo("airflow-core/src/airflow/x.py", "pass")
- manager = AllowlistManager(tmp_path / "allowlist.txt")
+ manager = ProvideSessionAllowlistManager(tmp_path / "allowlist.txt")
assert _expand_for_allowlist_edits([py], manager,
{"airflow-core/src/airflow/x.py": 1}) == [py]
def test_appends_allowlisted_files_when_allowlist_edited(self,
create_fake_repo, tmp_path):
allowlist_path = tmp_path / "allowlist.txt"
- manager = AllowlistManager(allowlist_path)
+ manager = ProvideSessionAllowlistManager(allowlist_path)
listed = create_fake_repo("airflow-core/src/airflow/listed.py", "pass")
# Pass a resolved path — matches production behavior (``main()``
resolves argv).
result = _expand_for_allowlist_edits(
@@ -383,7 +383,7 @@ class TestExpandForAllowlistEdits:
def test_detection_robust_to_symlinked_allowlist(self, create_fake_repo,
tmp_path):
"""A symlink pointing at the allowlist file must still trigger
expansion."""
allowlist_path = tmp_path / "allowlist.txt"
- manager = AllowlistManager(allowlist_path)
+ manager = ProvideSessionAllowlistManager(allowlist_path)
listed = create_fake_repo("airflow-core/src/airflow/listed.py", "pass")
manager.save({"airflow-core/src/airflow/listed.py": 1})
@@ -410,7 +410,7 @@ class TestExpandForAllowlistEdits:
""",
)
allowlist_path = tmp_path / "allowlist.txt"
- manager = AllowlistManager(allowlist_path)
+ manager = ProvideSessionAllowlistManager(allowlist_path)
manager.save({rel: 1})
create_git_repo("seed allowlist at HEAD")
@@ -429,7 +429,7 @@ class TestExpandForAllowlistEdits:
@pytest.mark.usefixtures("create_fake_repo")
def test_parse_tracked_allowlist_empty_when_no_git_history(self, tmp_path):
"""Without a git repo the git-tracked allowlist lookup returns empty
and does not crash."""
- manager = AllowlistManager(tmp_path / "allowlist.txt")
+ manager = ProvideSessionAllowlistManager(tmp_path / "allowlist.txt")
assert _parse_tracked_allowlist(manager) == {}
def test_re_validates_listed_files_so_loosening_cannot_bypass(self,
create_fake_repo, tmp_path, capsys):
@@ -448,7 +448,7 @@ class TestExpandForAllowlistEdits:
""",
)
allowlist_path = tmp_path / "allowlist.txt"
- manager = AllowlistManager(allowlist_path)
+ manager = ProvideSessionAllowlistManager(allowlist_path)
# Allowlist loosened to 5 although file only has 2 positional sessions.
allowlist = {rel: 5}
manager.save(allowlist)
@@ -467,7 +467,7 @@ class TestCleanup:
def test_cleanup_removes_stale_entries(self, create_fake_repo, tmp_path):
create_fake_repo("airflow-core/src/airflow/keeper.py", "pass")
allowlist_path = tmp_path / "allowlist.txt"
- manager = AllowlistManager(allowlist_path)
+ manager = ProvideSessionAllowlistManager(allowlist_path)
manager.save(
{
"airflow-core/src/airflow/keeper.py": 1,
@@ -478,5 +478,5 @@ class TestCleanup:
assert manager.load() == {"airflow-core/src/airflow/keeper.py": 1}
def test_cleanup_empty_allowlist(self, tmp_path):
- manager = AllowlistManager(tmp_path / "allowlist.txt")
+ manager = ProvideSessionAllowlistManager(tmp_path / "allowlist.txt")
assert manager.cleanup() == 0
diff --git a/scripts/tests/ci/prek/test_check_sdk_imports_in_core.py
b/scripts/tests/ci/prek/test_check_sdk_imports_in_core.py
index ed1bcbe057f..3a430c3c2c1 100644
--- a/scripts/tests/ci/prek/test_check_sdk_imports_in_core.py
+++ b/scripts/tests/ci/prek/test_check_sdk_imports_in_core.py
@@ -19,8 +19,23 @@ from __future__ import annotations
import textwrap
from pathlib import Path
+import check_sdk_imports_in_core as hook
import pytest
-from check_sdk_imports_in_core import check_file_for_sdk_imports
+from check_sdk_imports_in_core import SdkImportsAllowlistManager,
check_file_for_sdk_imports
+
+
[email protected]
+def create_fake_core_repo(tmp_path, monkeypatch):
+ monkeypatch.setattr(hook, "REPO_ROOT", tmp_path)
+ monkeypatch.setattr(hook, "CORE_SRC_ROOT", tmp_path / "airflow-core" /
"src" / "airflow")
+
+ def _write(rel: str, code: str) -> Path:
+ path = tmp_path / "airflow-core" / "src" / "airflow" / rel
+ path.parent.mkdir(parents=True, exist_ok=True)
+ path.write_text(textwrap.dedent(code))
+ return path
+
+ return _write
class TestCheckFileForSdkImports:
@@ -178,3 +193,218 @@ class TestNocheckMarker:
f = tmp_path / "example.py"
f.write_text(code)
assert check_file_for_sdk_imports(f) == expected
+
+
+class TestSdkImportsAllowlistRatchet:
+ def test_no_violations_passes(self, create_fake_core_repo, tmp_path):
+ path = create_fake_core_repo(
+ "models/clean.py",
+ """\
+ from airflow.models import DagRun
+ """,
+ )
+ manager = SdkImportsAllowlistManager(tmp_path / "allowlist.txt")
+ assert manager.check([path], {}) == 0
+
+ def test_new_violation_fails(self, create_fake_core_repo, tmp_path):
+ path = create_fake_core_repo(
+ "models/bad.py",
+ """\
+ from airflow.sdk import DAG
+ """,
+ )
+ manager = SdkImportsAllowlistManager(tmp_path / "allowlist.txt")
+ assert manager.check([path], {}) == 1
+
+ def test_violation_within_allowlist_passes(self, create_fake_core_repo,
tmp_path):
+ path = create_fake_core_repo(
+ "models/grandfathered.py",
+ """\
+ from airflow.sdk import DAG
+ """,
+ )
+ manager = SdkImportsAllowlistManager(tmp_path / "allowlist.txt")
+ allowlist = {"airflow-core/src/airflow/models/grandfathered.py": 1}
+ assert manager.check([path], allowlist) == 0
+
+ def test_exceeding_allowlist_fails(self, create_fake_core_repo, tmp_path):
+ path = create_fake_core_repo(
+ "models/grew.py",
+ """\
+ from airflow.sdk import DAG
+ from airflow.sdk.definitions.deadline import VariableInterval
+ """,
+ )
+ manager = SdkImportsAllowlistManager(tmp_path / "allowlist.txt")
+ allowlist = {"airflow-core/src/airflow/models/grew.py": 1}
+ assert manager.check([path], allowlist) == 1
+
+ def test_reducing_violations_tightens_allowlist(self,
create_fake_core_repo, tmp_path):
+ path = create_fake_core_repo(
+ "models/improved.py",
+ """\
+ from airflow.sdk import DAG
+ """,
+ )
+ manager = SdkImportsAllowlistManager(tmp_path / "allowlist.txt")
+ allowlist = {"airflow-core/src/airflow/models/improved.py": 2}
+ assert manager.check([path], allowlist) == 1
+ assert manager.load() ==
{"airflow-core/src/airflow/models/improved.py": 1}
+
+ def test_unparseable_file_does_not_tighten_allowlist(self,
create_fake_core_repo, tmp_path):
+ """A syntax error must not look like a legitimate drop to zero imports.
+
+ ``find_import_violations`` reports zero violations for a file it
cannot parse,
+ indistinguishable from a file that genuinely has none -- the allowlist
entry must
+ be left alone rather than tightened or deleted.
+ """
+ path = create_fake_core_repo(
+ "models/broken.py",
+ """\
+ from airflow.sdk import DAG
+ def broken(
+ """,
+ )
+ manager = SdkImportsAllowlistManager(tmp_path / "allowlist.txt")
+ manager.save({"airflow-core/src/airflow/models/broken.py": 1})
+
+ assert manager.check([path], manager.load()) == 0
+ assert manager.load() == {"airflow-core/src/airflow/models/broken.py":
1}
+
+ def test_fixing_all_violations_removes_entry(self, create_fake_core_repo,
tmp_path):
+ path = create_fake_core_repo(
+ "models/fixed.py",
+ """\
+ from airflow.models import DagRun
+ """,
+ )
+ manager = SdkImportsAllowlistManager(tmp_path / "allowlist.txt")
+ allowlist = {"airflow-core/src/airflow/models/fixed.py": 1}
+ assert manager.check([path], allowlist) == 1
+ assert manager.load() == {}
+
+ def test_noqa_marker_avoids_ratchet_entirely(self, create_fake_core_repo,
tmp_path):
+ path = create_fake_core_repo(
+ "models/one_off.py",
+ """\
+ from airflow.sdk import DAG # noqa: SDK001
+ """,
+ )
+ manager = SdkImportsAllowlistManager(tmp_path / "allowlist.txt")
+ assert manager.check([path], {}) == 0
+
+ def test_example_dags_are_not_checked(self, create_fake_core_repo,
tmp_path):
+ path = create_fake_core_repo(
+ "example_dags/example_simplest_dag.py",
+ """\
+ from airflow.sdk import DAG
+ from airflow.sdk import task
+ """,
+ )
+ manager = SdkImportsAllowlistManager(tmp_path / "allowlist.txt")
+ assert manager.check([path], {}) == 0
+
+ def test_non_python_file_is_skipped(self, create_fake_core_repo, tmp_path):
+ path = create_fake_core_repo(
+ "models/not_python.txt",
+ "from airflow.sdk import DAG\n",
+ )
+ manager = SdkImportsAllowlistManager(tmp_path / "allowlist.txt")
+ assert manager.check([path], {}) == 0
+
+
+class TestSdkImportsAllowlistCleanup:
+ def test_cleanup_removes_stale_entries(self, create_fake_core_repo,
tmp_path):
+ create_fake_core_repo("models/keeper.py", "from airflow.models import
DagRun\n")
+ allowlist_path = tmp_path / "allowlist.txt"
+ manager = SdkImportsAllowlistManager(allowlist_path)
+ manager.save(
+ {
+ "airflow-core/src/airflow/models/keeper.py": 0,
+ "airflow-core/src/airflow/models/gone.py": 1,
+ }
+ )
+ assert manager.cleanup() == 0
+ assert manager.load() == {"airflow-core/src/airflow/models/keeper.py":
0}
+
+ def test_cleanup_empty_allowlist(self, tmp_path):
+ manager = SdkImportsAllowlistManager(tmp_path / "allowlist.txt")
+ assert manager.cleanup() == 0
+
+
+class TestSdkImportsAllowlistGenerate:
+ def test_generate_records_current_occurrences(self, create_fake_core_repo,
tmp_path):
+ create_fake_core_repo("models/a.py", "from airflow.sdk import DAG\n")
+ create_fake_core_repo("models/b.py", "from airflow.models import
DagRun\n")
+ allowlist_path = tmp_path / "allowlist.txt"
+ manager = SdkImportsAllowlistManager(allowlist_path)
+ assert manager.generate() == 0
+ assert manager.load() == {"airflow-core/src/airflow/models/a.py": 1}
+
+ def test_main_generate_with_files_is_scoped(self, create_fake_core_repo,
tmp_path):
+ path = create_fake_core_repo("models/a.py", "from airflow.sdk import
DAG\n")
+ (tmp_path / "generated").mkdir()
+ allowlist_path = tmp_path / "generated" /
"known_sdk_imports_in_core.txt"
+
allowlist_path.write_text("airflow-core/src/airflow/models/other.py::2\n")
+
+ assert hook.main(["--generate", str(path)]) == 0
+ assert allowlist_path.read_text().splitlines() == [
+ "airflow-core/src/airflow/models/a.py::1",
+ "airflow-core/src/airflow/models/other.py::2",
+ ]
+
+ def test_generate_for_leaves_other_entries_untouched(self,
create_fake_core_repo, tmp_path):
+ target = create_fake_core_repo(
+ "models/target.py",
+ """\
+ from airflow.sdk import DAG
+ from airflow.sdk import task
+ """,
+ )
+ create_fake_core_repo("models/drifted.py", "from airflow.sdk import
DAG\n")
+ manager = SdkImportsAllowlistManager(tmp_path / "allowlist.txt")
+ manager.save({"airflow-core/src/airflow/models/target.py": 1})
+
+ assert manager.generate_for([target]) == 0
+ assert manager.load() == {"airflow-core/src/airflow/models/target.py":
2}
+
+ def test_generate_for_drops_entry_when_file_is_clean(self,
create_fake_core_repo, tmp_path):
+ path = create_fake_core_repo("models/fixed.py", "from airflow.models
import DagRun\n")
+ manager = SdkImportsAllowlistManager(tmp_path / "allowlist.txt")
+ manager.save(
+ {
+ "airflow-core/src/airflow/models/fixed.py": 1,
+ "airflow-core/src/airflow/models/other.py": 3,
+ }
+ )
+
+ assert manager.generate_for([path]) == 0
+ assert manager.load() == {"airflow-core/src/airflow/models/other.py":
3}
+
+ def test_generate_for_leaves_unparseable_file_untouched(self,
create_fake_core_repo, tmp_path):
+ path = create_fake_core_repo(
+ "models/broken.py",
+ """\
+ from airflow.sdk import DAG
+ def broken(
+ """,
+ )
+ manager = SdkImportsAllowlistManager(tmp_path / "allowlist.txt")
+ manager.save({"airflow-core/src/airflow/models/broken.py": 1})
+
+ assert manager.generate_for([path]) == 0
+ assert manager.load() == {"airflow-core/src/airflow/models/broken.py":
1}
+
+ def test_generate_for_skips_example_dags(self, create_fake_core_repo,
tmp_path):
+ path = create_fake_core_repo("example_dags/tutorial.py", "from
airflow.sdk import DAG\n")
+ manager = SdkImportsAllowlistManager(tmp_path / "allowlist.txt")
+
+ assert manager.generate_for([path]) == 0
+ assert manager.load() == {}
+
+ def test_generate_skips_example_dags(self, create_fake_core_repo,
tmp_path):
+ create_fake_core_repo("models/a.py", "from airflow.sdk import DAG\n")
+ create_fake_core_repo("example_dags/tutorial.py", "from airflow.sdk
import DAG\n")
+ manager = SdkImportsAllowlistManager(tmp_path / "allowlist.txt")
+ assert manager.generate() == 0
+ assert manager.load() == {"airflow-core/src/airflow/models/a.py": 1}
