This is an automated email from the ASF dual-hosted git repository.
jason810496 pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/airflow.git
The following commit(s) were added to refs/heads/main by this push:
new ff140cd0b4a Support inline ignore marker for check_core_imports_in_sdk
hook (#65358)
ff140cd0b4a is described below
commit ff140cd0b4a048287ed0c069b8de9659aebf6c34
Author: Jason(Zhe-You) Liu <[email protected]>
AuthorDate: Sat May 16 10:58:35 2026 +0800
Support inline ignore marker for check_core_imports_in_sdk hook (#65358)
* Support # nocheck: core-imports inline marker for
check_core_imports_in_sdk hook
* Address review comment
* refactor: streamline import checks and update nocheck markers to noqa
* refactor: streamline import checks and update nocheck markers to noqa
* Restrict noqa code parsing to the leading comma-separated list
A comment like `# noqa: F401 - see SDK002 docs` no longer
accidentally suppresses SDK002 — codes mentioned in the trailing
explanation are ignored.
* Add positive test cases for noqa codes with trailing explanation
Cover `# noqa: F401, SDK002 - needed for compat` (suppresses SDK002
because it's part of the leading code list) and the symmetric case
for the SDK001 hook.
* Tighten noqa code parsing and detect bare-module attribute imports
- Anchor `_NOQA_CODE_RE` with a word boundary so prefixes like
`SDK002x` no longer suppress `SDK002`.
- Extend the ImportFrom check to also test `<module>.<name>` so that
`from airflow import settings` is reported as a core import. Mixed
imports like `from airflow import sdk, settings` report only the
offending names.
* Print noqa suppression hint when an import violation is reported
After listing offending imports, ``report_import_violations`` now
echoes the ``# noqa: <code>`` escape hatch so users know the exact
marker to add for an intentional import without having to dig into
the hook source.
* fixup: Add SDK002 ignore in task-sdk
* Rename check_sdk_imports to check_sdk_imports_in_core
Mirrors the naming of the sibling check_core_imports_in_sdk hook so the
direction of each guard (what's checked in where) is obvious from the
filename.
---
airflow-core/.pre-commit-config.yaml | 2 +-
scripts/ci/prek/check_core_imports_in_sdk.py | 62 ++----
scripts/ci/prek/check_sdk_imports.py | 88 ---------
scripts/ci/prek/check_sdk_imports_in_core.py | 64 ++++++
scripts/ci/prek/common_prek_utils.py | 129 +++++++++++-
.../ci/prek/test_check_core_imports_in_sdk.py | 220 +++++++++++++++++++++
...mports.py => test_check_sdk_imports_in_core.py} | 42 +++-
task-sdk/src/airflow/sdk/execution_time/context.py | 2 +-
task-sdk/src/airflow/sdk/plugins_manager.py | 2 +-
9 files changed, 470 insertions(+), 141 deletions(-)
diff --git a/airflow-core/.pre-commit-config.yaml
b/airflow-core/.pre-commit-config.yaml
index 838747ab242..aea65d35b0e 100644
--- a/airflow-core/.pre-commit-config.yaml
+++ b/airflow-core/.pre-commit-config.yaml
@@ -299,7 +299,7 @@ repos:
^tests/unit/models/test_variable\.py$
- id: check-sdk-imports
name: Check for SDK imports in core files
- entry: ../scripts/ci/prek/check_sdk_imports.py
+ entry: ../scripts/ci/prek/check_sdk_imports_in_core.py
language: python
types: [python]
files: ^src/airflow/
diff --git a/scripts/ci/prek/check_core_imports_in_sdk.py
b/scripts/ci/prek/check_core_imports_in_sdk.py
index 700d81f5ea1..abdb194cd75 100755
--- a/scripts/ci/prek/check_core_imports_in_sdk.py
+++ b/scripts/ci/prek/check_core_imports_in_sdk.py
@@ -25,44 +25,24 @@
from __future__ import annotations
import argparse
-import ast
import sys
from pathlib import Path
-from common_prek_utils import console
+from common_prek_utils import find_import_violations, report_import_violations
+
+NOCHECK_CODE = "SDK002"
def check_file_for_core_imports(file_path: Path) -> list[tuple[int, str]]:
"""Check file for airflow-core imports (anything except airflow.sdk).
Returns list of (line_num, import_statement)."""
- try:
- source = file_path.read_text(encoding="utf-8")
- tree = ast.parse(source, filename=str(file_path))
- except (OSError, UnicodeDecodeError, SyntaxError):
- return []
-
- mismatches = []
-
- for node in ast.walk(tree):
- # for `from airflow.x import y` statements
- if isinstance(node, ast.ImportFrom):
- if (
- node.module
- and node.module.startswith("airflow.")
- and not node.module.startswith("airflow.sdk")
- ):
- import_names = ", ".join(alias.name for alias in node.names)
- statement = f"from {node.module} import {import_names}"
- mismatches.append((node.lineno, statement))
- # for `import airflow.x` statements
- elif isinstance(node, ast.Import):
- for alias in node.names:
- if alias.name.startswith("airflow.") and not
alias.name.startswith("airflow.sdk"):
- statement = f"import {alias.name}"
- if alias.asname:
- statement += f" as {alias.asname}"
- mismatches.append((node.lineno, statement))
-
- return mismatches
+ return find_import_violations(
+ file_path,
+ is_violating_module=lambda module: (
+ module.startswith("airflow.") and not
module.startswith("airflow.sdk")
+ ),
+ nocheck_code=NOCHECK_CODE,
+ check_plain_imports=True,
+ )
def main():
@@ -73,20 +53,12 @@ def main():
if not args.files:
return
- total_violations = 0
-
- 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]:")
- 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)
+ report_import_violations(
+ args.files,
+ check_func=check_file_for_core_imports,
+ violation_label="core import(s) in task-sdk files",
+ nocheck_code=NOCHECK_CODE,
+ )
if __name__ == "__main__":
diff --git a/scripts/ci/prek/check_sdk_imports.py
b/scripts/ci/prek/check_sdk_imports.py
deleted file mode 100755
index b35ebc9db8f..00000000000
--- a/scripts/ci/prek/check_sdk_imports.py
+++ /dev/null
@@ -1,88 +0,0 @@
-#!/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
-
-from common_prek_utils import console, has_nocheck_marker
-
-NOCHECK_MARKER = "# noqa: SDK001"
-
-
-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)."""
- 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 []
-
- source_lines = source.splitlines()
- mismatches = []
-
- for node in ast.walk(tree):
- if isinstance(node, ast.ImportFrom):
- if node.module and ("airflow.sdk" in node.module):
- if has_nocheck_marker(source_lines, node, NOCHECK_MARKER):
- continue
- 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 SDK imports in
airflow-core 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_sdk_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} SDK import(s) in core
files[/red]")
- sys.exit(1)
-
-
-if __name__ == "__main__":
- main()
- sys.exit(0)
diff --git a/scripts/ci/prek/check_sdk_imports_in_core.py
b/scripts/ci/prek/check_sdk_imports_in_core.py
new file mode 100755
index 00000000000..03192728930
--- /dev/null
+++ b/scripts/ci/prek/check_sdk_imports_in_core.py
@@ -0,0 +1,64 @@
+#!/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 sys
+from pathlib import Path
+
+from common_prek_utils import find_import_violations, report_import_violations
+
+NOCHECK_CODE = "SDK001"
+
+
+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)."""
+ return find_import_violations(
+ file_path,
+ is_violating_module=lambda module: "airflow.sdk" in module,
+ nocheck_code=NOCHECK_CODE,
+ )
+
+
+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()
+
+ if not args.files:
+ return
+
+ 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,
+ )
+
+
+if __name__ == "__main__":
+ main()
+ sys.exit(0)
diff --git a/scripts/ci/prek/common_prek_utils.py
b/scripts/ci/prek/common_prek_utils.py
index f59828bd297..7bcdfa3ab27 100644
--- a/scripts/ci/prek/common_prek_utils.py
+++ b/scripts/ci/prek/common_prek_utils.py
@@ -26,7 +26,7 @@ import subprocess
import sys
import textwrap
import time
-from collections.abc import Generator
+from collections.abc import Callable, Generator
from contextlib import contextmanager
from pathlib import Path
from tempfile import NamedTemporaryFile, _TemporaryFileWrapper
@@ -526,16 +526,137 @@ def get_all_provider_info_dicts() -> dict[str, dict]:
return providers
-def has_nocheck_marker(source_lines: list[str], node: ast.ImportFrom, marker:
str) -> bool:
- """Check if the import statement has the given nocheck marker comment on
any of its lines."""
+_NOQA_RE = re.compile(r"#\s*noqa\s*:\s*([^\n]*)", re.IGNORECASE)
+_NOQA_CODE_RE = re.compile(r"[A-Z]+\d+\b")
+
+
+def _parse_noqa_codes(line: str) -> set[str]:
+ """Extract codes from the leading comma-separated list in a ``# noqa:
<codes>`` comment.
+
+ Each code must be terminated by a word boundary, so tokens like ``SDK002x``
+ or ``F401foo`` are not treated as the corresponding code.
+
+ Anything after the first non-code token is treated as explanatory text and
+ ignored, so ``# noqa: F401 - see SDK002 docs`` only yields ``{"F401"}``.
+ """
+ match = _NOQA_RE.search(line)
+ if not match:
+ return set()
+ codes: set[str] = set()
+ for raw in match.group(1).split(","):
+ code_match = _NOQA_CODE_RE.match(raw.strip())
+ if not code_match:
+ break
+ codes.add(code_match.group(0))
+ return codes
+
+
+def has_nocheck_marker(source_lines: list[str], node: ast.ImportFrom |
ast.Import, nocheck_code: str) -> bool:
+ """
+ Check if the import statement has a ``# noqa: <codes>`` comment that lists
+ ``nocheck_code`` on any of its lines. The code may appear anywhere in the
+ comma-separated code list (e.g. ``# noqa: F401, SDK002``).
+ """
start = node.lineno
end = node.end_lineno or start
for lineno in range(start, end + 1):
- if lineno <= len(source_lines) and marker in source_lines[lineno - 1]:
+ if lineno <= len(source_lines) and nocheck_code in
_parse_noqa_codes(source_lines[lineno - 1]):
return True
return False
+def find_import_violations(
+ file_path: Path,
+ *,
+ is_violating_module: Callable[[str], bool],
+ nocheck_code: str,
+ check_plain_imports: bool = False,
+) -> list[tuple[int, str]]:
+ """
+ Walk imports in ``file_path`` and return ``(lineno, statement)`` for each
+ that matches ``is_violating_module`` and is not suppressed by a
+ ``# noqa: <nocheck_code>`` comment.
+
+ :param check_plain_imports: also check ``import x`` statements (in addition
+ to ``from x import y``).
+ """
+ try:
+ source = file_path.read_text(encoding="utf-8")
+ tree = ast.parse(source, filename=str(file_path))
+ except (OSError, UnicodeDecodeError, SyntaxError):
+ return []
+
+ source_lines = source.splitlines()
+ violations: list[tuple[int, str]] = []
+
+ for node in ast.walk(tree):
+ if isinstance(node, ast.ImportFrom):
+ if not node.module:
+ continue
+ if is_violating_module(node.module):
+ violating_names = [alias.name for alias in node.names]
+ else:
+ # Catch ``from airflow import settings`` style imports where
the
+ # offending module is the dotted ``<module>.<name>`` path.
+ violating_names = [
+ alias.name for alias in node.names if
is_violating_module(f"{node.module}.{alias.name}")
+ ]
+ if not violating_names:
+ continue
+ if has_nocheck_marker(source_lines, node, nocheck_code):
+ continue
+ statement = f"from {node.module} import {',
'.join(violating_names)}"
+ violations.append((node.lineno, statement))
+ elif check_plain_imports and isinstance(node, ast.Import):
+ for alias in node.names:
+ if is_violating_module(alias.name):
+ if has_nocheck_marker(source_lines, node, nocheck_code):
+ continue
+ statement = f"import {alias.name}"
+ if alias.asname:
+ statement += f" as {alias.asname}"
+ violations.append((node.lineno, statement))
+
+ return violations
+
+
+def report_import_violations(
+ files: list[str],
+ *,
+ check_func: Callable[[Path], list[tuple[int, str]]],
+ violation_label: str,
+ nocheck_code: str | None = None,
+ only_python_files: bool = False,
+) -> None:
+ """Run ``check_func`` on each file, print violations, and exit(1) if any
are found.
+
+ When ``nocheck_code`` is given, a hint pointing at the ``# noqa: <code>``
+ escape hatch is printed alongside the failure summary.
+ """
+ file_paths = [Path(f) for f in files if not only_python_files or
f.endswith(".py")]
+ total_violations = 0
+
+ for file_path in file_paths:
+ mismatches = check_func(file_path)
+ if mismatches:
+ console.print(f"[red]{file_path}[/red]:")
+ for line_num, statement in mismatches:
+ console.print(f" [yellow]Line {line_num}[/yellow]:
{statement}")
+ total_violations += len(mismatches)
+
+ if total_violations:
+ console.print()
+ console.print(f"[red]Found {total_violations} {violation_label}[/red]")
+ if nocheck_code:
+ console.print(
+ f"[yellow]Hint:[/yellow] if an import above is intentional,
append "
+ f"`# noqa: {nocheck_code}` to the import line (single-line
imports) "
+ f"or to the opening/closing paren line (multi-line imports) to
"
+ f"suppress this check for that statement."
+ )
+ sys.exit(1)
+
+
def get_imports_from_file(file_path: Path, *, only_top_level: bool) ->
list[str]:
"""
Returns list of all imports in file.
diff --git a/scripts/tests/ci/prek/test_check_core_imports_in_sdk.py
b/scripts/tests/ci/prek/test_check_core_imports_in_sdk.py
new file mode 100644
index 00000000000..70b3571ac53
--- /dev/null
+++ b/scripts/tests/ci/prek/test_check_core_imports_in_sdk.py
@@ -0,0 +1,220 @@
+# 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 check_core_imports_in_sdk import check_file_for_core_imports
+
+
+class TestCheckFileForCoreImports:
+ @pytest.mark.parametrize(
+ "code, expected",
+ [
+ pytest.param(
+ "from airflow.models import DagRun\n",
+ [(1, "from airflow.models import DagRun")],
+ id="from-import-core",
+ ),
+ pytest.param(
+ "import airflow.models\n",
+ [(1, "import airflow.models")],
+ id="import-core",
+ ),
+ pytest.param(
+ "import airflow.models as models\n",
+ [(1, "import airflow.models as models")],
+ id="import-core-aliased",
+ ),
+ pytest.param(
+ "from airflow.sdk import DAG\n",
+ [],
+ id="sdk-import-allowed",
+ ),
+ pytest.param(
+ "from airflow.sdk.definitions import dag\n",
+ [],
+ id="sdk-submodule-allowed",
+ ),
+ pytest.param(
+ "import airflow.sdk\n",
+ [],
+ id="import-sdk-allowed",
+ ),
+ pytest.param(
+ "import os\nimport sys\n",
+ [],
+ id="stdlib-only",
+ ),
+ pytest.param(
+ "from airflow import settings\n",
+ [(1, "from airflow import settings")],
+ id="from-airflow-import-name",
+ ),
+ pytest.param(
+ "from airflow import sdk\n",
+ [],
+ id="from-airflow-import-sdk-allowed",
+ ),
+ pytest.param(
+ "from airflow import settings, models\n",
+ [(1, "from airflow import settings, models")],
+ id="from-airflow-import-multiple-names",
+ ),
+ pytest.param(
+ "from airflow import sdk, settings\n",
+ [(1, "from airflow import settings")],
+ id="from-airflow-import-mixed-names",
+ ),
+ ],
+ )
+ def test_detects_core_imports(self, tmp_path: Path, code: str, expected:
list[tuple[int, str]]):
+ f = tmp_path / "example.py"
+ f.write_text(code)
+ assert check_file_for_core_imports(f) == expected
+
+
+class TestNocheckMarker:
+ @pytest.mark.parametrize(
+ "code, expected",
+ [
+ pytest.param(
+ "from airflow.models import DagRun # noqa: SDK002\n",
+ [],
+ id="from-import-suppressed",
+ ),
+ pytest.param(
+ "import airflow.models # noqa: SDK002\n",
+ [],
+ id="import-suppressed",
+ ),
+ pytest.param(
+ "from airflow.models import DagRun # noqa: SDK002 - needed
for compat\n",
+ [],
+ id="marker-with-extra-text",
+ ),
+ pytest.param(
+ textwrap.dedent("""\
+ from airflow.models import (
+ DagRun,
+ TaskInstance,
+ ) # noqa: SDK002
+ """),
+ [],
+ id="multiline-marker-on-closing-paren",
+ ),
+ pytest.param(
+ textwrap.dedent("""\
+ from airflow.models import ( # noqa: SDK002
+ DagRun,
+ TaskInstance,
+ )
+ """),
+ [],
+ id="multiline-marker-on-first-line",
+ ),
+ pytest.param(
+ textwrap.dedent("""\
+ from airflow.models import (
+ DagRun, # noqa: SDK002
+ TaskInstance,
+ )
+ """),
+ [],
+ id="multiline-marker-on-middle-line",
+ ),
+ pytest.param(
+ "from airflow.models import DagRun # noqa: E402\n",
+ [(1, "from airflow.models import DagRun")],
+ id="wrong-marker-not-suppressed",
+ ),
+ pytest.param(
+ textwrap.dedent("""\
+ from airflow.models import (
+ DagRun,
+ TaskInstance,
+ )
+ """),
+ [(1, "from airflow.models import DagRun, TaskInstance")],
+ id="multiline-without-marker-detected",
+ ),
+ pytest.param(
+ textwrap.dedent("""\
+ from airflow.models import DagRun # noqa: SDK002
+ from airflow.jobs import BaseJob
+ """),
+ [(2, "from airflow.jobs import BaseJob")],
+ id="only-marked-line-suppressed",
+ ),
+ pytest.param(
+ "from airflow.models import DagRun # noqa: F401, SDK002\n",
+ [],
+ id="combined-codes-target-last",
+ ),
+ pytest.param(
+ "from airflow.models import DagRun # noqa: SDK002, F401\n",
+ [],
+ id="combined-codes-target-first",
+ ),
+ pytest.param(
+ "from airflow.models import DagRun # noqa: E402, SDK002,
F401\n",
+ [],
+ id="combined-codes-target-middle",
+ ),
+ pytest.param(
+ "from airflow.models import DagRun # noqa:SDK002\n",
+ [],
+ id="no-space-after-colon",
+ ),
+ pytest.param(
+ "from airflow.models import DagRun # noqa: F401\n",
+ [(1, "from airflow.models import DagRun")],
+ id="other-code-only-not-suppressed",
+ ),
+ pytest.param(
+ "from airflow.models import DagRun # noqa: F401 - see SDK002
docs\n",
+ [(1, "from airflow.models import DagRun")],
+ id="code-in-explanation-not-suppressed",
+ ),
+ pytest.param(
+ "from airflow.models import DagRun # noqa: F401, SDK002 -
needed for compat\n",
+ [],
+ id="combined-codes-with-explanation-suppressed",
+ ),
+ pytest.param(
+ "from airflow.models import DagRun # noqa: SDK002x\n",
+ [(1, "from airflow.models import DagRun")],
+ id="partial-code-match-not-suppressed",
+ ),
+ pytest.param(
+ "from airflow import settings # noqa: SDK002\n",
+ [],
+ id="from-airflow-import-name-suppressed",
+ ),
+ pytest.param(
+ "from airflow.models import DagRun # noqa\n",
+ [(1, "from airflow.models import DagRun")],
+ id="bare-noqa-not-suppressed",
+ ),
+ ],
+ )
+ def test_nocheck_marker(self, tmp_path: Path, code: str, expected:
list[tuple[int, str]]):
+ f = tmp_path / "example.py"
+ f.write_text(code)
+ assert check_file_for_core_imports(f) == expected
diff --git a/scripts/tests/ci/prek/test_check_sdk_imports.py
b/scripts/tests/ci/prek/test_check_sdk_imports_in_core.py
similarity index 74%
rename from scripts/tests/ci/prek/test_check_sdk_imports.py
rename to scripts/tests/ci/prek/test_check_sdk_imports_in_core.py
index cf097027dc3..ed1bcbe057f 100644
--- a/scripts/tests/ci/prek/test_check_sdk_imports.py
+++ b/scripts/tests/ci/prek/test_check_sdk_imports_in_core.py
@@ -20,7 +20,7 @@ import textwrap
from pathlib import Path
import pytest
-from check_sdk_imports import check_file_for_sdk_imports
+from check_sdk_imports_in_core import check_file_for_sdk_imports
class TestCheckFileForSdkImports:
@@ -132,6 +132,46 @@ class TestNocheckMarker:
[(2, "from airflow.sdk.definitions import dag")],
id="only-marked-line-suppressed",
),
+ pytest.param(
+ "from airflow.sdk import DAG # noqa: F401, SDK001\n",
+ [],
+ id="combined-codes-target-last",
+ ),
+ pytest.param(
+ "from airflow.sdk import DAG # noqa: SDK001, F401\n",
+ [],
+ id="combined-codes-target-first",
+ ),
+ pytest.param(
+ "from airflow.sdk import DAG # noqa: E402, SDK001, F401\n",
+ [],
+ id="combined-codes-target-middle",
+ ),
+ pytest.param(
+ "from airflow.sdk import DAG # noqa:SDK001\n",
+ [],
+ id="no-space-after-colon",
+ ),
+ pytest.param(
+ "from airflow.sdk import DAG # noqa: F401\n",
+ [(1, "from airflow.sdk import DAG")],
+ id="other-code-only-not-suppressed",
+ ),
+ pytest.param(
+ "from airflow.sdk import DAG # noqa: F401 - see SDK001
docs\n",
+ [(1, "from airflow.sdk import DAG")],
+ id="code-in-explanation-not-suppressed",
+ ),
+ pytest.param(
+ "from airflow.sdk import DAG # noqa: F401, SDK001 - needed
for compat\n",
+ [],
+ id="combined-codes-with-explanation-suppressed",
+ ),
+ pytest.param(
+ "from airflow.sdk import DAG # noqa\n",
+ [(1, "from airflow.sdk import DAG")],
+ id="bare-noqa-not-suppressed",
+ ),
],
)
def test_nocheck_marker(self, tmp_path: Path, code: str, expected:
list[tuple[int, str]]):
diff --git a/task-sdk/src/airflow/sdk/execution_time/context.py
b/task-sdk/src/airflow/sdk/execution_time/context.py
index 1aaa68ed795..cdd25803989 100644
--- a/task-sdk/src/airflow/sdk/execution_time/context.py
+++ b/task-sdk/src/airflow/sdk/execution_time/context.py
@@ -1101,7 +1101,7 @@ def context_to_airflow_vars(context: Mapping[str, Any],
in_env_var_format: bool
"""
from datetime import datetime
- from airflow import settings
+ from airflow import settings # noqa: SDK002
params = {}
if in_env_var_format:
diff --git a/task-sdk/src/airflow/sdk/plugins_manager.py
b/task-sdk/src/airflow/sdk/plugins_manager.py
index 8419a7dbba9..dca92ea0be8 100644
--- a/task-sdk/src/airflow/sdk/plugins_manager.py
+++ b/task-sdk/src/airflow/sdk/plugins_manager.py
@@ -23,7 +23,7 @@ import logging
from functools import cache
from typing import TYPE_CHECKING
-from airflow import settings
+from airflow import settings # noqa: SDK002
from airflow.sdk._shared.module_loading import import_string
from airflow.sdk._shared.observability.metrics import stats
from airflow.sdk._shared.plugins_manager import (
