From 1e021fb894450470d59e42b8ed460d38205f0125 Mon Sep 17 00:00:00 2001 From: Amogh Desai Date: Sat, 17 Jan 2026 09:59:24 +0530 Subject: [PATCH 1/5] Introducing prek hook to detect airflow-core imports in task-sdk --- scripts/ci/prek/check_core_imports_in_sdk.py | 88 ++++++++++++++++++++ task-sdk/.pre-commit-config.yaml | 31 +++++++ 2 files changed, 119 insertions(+) create mode 100755 scripts/ci/prek/check_core_imports_in_sdk.py diff --git a/scripts/ci/prek/check_core_imports_in_sdk.py b/scripts/ci/prek/check_core_imports_in_sdk.py new file mode 100755 index 0000000000000..8cbdb8683c414 --- /dev/null +++ b/scripts/ci/prek/check_core_imports_in_sdk.py @@ -0,0 +1,88 @@ +#!/usr/bin/env python +# +# 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. +# /// script +# requires-python = ">=3.10,<3.11" +# dependencies = [ +# "rich>=13.6.0", +# ] +# /// +from __future__ import annotations + +import argparse +import ast +import sys +from pathlib import Path + +sys.path.insert(0, str(Path(__file__).parent.resolve())) +from common_prek_utils import console + + +def check_file_for_core_imports(file_path: Path) -> list[tuple[int, str]]: + """Check file for airflow-core imports (anything except airflow.sdk). Returns list of (line_num, import_statement).""" + try: + with open(file_path, encoding="utf-8") as f: + source = f.read() + tree = ast.parse(source, filename=str(file_path)) + except (OSError, UnicodeDecodeError, SyntaxError): + return [] + + mismatches = [] + + for node in ast.walk(tree): + if isinstance(node, ast.ImportFrom): + if ( + node.module + and node.module.startswith("airflow.") + and not node.module.startswith("airflow.sdk") + ): + import_names = ", ".join(alias.name for alias in node.names) + statement = f"from {node.module} import {import_names}" + mismatches.append((node.lineno, statement)) + + return mismatches + + +def main(): + parser = argparse.ArgumentParser(description="Check for core imports in task-sdk files") + parser.add_argument("files", nargs="*", help="Files to check") + args = parser.parse_args() + + if not args.files: + return + + files_to_check = [Path(f) for f in args.files if f.endswith(".py")] + total_violations = 0 + + for file_path in files_to_check: + mismatches = check_file_for_core_imports(file_path) + if mismatches: + console.print(f"[red]{file_path}[/red]:") + for line_num, statement in mismatches: + console.print(f" [yellow]Line {line_num}[/yellow]: {statement}") + total_violations += len(mismatches) + + if total_violations: + console.print() + console.print(f"[red]Found {total_violations} core import(s) in task-sdk files[/red]") + sys.exit(1) + + +if __name__ == "__main__": + main() + sys.exit(0) diff --git a/task-sdk/.pre-commit-config.yaml b/task-sdk/.pre-commit-config.yaml index 6947aa87caa3f..b7d43b535601f 100644 --- a/task-sdk/.pre-commit-config.yaml +++ b/task-sdk/.pre-commit-config.yaml @@ -24,6 +24,37 @@ default_language_version: repos: - repo: local hooks: + - id: check-core-imports + name: Check for core imports in task-sdk files + entry: ../scripts/ci/prek/check_core_imports_in_sdk.py + language: python + types: [python] + files: ^src/airflow/ + exclude: | + (?x) + # TODO: These files need to be refactored to remove core coupling + ^src/airflow/sdk/plugins_manager\.py$| + ^src/airflow/sdk/io/fs\.py$| + ^src/airflow/sdk/definitions/decorators/__init__\.py$| + ^src/airflow/sdk/definitions/decorators/__init__\.pyi$| + ^src/airflow/sdk/definitions/decorators/setup_teardown\.py$| + ^src/airflow/sdk/definitions/connection\.py$| + ^src/airflow/sdk/definitions/asset/__init__\.py$| + ^src/airflow/sdk/definitions/asset/decorators\.py$| + ^src/airflow/sdk/definitions/taskgroup\.py$| + ^src/airflow/sdk/definitions/mappedoperator\.py$| + ^src/airflow/sdk/definitions/deadline\.py$| + ^src/airflow/sdk/definitions/dag\.py$| + ^src/airflow/sdk/definitions/_internal/types\.py$| + ^src/airflow/sdk/definitions/_internal/node\.py$| + ^src/airflow/sdk/serde/serializers/kubernetes\.py$| + ^src/airflow/sdk/serde/serializers/iceberg\.py$| + ^src/airflow/sdk/serde/serializers/deltalake\.py$| + ^src/airflow/sdk/execution_time/task_runner\.py$| + ^src/airflow/sdk/execution_time/supervisor\.py$| + ^src/airflow/sdk/execution_time/secrets_masker\.py$| + ^src/airflow/sdk/bases/operator\.py$| + ^src/airflow/sdk/bases/operatorlink\.py$ - id: check-init-decorator-arguments name: Sync model __init__ and decorator arguments language: python From f44963952156cf6e9d31a619099099ec78a6c355 Mon Sep 17 00:00:00 2001 From: Amogh Desai Date: Sat, 17 Jan 2026 12:41:44 +0530 Subject: [PATCH 2/5] Introducing prek hook to detect airflow-core imports in task-sdk --- task-sdk/.pre-commit-config.yaml | 3 +++ 1 file changed, 3 insertions(+) diff --git a/task-sdk/.pre-commit-config.yaml b/task-sdk/.pre-commit-config.yaml index b7d43b535601f..2bca980a588e4 100644 --- a/task-sdk/.pre-commit-config.yaml +++ b/task-sdk/.pre-commit-config.yaml @@ -45,6 +45,9 @@ repos: ^src/airflow/sdk/definitions/mappedoperator\.py$| ^src/airflow/sdk/definitions/deadline\.py$| ^src/airflow/sdk/definitions/dag\.py$| + ^src/airflow/sdk/log.py$| + ^src/airflow/sdk/io/path.py$| + ^src/airflow/sdk/execution_time/execute_workload.py$| ^src/airflow/sdk/definitions/_internal/types\.py$| ^src/airflow/sdk/definitions/_internal/node\.py$| ^src/airflow/sdk/serde/serializers/kubernetes\.py$| From d5fae93bd2f49621411fb9b3b4602d80cd286c73 Mon Sep 17 00:00:00 2001 From: Amogh Desai Date: Sat, 17 Jan 2026 16:50:47 +0530 Subject: [PATCH 3/5] catch import airflow.x too --- scripts/ci/prek/check_core_imports_in_sdk.py | 9 +++++++++ 1 file changed, 9 insertions(+) diff --git a/scripts/ci/prek/check_core_imports_in_sdk.py b/scripts/ci/prek/check_core_imports_in_sdk.py index 8cbdb8683c414..76b74dad06ef9 100755 --- a/scripts/ci/prek/check_core_imports_in_sdk.py +++ b/scripts/ci/prek/check_core_imports_in_sdk.py @@ -45,6 +45,7 @@ def check_file_for_core_imports(file_path: Path) -> list[tuple[int, str]]: mismatches = [] for node in ast.walk(tree): + # for `from airflow.x import y` statements if isinstance(node, ast.ImportFrom): if ( node.module @@ -54,6 +55,14 @@ def check_file_for_core_imports(file_path: Path) -> list[tuple[int, str]]: import_names = ", ".join(alias.name for alias in node.names) statement = f"from {node.module} import {import_names}" mismatches.append((node.lineno, statement)) + # for `import airflow.x` statements + elif isinstance(node, ast.Import): + for alias in node.names: + if alias.name.startswith("airflow.") and not alias.name.startswith("airflow.sdk"): + statement = f"import {alias.name}" + if alias.asname: + statement += f" as {alias.asname}" + mismatches.append((node.lineno, statement)) return mismatches From 888950918f46c6d708d71c196e2a25ea83389f98 Mon Sep 17 00:00:00 2001 From: Amogh Desai Date: Mon, 19 Jan 2026 12:55:33 +0530 Subject: [PATCH 4/5] suggestion from Prab-27 --- scripts/ci/prek/check_core_imports_in_sdk.py | 5 ++--- 1 file changed, 2 insertions(+), 3 deletions(-) diff --git a/scripts/ci/prek/check_core_imports_in_sdk.py b/scripts/ci/prek/check_core_imports_in_sdk.py index 76b74dad06ef9..619318b38fe68 100755 --- a/scripts/ci/prek/check_core_imports_in_sdk.py +++ b/scripts/ci/prek/check_core_imports_in_sdk.py @@ -36,9 +36,8 @@ def check_file_for_core_imports(file_path: Path) -> list[tuple[int, str]]: """Check file for airflow-core imports (anything except airflow.sdk). Returns list of (line_num, import_statement).""" try: - with open(file_path, encoding="utf-8") as f: - source = f.read() - tree = ast.parse(source, filename=str(file_path)) + source = file_path.read_text(encoding="utf-8") + tree = ast.parse(source, filename=str(file_path)) except (OSError, UnicodeDecodeError, SyntaxError): return [] From c80058c8dbbace090324e0086af0d45434c68108 Mon Sep 17 00:00:00 2001 From: Amogh Desai Date: Mon, 19 Jan 2026 12:59:36 +0530 Subject: [PATCH 5/5] suggestion from Prab-27 --- scripts/ci/prek/check_core_imports_in_sdk.py | 3 +-- task-sdk/.pre-commit-config.yaml | 2 +- 2 files changed, 2 insertions(+), 3 deletions(-) diff --git a/scripts/ci/prek/check_core_imports_in_sdk.py b/scripts/ci/prek/check_core_imports_in_sdk.py index 619318b38fe68..393ae5a4c7b11 100755 --- a/scripts/ci/prek/check_core_imports_in_sdk.py +++ b/scripts/ci/prek/check_core_imports_in_sdk.py @@ -74,10 +74,9 @@ def main(): if not args.files: return - files_to_check = [Path(f) for f in args.files if f.endswith(".py")] total_violations = 0 - for file_path in files_to_check: + for file_path in [Path(f) for f in args.files]: mismatches = check_file_for_core_imports(file_path) if mismatches: console.print(f"[red]{file_path}[/red]:") diff --git a/task-sdk/.pre-commit-config.yaml b/task-sdk/.pre-commit-config.yaml index 2bca980a588e4..ecab133e4deef 100644 --- a/task-sdk/.pre-commit-config.yaml +++ b/task-sdk/.pre-commit-config.yaml @@ -29,7 +29,7 @@ repos: entry: ../scripts/ci/prek/check_core_imports_in_sdk.py language: python types: [python] - files: ^src/airflow/ + files: ^src/airflow/.*\.py$ exclude: | (?x) # TODO: These files need to be refactored to remove core coupling