From 52810403be19cd17deed79e0f600730ff1ff3774 Mon Sep 17 00:00:00 2001 From: Kaxil Naik Date: Mon, 20 Oct 2025 01:09:47 +0100 Subject: [PATCH] Move Standard provider to use lazy_compat for compatibility imports MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Consolidate the Standard provider's compatibility logic by replacing imports from version_compat.py with direct imports from lazy_compat, providing a single source of truth for all Airflow 2.x ↔ 3.x compatibility. Follow-up of https://github.com/apache/airflow/pull/56790 & https://github.com/apache/airflow/pull/56793 --- .pre-commit-config.yaml | 4 +- .../airflow/providers/common/compat/sdk.py | 63 +++ .../airflow/providers/common/compat/sdk.pyi | 120 ----- providers/standard/pyproject.toml | 2 + .../providers/standard/decorators/bash.py | 20 +- .../decorators/branch_external_python.py | 10 +- .../standard/decorators/branch_python.py | 9 +- .../standard/decorators/branch_virtualenv.py | 9 +- .../standard/decorators/external_python.py | 9 +- .../providers/standard/decorators/python.py | 9 +- .../standard/decorators/python_virtualenv.py | 11 +- .../providers/standard/decorators/sensor.py | 11 +- .../standard/decorators/short_circuit.py | 10 +- .../providers/standard/decorators/stub.py | 18 +- .../example_dags/example_bash_decorator.py | 7 +- .../example_dags/example_branch_operator.py | 7 +- .../example_branch_operator_decorator.py | 7 +- ...example_external_task_parent_deferrable.py | 8 +- .../example_dags/example_hitl_operator.py | 2 +- .../standard/example_dags/example_sensors.py | 7 +- .../example_short_circuit_decorator.py | 7 +- .../example_short_circuit_operator.py | 7 +- .../providers/standard/hooks/filesystem.py | 2 +- .../providers/standard/hooks/package_index.py | 2 +- .../providers/standard/hooks/subprocess.py | 2 +- .../providers/standard/operators/bash.py | 15 +- .../providers/standard/operators/branch.py | 2 +- .../providers/standard/operators/datetime.py | 8 +- .../providers/standard/operators/empty.py | 2 +- .../providers/standard/operators/hitl.py | 2 +- .../standard/operators/latest_only.py | 7 +- .../providers/standard/operators/python.py | 9 +- .../providers/standard/operators/smooth.py | 2 +- .../standard/operators/trigger_dagrun.py | 20 +- .../providers/standard/operators/weekday.py | 8 +- .../providers/standard/sensors/bash.py | 8 +- .../providers/standard/sensors/date_time.py | 8 +- .../standard/sensors/external_task.py | 9 +- .../providers/standard/sensors/filesystem.py | 2 +- .../providers/standard/sensors/python.py | 8 +- .../providers/standard/sensors/time.py | 7 +- .../providers/standard/sensors/time_delta.py | 10 +- .../providers/standard/sensors/weekday.py | 9 +- .../providers/standard/triggers/temporal.py | 6 +- .../providers/standard/utils/skipmixin.py | 8 +- .../providers/standard/version_compat.py | 24 +- .../prek/check_common_compat_lazy_imports.py | 483 +++++------------- 47 files changed, 263 insertions(+), 757 deletions(-) delete mode 100644 providers/common/compat/src/airflow/providers/common/compat/sdk.pyi diff --git a/.pre-commit-config.yaml b/.pre-commit-config.yaml index ba96c586bac8a..8291a9ca6dec9 100644 --- a/.pre-commit-config.yaml +++ b/.pre-commit-config.yaml @@ -466,9 +466,9 @@ repos: --pattern '^openlineage\.client\.(facet|run)' --message "You should import from `airflow.providers.common.compat.openlineage.facet` instead." - id: check-common-compat-sdk-imports-in-sync - name: Check common.compat sdk.pyi is in sync + name: Check common.compat sdk TYPE_CHECKING matches runtime maps language: python - files: ^providers/common/compat/src/airflow/providers/common/compat/sdk\.(py|pyi)$ + files: ^providers/common/compat/src/airflow/providers/common/compat/sdk\.py$ pass_filenames: false entry: ./scripts/ci/prek/check_common_compat_lazy_imports.py - id: check-airflow-providers-bug-report-template diff --git a/providers/common/compat/src/airflow/providers/common/compat/sdk.py b/providers/common/compat/src/airflow/providers/common/compat/sdk.py index 5bde8567bd1be..3ccda1e1929d4 100644 --- a/providers/common/compat/src/airflow/providers/common/compat/sdk.py +++ b/providers/common/compat/src/airflow/providers/common/compat/sdk.py @@ -23,6 +23,64 @@ from __future__ import annotations +from typing import TYPE_CHECKING + +if TYPE_CHECKING: + import airflow.sdk.io as io # noqa: F401 + import airflow.sdk.timezone as timezone # noqa: F401 + from airflow.models.xcom import XCOM_RETURN_KEY as XCOM_RETURN_KEY + from airflow.sdk import ( + DAG as DAG, + Asset as Asset, + AssetAlias as AssetAlias, + AssetAll as AssetAll, + AssetAny as AssetAny, + BaseHook as BaseHook, + BaseNotifier as BaseNotifier, + BaseOperator as BaseOperator, + BaseOperatorLink as BaseOperatorLink, + BaseSensorOperator as BaseSensorOperator, + Connection as Connection, + Context as Context, + DagRunState as DagRunState, + EdgeModifier as EdgeModifier, + Label as Label, + Metadata as Metadata, + ObjectStoragePath as ObjectStoragePath, + Param as Param, + PokeReturnValue as PokeReturnValue, + TaskGroup as TaskGroup, + TaskInstanceState as TaskInstanceState, + TriggerRule as TriggerRule, + Variable as Variable, + WeightRule as WeightRule, + XComArg as XComArg, + chain as chain, + chain_linear as chain_linear, + cross_downstream as cross_downstream, + dag as dag, + get_current_context as get_current_context, + get_parsing_context as get_parsing_context, + setup as setup, + task as task, + task_group as task_group, + teardown as teardown, + ) + from airflow.sdk.bases.decorator import ( + DecoratedMappedOperator as DecoratedMappedOperator, + DecoratedOperator as DecoratedOperator, + TaskDecorator as TaskDecorator, + get_unique_task_id as get_unique_task_id, + task_decorator_factory as task_decorator_factory, + ) + from airflow.sdk.bases.sensor import poke_mode_only as poke_mode_only + from airflow.sdk.definitions.context import context_merge as context_merge + from airflow.sdk.definitions.mappedoperator import MappedOperator as MappedOperator + from airflow.sdk.definitions.template import literal as literal + from airflow.sdk.execution_time.context import context_to_airflow_vars as context_to_airflow_vars + from airflow.sdk.execution_time.timeout import timeout as timeout + from airflow.sdk.execution_time.xcom import XCom as XCom + from airflow.providers.common.compat._compat_utils import create_module_getattr # Rename map for classes that changed names between Airflow 2.x and 3.x @@ -63,6 +121,8 @@ "setup": ("airflow.sdk", "airflow.decorators"), "teardown": ("airflow.sdk", "airflow.decorators"), "TaskDecorator": ("airflow.sdk.bases.decorator", "airflow.decorators"), + "task_decorator_factory": ("airflow.sdk.bases.decorator", "airflow.decorators.base"), + "get_unique_task_id": ("airflow.sdk.bases.decorator", "airflow.decorators.base"), # ============================================================================ # Models # ============================================================================ @@ -74,6 +134,7 @@ "XComArg": ("airflow.sdk", "airflow.models.xcom_arg"), "DecoratedOperator": ("airflow.sdk.bases.decorator", "airflow.decorators.base"), "DecoratedMappedOperator": ("airflow.sdk.bases.decorator", "airflow.decorators.base"), + "MappedOperator": ("airflow.sdk.definitions.mappedoperator", "airflow.models.mappedoperator"), # ============================================================================ # Assets (Dataset → Asset rename in Airflow 3.0) # ============================================================================ @@ -119,6 +180,8 @@ # Context & Utilities # ============================================================================ "Context": ("airflow.sdk", "airflow.utils.context"), + "context_merge": ("airflow.sdk.definitions.context", "airflow.utils.context"), + "context_to_airflow_vars": ("airflow.sdk.execution_time.context", "airflow.utils.operator_helpers"), "get_current_context": ("airflow.sdk", "airflow.operators.python"), "get_parsing_context": ("airflow.sdk", "airflow.utils.dag_parsing_context"), # ============================================================================ diff --git a/providers/common/compat/src/airflow/providers/common/compat/sdk.pyi b/providers/common/compat/src/airflow/providers/common/compat/sdk.pyi deleted file mode 100644 index 99593f5d20c27..0000000000000 --- a/providers/common/compat/src/airflow/providers/common/compat/sdk.pyi +++ /dev/null @@ -1,120 +0,0 @@ -# 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. -""" -Type stubs for IDE autocomplete - always uses Airflow 3 paths. - -This file is auto-generated from sdk.py. - - run scripts/ci/prek/check_common_compat_lazy_imports.py --generate instead. -""" - -import airflow.sdk.io as io -import airflow.sdk.timezone as timezone -from airflow.models.xcom import XCOM_RETURN_KEY as XCOM_RETURN_KEY -from airflow.sdk import ( - DAG as DAG, - Asset as Asset, - AssetAlias as AssetAlias, - AssetAll as AssetAll, - AssetAny as AssetAny, - BaseHook as BaseHook, - BaseNotifier as BaseNotifier, - BaseOperator as BaseOperator, - BaseOperatorLink as BaseOperatorLink, - BaseSensorOperator as BaseSensorOperator, - Connection as Connection, - Context as Context, - DagRunState as DagRunState, - EdgeModifier as EdgeModifier, - Label as Label, - Metadata as Metadata, - ObjectStoragePath as ObjectStoragePath, - Param as Param, - PokeReturnValue as PokeReturnValue, - TaskGroup as TaskGroup, - TaskInstanceState as TaskInstanceState, - TriggerRule as TriggerRule, - Variable as Variable, - WeightRule as WeightRule, - XComArg as XComArg, - chain as chain, - chain_linear as chain_linear, - cross_downstream as cross_downstream, - dag as dag, - get_current_context as get_current_context, - get_parsing_context as get_parsing_context, - setup as setup, - task as task, - task_group as task_group, - teardown as teardown, -) -from airflow.sdk.bases.decorator import ( - DecoratedMappedOperator as DecoratedMappedOperator, - DecoratedOperator as DecoratedOperator, - TaskDecorator as TaskDecorator, -) -from airflow.sdk.bases.sensor import poke_mode_only as poke_mode_only -from airflow.sdk.definitions.template import literal as literal -from airflow.sdk.execution_time.timeout import timeout as timeout -from airflow.sdk.execution_time.xcom import XCom as XCom - -__all__: list[str] = [ - "Asset", - "AssetAlias", - "AssetAll", - "AssetAny", - "BaseHook", - "BaseNotifier", - "BaseOperator", - "BaseOperatorLink", - "BaseSensorOperator", - "Connection", - "Context", - "DAG", - "DagRunState", - "DecoratedMappedOperator", - "DecoratedOperator", - "EdgeModifier", - "Label", - "Metadata", - "ObjectStoragePath", - "Param", - "PokeReturnValue", - "TaskDecorator", - "TaskGroup", - "TaskInstanceState", - "TriggerRule", - "Variable", - "WeightRule", - "XCOM_RETURN_KEY", - "XCom", - "XComArg", - "chain", - "chain_linear", - "cross_downstream", - "dag", - "get_current_context", - "get_parsing_context", - "io", - "literal", - "poke_mode_only", - "setup", - "task", - "task_group", - "teardown", - "timeout", - "timezone", -] diff --git a/providers/standard/pyproject.toml b/providers/standard/pyproject.toml index ffe78afd6d65a..b8334ae4c2637 100644 --- a/providers/standard/pyproject.toml +++ b/providers/standard/pyproject.toml @@ -58,6 +58,7 @@ requires-python = ">=3.10" # After you modify the dependencies, and rebuild your Breeze CI image with ``breeze ci-image build`` dependencies = [ "apache-airflow>=2.10.0", + "apache-airflow-providers-common-compat>=1.7.4", # + TODO: bump to next version ] [dependency-groups] @@ -65,6 +66,7 @@ dev = [ "apache-airflow", "apache-airflow-task-sdk", "apache-airflow-devel-common", + "apache-airflow-providers-common-compat", # Additional devel dependencies (do not remove this line and add extra development dependencies) "apache-airflow-providers-mysql", ] diff --git a/providers/standard/src/airflow/providers/standard/decorators/bash.py b/providers/standard/src/airflow/providers/standard/decorators/bash.py index 71962b0da13be..b53c8b0242d96 100644 --- a/providers/standard/src/airflow/providers/standard/decorators/bash.py +++ b/providers/standard/src/airflow/providers/standard/decorators/bash.py @@ -21,24 +21,18 @@ from collections.abc import Callable, Collection, Mapping, Sequence from typing import TYPE_CHECKING, Any, ClassVar -from airflow.providers.standard.version_compat import AIRFLOW_V_3_0_PLUS - -if AIRFLOW_V_3_0_PLUS: - from airflow.sdk.bases.decorator import DecoratedOperator, TaskDecorator, task_decorator_factory -else: - from airflow.decorators.base import ( # type: ignore[no-redef] - DecoratedOperator, - TaskDecorator, - task_decorator_factory, - ) - +from airflow.providers.common.compat.sdk import ( + DecoratedOperator, + TaskDecorator, + context_merge, + task_decorator_factory, +) from airflow.providers.standard.operators.bash import BashOperator -from airflow.providers.standard.version_compat import context_merge from airflow.sdk.definitions._internal.types import SET_DURING_EXECUTION from airflow.utils.operator_helpers import determine_kwargs if TYPE_CHECKING: - from airflow.sdk.definitions.context import Context + from airflow.providers.common.compat.sdk import Context class _BashDecoratedOperator(DecoratedOperator, BashOperator): diff --git a/providers/standard/src/airflow/providers/standard/decorators/branch_external_python.py b/providers/standard/src/airflow/providers/standard/decorators/branch_external_python.py index 3241d1bbded59..c724e4840b5bc 100644 --- a/providers/standard/src/airflow/providers/standard/decorators/branch_external_python.py +++ b/providers/standard/src/airflow/providers/standard/decorators/branch_external_python.py @@ -19,18 +19,12 @@ from collections.abc import Callable from typing import TYPE_CHECKING -from airflow.providers.standard.version_compat import AIRFLOW_V_3_0_PLUS - -if AIRFLOW_V_3_0_PLUS: - from airflow.sdk.bases.decorator import task_decorator_factory -else: - from airflow.decorators.base import task_decorator_factory # type: ignore[no-redef] - +from airflow.providers.common.compat.sdk import task_decorator_factory from airflow.providers.standard.decorators.python import _PythonDecoratedOperator from airflow.providers.standard.operators.python import BranchExternalPythonOperator if TYPE_CHECKING: - from airflow.sdk.bases.decorator import TaskDecorator + from airflow.providers.common.compat.sdk import TaskDecorator class _BranchExternalPythonDecoratedOperator(_PythonDecoratedOperator, BranchExternalPythonOperator): diff --git a/providers/standard/src/airflow/providers/standard/decorators/branch_python.py b/providers/standard/src/airflow/providers/standard/decorators/branch_python.py index 719037ac054e5..646b9a67a04d7 100644 --- a/providers/standard/src/airflow/providers/standard/decorators/branch_python.py +++ b/providers/standard/src/airflow/providers/standard/decorators/branch_python.py @@ -19,17 +19,12 @@ from collections.abc import Callable from typing import TYPE_CHECKING -from airflow.providers.standard.version_compat import AIRFLOW_V_3_0_PLUS - -if AIRFLOW_V_3_0_PLUS: - from airflow.sdk.bases.decorator import task_decorator_factory -else: - from airflow.decorators.base import task_decorator_factory # type: ignore[no-redef] +from airflow.providers.common.compat.sdk import task_decorator_factory from airflow.providers.standard.decorators.python import _PythonDecoratedOperator from airflow.providers.standard.operators.python import BranchPythonOperator if TYPE_CHECKING: - from airflow.sdk.bases.decorator import TaskDecorator + from airflow.providers.common.compat.sdk import TaskDecorator class _BranchPythonDecoratedOperator(_PythonDecoratedOperator, BranchPythonOperator): diff --git a/providers/standard/src/airflow/providers/standard/decorators/branch_virtualenv.py b/providers/standard/src/airflow/providers/standard/decorators/branch_virtualenv.py index f138f90bda25e..4cd52ea6140d1 100644 --- a/providers/standard/src/airflow/providers/standard/decorators/branch_virtualenv.py +++ b/providers/standard/src/airflow/providers/standard/decorators/branch_virtualenv.py @@ -19,17 +19,12 @@ from collections.abc import Callable from typing import TYPE_CHECKING -from airflow.providers.standard.version_compat import AIRFLOW_V_3_0_PLUS - -if AIRFLOW_V_3_0_PLUS: - from airflow.sdk.bases.decorator import task_decorator_factory -else: - from airflow.decorators.base import task_decorator_factory # type: ignore[no-redef] +from airflow.providers.common.compat.sdk import task_decorator_factory from airflow.providers.standard.decorators.python import _PythonDecoratedOperator from airflow.providers.standard.operators.python import BranchPythonVirtualenvOperator if TYPE_CHECKING: - from airflow.sdk.bases.decorator import TaskDecorator + from airflow.providers.common.compat.sdk import TaskDecorator class _BranchPythonVirtualenvDecoratedOperator(_PythonDecoratedOperator, BranchPythonVirtualenvOperator): diff --git a/providers/standard/src/airflow/providers/standard/decorators/external_python.py b/providers/standard/src/airflow/providers/standard/decorators/external_python.py index fa399cdb1ee4d..c3a3a5efff18e 100644 --- a/providers/standard/src/airflow/providers/standard/decorators/external_python.py +++ b/providers/standard/src/airflow/providers/standard/decorators/external_python.py @@ -19,17 +19,12 @@ from collections.abc import Callable from typing import TYPE_CHECKING -from airflow.providers.standard.version_compat import AIRFLOW_V_3_0_PLUS - -if AIRFLOW_V_3_0_PLUS: - from airflow.sdk.bases.decorator import task_decorator_factory -else: - from airflow.decorators.base import task_decorator_factory # type: ignore[no-redef] +from airflow.providers.common.compat.sdk import task_decorator_factory from airflow.providers.standard.decorators.python import _PythonDecoratedOperator from airflow.providers.standard.operators.python import ExternalPythonOperator if TYPE_CHECKING: - from airflow.sdk.bases.decorator import TaskDecorator + from airflow.providers.common.compat.sdk import TaskDecorator class _PythonExternalDecoratedOperator(_PythonDecoratedOperator, ExternalPythonOperator): diff --git a/providers/standard/src/airflow/providers/standard/decorators/python.py b/providers/standard/src/airflow/providers/standard/decorators/python.py index 73dab343d719e..91e6df722d940 100644 --- a/providers/standard/src/airflow/providers/standard/decorators/python.py +++ b/providers/standard/src/airflow/providers/standard/decorators/python.py @@ -19,16 +19,11 @@ from collections.abc import Callable, Sequence from typing import TYPE_CHECKING +from airflow.providers.common.compat.sdk import DecoratedOperator, task_decorator_factory from airflow.providers.standard.operators.python import PythonOperator -from airflow.providers.standard.version_compat import AIRFLOW_V_3_0_PLUS - -if AIRFLOW_V_3_0_PLUS: - from airflow.sdk.bases.decorator import DecoratedOperator, task_decorator_factory -else: - from airflow.decorators.base import DecoratedOperator, task_decorator_factory # type: ignore[no-redef] if TYPE_CHECKING: - from airflow.sdk.bases.decorator import TaskDecorator + from airflow.providers.common.compat.sdk import TaskDecorator class _PythonDecoratedOperator(DecoratedOperator, PythonOperator): diff --git a/providers/standard/src/airflow/providers/standard/decorators/python_virtualenv.py b/providers/standard/src/airflow/providers/standard/decorators/python_virtualenv.py index 91ef7f03c0508..4683c1220a328 100644 --- a/providers/standard/src/airflow/providers/standard/decorators/python_virtualenv.py +++ b/providers/standard/src/airflow/providers/standard/decorators/python_virtualenv.py @@ -19,19 +19,12 @@ from collections.abc import Callable from typing import TYPE_CHECKING -from airflow.providers.standard.version_compat import AIRFLOW_V_3_0_PLUS - -if AIRFLOW_V_3_0_PLUS: - from airflow.sdk.bases.decorator import task_decorator_factory -else: - from airflow.decorators.base import task_decorator_factory # type: ignore[no-redef] - - +from airflow.providers.common.compat.sdk import task_decorator_factory from airflow.providers.standard.decorators.python import _PythonDecoratedOperator from airflow.providers.standard.operators.python import PythonVirtualenvOperator if TYPE_CHECKING: - from airflow.sdk.bases.decorator import TaskDecorator + from airflow.providers.common.compat.sdk import TaskDecorator class _PythonVirtualenvDecoratedOperator(_PythonDecoratedOperator, PythonVirtualenvOperator): diff --git a/providers/standard/src/airflow/providers/standard/decorators/sensor.py b/providers/standard/src/airflow/providers/standard/decorators/sensor.py index 5310d30695ed8..48fabd5d6c5c4 100644 --- a/providers/standard/src/airflow/providers/standard/decorators/sensor.py +++ b/providers/standard/src/airflow/providers/standard/decorators/sensor.py @@ -20,18 +20,11 @@ from collections.abc import Callable, Sequence from typing import TYPE_CHECKING, ClassVar -from airflow.providers.standard.version_compat import AIRFLOW_V_3_0_PLUS - -if AIRFLOW_V_3_0_PLUS: - from airflow.sdk.bases.decorator import get_unique_task_id, task_decorator_factory -else: - from airflow.decorators.base import get_unique_task_id, task_decorator_factory # type: ignore[no-redef] - - +from airflow.providers.common.compat.sdk import get_unique_task_id, task_decorator_factory from airflow.providers.standard.sensors.python import PythonSensor if TYPE_CHECKING: - from airflow.sdk.bases.decorator import TaskDecorator + from airflow.providers.common.compat.sdk import TaskDecorator class DecoratedSensorOperator(PythonSensor): diff --git a/providers/standard/src/airflow/providers/standard/decorators/short_circuit.py b/providers/standard/src/airflow/providers/standard/decorators/short_circuit.py index 148a6ed828fdc..2cf838534c172 100644 --- a/providers/standard/src/airflow/providers/standard/decorators/short_circuit.py +++ b/providers/standard/src/airflow/providers/standard/decorators/short_circuit.py @@ -19,18 +19,12 @@ from collections.abc import Callable from typing import TYPE_CHECKING -from airflow.providers.standard.version_compat import AIRFLOW_V_3_0_PLUS - -if AIRFLOW_V_3_0_PLUS: - from airflow.sdk.bases.decorator import task_decorator_factory -else: - from airflow.decorators.base import task_decorator_factory # type: ignore[no-redef] - +from airflow.providers.common.compat.sdk import task_decorator_factory from airflow.providers.standard.decorators.python import _PythonDecoratedOperator from airflow.providers.standard.operators.python import ShortCircuitOperator if TYPE_CHECKING: - from airflow.sdk.bases.decorator import TaskDecorator + from airflow.providers.common.compat.sdk import TaskDecorator class _ShortCircuitDecoratedOperator(_PythonDecoratedOperator, ShortCircuitOperator): diff --git a/providers/standard/src/airflow/providers/standard/decorators/stub.py b/providers/standard/src/airflow/providers/standard/decorators/stub.py index c771dcdb436c9..f29d123c740c1 100644 --- a/providers/standard/src/airflow/providers/standard/decorators/stub.py +++ b/providers/standard/src/airflow/providers/standard/decorators/stub.py @@ -21,20 +21,14 @@ from collections.abc import Callable from typing import TYPE_CHECKING, Any -if TYPE_CHECKING: - from airflow.sdk.bases.decorator import DecoratedOperator, TaskDecorator, task_decorator_factory -else: - try: - from airflow.sdk.bases.decorator import DecoratedOperator, TaskDecorator, task_decorator_factory - except ModuleNotFoundError: - from airflow.decorators.base import ( - DecoratedOperator, - TaskDecorator, - task_decorator_factory, - ) +from airflow.providers.common.compat.sdk import ( + DecoratedOperator, + TaskDecorator, + task_decorator_factory, +) if TYPE_CHECKING: - from airflow.sdk.definitions.context import Context + from airflow.providers.common.compat.sdk import Context class _StubOperator(DecoratedOperator): diff --git a/providers/standard/src/airflow/providers/standard/example_dags/example_bash_decorator.py b/providers/standard/src/airflow/providers/standard/example_dags/example_bash_decorator.py index dfe49c5221974..20225be4b39fd 100644 --- a/providers/standard/src/airflow/providers/standard/example_dags/example_bash_decorator.py +++ b/providers/standard/src/airflow/providers/standard/example_dags/example_bash_decorator.py @@ -20,16 +20,11 @@ import pendulum from airflow.exceptions import AirflowSkipException +from airflow.providers.common.compat.sdk import TriggerRule from airflow.providers.standard.operators.empty import EmptyOperator from airflow.providers.standard.utils.weekday import WeekDay from airflow.sdk import chain, dag, task -try: - from airflow.sdk import TriggerRule -except ImportError: - # Compatibility for Airflow < 3.1 - from airflow.utils.trigger_rule import TriggerRule # type: ignore[no-redef,attr-defined] - @dag(schedule=None, start_date=pendulum.datetime(2023, 1, 1, tz="UTC"), catchup=False) def example_bash_decorator(): diff --git a/providers/standard/src/airflow/providers/standard/example_dags/example_branch_operator.py b/providers/standard/src/airflow/providers/standard/example_dags/example_branch_operator.py index c3cf4856ba565..6a7fec26dd7a1 100644 --- a/providers/standard/src/airflow/providers/standard/example_dags/example_branch_operator.py +++ b/providers/standard/src/airflow/providers/standard/example_dags/example_branch_operator.py @@ -29,6 +29,7 @@ import pendulum +from airflow.providers.common.compat.sdk import TriggerRule from airflow.providers.standard.operators.empty import EmptyOperator from airflow.providers.standard.operators.python import ( BranchExternalPythonOperator, @@ -40,12 +41,6 @@ ) from airflow.sdk import DAG, Label -try: - from airflow.sdk import TriggerRule -except ImportError: - # Compatibility for Airflow < 3.1 - from airflow.utils.trigger_rule import TriggerRule # type: ignore[no-redef,attr-defined] - PATH_TO_PYTHON_BINARY = sys.executable with DAG( diff --git a/providers/standard/src/airflow/providers/standard/example_dags/example_branch_operator_decorator.py b/providers/standard/src/airflow/providers/standard/example_dags/example_branch_operator_decorator.py index 2419138aa9514..3c8c2ba16057f 100644 --- a/providers/standard/src/airflow/providers/standard/example_dags/example_branch_operator_decorator.py +++ b/providers/standard/src/airflow/providers/standard/example_dags/example_branch_operator_decorator.py @@ -30,15 +30,10 @@ import pendulum +from airflow.providers.common.compat.sdk import TriggerRule from airflow.providers.standard.operators.empty import EmptyOperator from airflow.sdk import DAG, Label, task -try: - from airflow.sdk import TriggerRule -except ImportError: - # Compatibility for Airflow < 3.1 - from airflow.utils.trigger_rule import TriggerRule # type: ignore[no-redef,attr-defined] - PATH_TO_PYTHON_BINARY = sys.executable with DAG( diff --git a/providers/standard/src/airflow/providers/standard/example_dags/example_external_task_parent_deferrable.py b/providers/standard/src/airflow/providers/standard/example_dags/example_external_task_parent_deferrable.py index 5b55fe566d722..1e1f218bab398 100644 --- a/providers/standard/src/airflow/providers/standard/example_dags/example_external_task_parent_deferrable.py +++ b/providers/standard/src/airflow/providers/standard/example_dags/example_external_task_parent_deferrable.py @@ -17,18 +17,14 @@ from __future__ import annotations from airflow import DAG +from airflow.providers.common.compat.sdk import timezone from airflow.providers.standard.operators.empty import EmptyOperator from airflow.providers.standard.operators.trigger_dagrun import TriggerDagRunOperator from airflow.providers.standard.sensors.external_task import ExternalTaskSensor -try: - from airflow.sdk.timezone import datetime -except ImportError: - from airflow.utils.timezone import datetime # type: ignore[no-redef] - with DAG( dag_id="example_external_task", - start_date=datetime(2022, 1, 1), + start_date=timezone.datetime(2022, 1, 1), schedule="@once", catchup=False, tags=["example", "async", "core"], diff --git a/providers/standard/src/airflow/providers/standard/example_dags/example_hitl_operator.py b/providers/standard/src/airflow/providers/standard/example_dags/example_hitl_operator.py index 44d77ba4959a2..fd93eccd84189 100644 --- a/providers/standard/src/airflow/providers/standard/example_dags/example_hitl_operator.py +++ b/providers/standard/src/airflow/providers/standard/example_dags/example_hitl_operator.py @@ -32,7 +32,7 @@ from airflow.sdk.bases.notifier import BaseNotifier if TYPE_CHECKING: - from airflow.sdk.definitions.context import Context + from airflow.providers.common.compat.sdk import Context # [START hitl_tutorial] diff --git a/providers/standard/src/airflow/providers/standard/example_dags/example_sensors.py b/providers/standard/src/airflow/providers/standard/example_dags/example_sensors.py index 50b8642b34173..73ccaadbe0f02 100644 --- a/providers/standard/src/airflow/providers/standard/example_dags/example_sensors.py +++ b/providers/standard/src/airflow/providers/standard/example_dags/example_sensors.py @@ -21,6 +21,7 @@ import pendulum +from airflow.providers.common.compat.sdk import TriggerRule from airflow.providers.standard.operators.bash import BashOperator from airflow.providers.standard.sensors.bash import BashSensor from airflow.providers.standard.sensors.filesystem import FileSensor @@ -31,12 +32,6 @@ from airflow.providers.standard.utils.weekday import WeekDay from airflow.sdk import DAG -try: - from airflow.sdk import TriggerRule -except ImportError: - # Compatibility for Airflow < 3.1 - from airflow.utils.trigger_rule import TriggerRule # type: ignore[no-redef,attr-defined] - # [START example_callables] def success_callable(): diff --git a/providers/standard/src/airflow/providers/standard/example_dags/example_short_circuit_decorator.py b/providers/standard/src/airflow/providers/standard/example_dags/example_short_circuit_decorator.py index 4bb248fb94ce0..598778edfd82e 100644 --- a/providers/standard/src/airflow/providers/standard/example_dags/example_short_circuit_decorator.py +++ b/providers/standard/src/airflow/providers/standard/example_dags/example_short_circuit_decorator.py @@ -20,15 +20,10 @@ import pendulum +from airflow.providers.common.compat.sdk import TriggerRule from airflow.providers.standard.operators.empty import EmptyOperator from airflow.sdk import chain, dag, task -try: - from airflow.sdk import TriggerRule -except ImportError: - # Compatibility for Airflow < 3.1 - from airflow.utils.trigger_rule import TriggerRule # type: ignore[no-redef,attr-defined] - @dag(schedule=None, start_date=pendulum.datetime(2021, 1, 1, tz="UTC"), catchup=False, tags=["example"]) def example_short_circuit_decorator(): diff --git a/providers/standard/src/airflow/providers/standard/example_dags/example_short_circuit_operator.py b/providers/standard/src/airflow/providers/standard/example_dags/example_short_circuit_operator.py index b054d9ec6dd51..9d028f21be0ff 100644 --- a/providers/standard/src/airflow/providers/standard/example_dags/example_short_circuit_operator.py +++ b/providers/standard/src/airflow/providers/standard/example_dags/example_short_circuit_operator.py @@ -21,16 +21,11 @@ import pendulum +from airflow.providers.common.compat.sdk import TriggerRule from airflow.providers.standard.operators.empty import EmptyOperator from airflow.providers.standard.operators.python import ShortCircuitOperator from airflow.sdk import DAG, chain -try: - from airflow.sdk import TriggerRule -except ImportError: - # Compatibility for Airflow < 3.1 - from airflow.utils.trigger_rule import TriggerRule # type: ignore[no-redef,attr-defined] - with DAG( dag_id="example_short_circuit_operator", schedule=None, diff --git a/providers/standard/src/airflow/providers/standard/hooks/filesystem.py b/providers/standard/src/airflow/providers/standard/hooks/filesystem.py index 090ffc4789035..689b1b5e3b269 100644 --- a/providers/standard/src/airflow/providers/standard/hooks/filesystem.py +++ b/providers/standard/src/airflow/providers/standard/hooks/filesystem.py @@ -20,7 +20,7 @@ from pathlib import Path from typing import Any -from airflow.providers.standard.version_compat import BaseHook +from airflow.providers.common.compat.sdk import BaseHook class FSHook(BaseHook): diff --git a/providers/standard/src/airflow/providers/standard/hooks/package_index.py b/providers/standard/src/airflow/providers/standard/hooks/package_index.py index da2dc17c7afb8..e7f3c63f1c975 100644 --- a/providers/standard/src/airflow/providers/standard/hooks/package_index.py +++ b/providers/standard/src/airflow/providers/standard/hooks/package_index.py @@ -23,7 +23,7 @@ from typing import Any from urllib.parse import quote, urlparse -from airflow.providers.standard.version_compat import BaseHook +from airflow.providers.common.compat.sdk import BaseHook class PackageIndexHook(BaseHook): diff --git a/providers/standard/src/airflow/providers/standard/hooks/subprocess.py b/providers/standard/src/airflow/providers/standard/hooks/subprocess.py index aba2c3ede2e26..601fe6815554e 100644 --- a/providers/standard/src/airflow/providers/standard/hooks/subprocess.py +++ b/providers/standard/src/airflow/providers/standard/hooks/subprocess.py @@ -24,7 +24,7 @@ from subprocess import PIPE, STDOUT, Popen from tempfile import TemporaryDirectory, gettempdir -from airflow.providers.standard.version_compat import BaseHook +from airflow.providers.common.compat.sdk import BaseHook SubprocessResult = namedtuple("SubprocessResult", ["exit_code", "output"]) diff --git a/providers/standard/src/airflow/providers/standard/operators/bash.py b/providers/standard/src/airflow/providers/standard/operators/bash.py index e10192a05b228..533b1dee5fbb1 100644 --- a/providers/standard/src/airflow/providers/standard/operators/bash.py +++ b/providers/standard/src/airflow/providers/standard/operators/bash.py @@ -25,23 +25,14 @@ from typing import TYPE_CHECKING, Any, cast from airflow.exceptions import AirflowException, AirflowSkipException +from airflow.providers.common.compat.sdk import context_to_airflow_vars from airflow.providers.standard.hooks.subprocess import SubprocessHook, SubprocessResult, working_directory -from airflow.providers.standard.version_compat import AIRFLOW_V_3_0_PLUS, BaseOperator - -if AIRFLOW_V_3_0_PLUS: - from airflow.sdk.execution_time.context import context_to_airflow_vars -else: - from airflow.utils.operator_helpers import context_to_airflow_vars # type: ignore[no-redef, attr-defined] +from airflow.providers.standard.version_compat import BaseOperator if TYPE_CHECKING: + from airflow.providers.common.compat.sdk import Context from airflow.utils.types import ArgNotSet - try: - from airflow.sdk.definitions.context import Context - except ImportError: - # TODO: Remove once provider drops support for Airflow 2 - from airflow.utils.context import Context - class BashOperator(BaseOperator): r""" diff --git a/providers/standard/src/airflow/providers/standard/operators/branch.py b/providers/standard/src/airflow/providers/standard/operators/branch.py index cc405b365c902..64796bd0f7a36 100644 --- a/providers/standard/src/airflow/providers/standard/operators/branch.py +++ b/providers/standard/src/airflow/providers/standard/operators/branch.py @@ -30,7 +30,7 @@ from airflow.models.skipmixin import SkipMixin if TYPE_CHECKING: - from airflow.sdk.definitions.context import Context + from airflow.providers.common.compat.sdk import Context from airflow.sdk.types import RuntimeTaskInstanceProtocol diff --git a/providers/standard/src/airflow/providers/standard/operators/datetime.py b/providers/standard/src/airflow/providers/standard/operators/datetime.py index a549b13354ced..5613951cb4177 100644 --- a/providers/standard/src/airflow/providers/standard/operators/datetime.py +++ b/providers/standard/src/airflow/providers/standard/operators/datetime.py @@ -21,15 +21,11 @@ from typing import TYPE_CHECKING from airflow.exceptions import AirflowException +from airflow.providers.common.compat.sdk import timezone from airflow.providers.standard.operators.branch import BaseBranchOperator -from airflow.utils import timezone if TYPE_CHECKING: - try: - from airflow.sdk.definitions.context import Context - except ImportError: - # TODO: Remove once provider drops support for Airflow 2 - from airflow.utils.context import Context + from airflow.providers.common.compat.sdk import Context class BranchDateTimeOperator(BaseBranchOperator): diff --git a/providers/standard/src/airflow/providers/standard/operators/empty.py b/providers/standard/src/airflow/providers/standard/operators/empty.py index b9a480bc86ea9..f1d877b4c9d1a 100644 --- a/providers/standard/src/airflow/providers/standard/operators/empty.py +++ b/providers/standard/src/airflow/providers/standard/operators/empty.py @@ -21,7 +21,7 @@ from airflow.providers.standard.version_compat import BaseOperator if TYPE_CHECKING: - from airflow.sdk.definitions.context import Context + from airflow.providers.common.compat.sdk import Context class EmptyOperator(BaseOperator): diff --git a/providers/standard/src/airflow/providers/standard/operators/hitl.py b/providers/standard/src/airflow/providers/standard/operators/hitl.py index f76d62d8a16a1..5a6b366957932 100644 --- a/providers/standard/src/airflow/providers/standard/operators/hitl.py +++ b/providers/standard/src/airflow/providers/standard/operators/hitl.py @@ -40,7 +40,7 @@ from airflow.sdk.timezone import utcnow if TYPE_CHECKING: - from airflow.sdk.definitions.context import Context + from airflow.providers.common.compat.sdk import Context from airflow.sdk.execution_time.hitl import HITLUser from airflow.sdk.types import RuntimeTaskInstanceProtocol diff --git a/providers/standard/src/airflow/providers/standard/operators/latest_only.py b/providers/standard/src/airflow/providers/standard/operators/latest_only.py index cf8d9038128b7..d2263ff66bc20 100644 --- a/providers/standard/src/airflow/providers/standard/operators/latest_only.py +++ b/providers/standard/src/airflow/providers/standard/operators/latest_only.py @@ -33,12 +33,7 @@ from pendulum.datetime import DateTime from airflow.models import DagRun - - try: - from airflow.sdk.definitions.context import Context - except ImportError: - # TODO: Remove once provider drops support for Airflow 2 - from airflow.utils.context import Context + from airflow.providers.common.compat.sdk import Context class LatestOnlyOperator(BaseBranchOperator): diff --git a/providers/standard/src/airflow/providers/standard/operators/python.py b/providers/standard/src/airflow/providers/standard/operators/python.py index da6214837b9ff..8436592e38864 100644 --- a/providers/standard/src/airflow/providers/standard/operators/python.py +++ b/providers/standard/src/airflow/providers/standard/operators/python.py @@ -49,9 +49,10 @@ DeserializingResultError, ) from airflow.models.variable import Variable +from airflow.providers.common.compat.sdk import context_merge from airflow.providers.standard.hooks.package_index import PackageIndexHook from airflow.providers.standard.utils.python_virtualenv import prepare_virtualenv, write_python_script -from airflow.providers.standard.version_compat import AIRFLOW_V_3_0_PLUS, BaseOperator, context_merge +from airflow.providers.standard.version_compat import AIRFLOW_V_3_0_PLUS, BaseOperator from airflow.utils import hashlib_wrapper from airflow.utils.file import get_unique_dag_module_name from airflow.utils.operator_helpers import KeywordParameters @@ -72,14 +73,10 @@ from pendulum.datetime import DateTime + from airflow.providers.common.compat.sdk import Context from airflow.sdk.execution_time.callback_runner import ExecutionCallableRunner from airflow.sdk.execution_time.context import OutletEventAccessorsProtocol - try: - from airflow.sdk.definitions.context import Context - except ImportError: # TODO: Remove once provider drops support for Airflow 2 - from airflow.utils.context import Context - _SerializerTypeDef = Literal["pickle", "cloudpickle", "dill"] diff --git a/providers/standard/src/airflow/providers/standard/operators/smooth.py b/providers/standard/src/airflow/providers/standard/operators/smooth.py index 9b3d1ff181e26..a6fd075710366 100644 --- a/providers/standard/src/airflow/providers/standard/operators/smooth.py +++ b/providers/standard/src/airflow/providers/standard/operators/smooth.py @@ -22,7 +22,7 @@ from airflow.providers.standard.version_compat import BaseOperator if TYPE_CHECKING: - from airflow.sdk.definitions.context import Context + from airflow.providers.common.compat.sdk import Context class SmoothOperator(BaseOperator): diff --git a/providers/standard/src/airflow/providers/standard/operators/trigger_dagrun.py b/providers/standard/src/airflow/providers/standard/operators/trigger_dagrun.py index 7fadc968057d3..b70841525303b 100644 --- a/providers/standard/src/airflow/providers/standard/operators/trigger_dagrun.py +++ b/providers/standard/src/airflow/providers/standard/operators/trigger_dagrun.py @@ -37,13 +37,9 @@ from airflow.models.dag import DagModel from airflow.models.dagrun import DagRun from airflow.models.serialized_dag import SerializedDagModel +from airflow.providers.common.compat.sdk import BaseOperatorLink, XCom, timezone from airflow.providers.standard.triggers.external_task import DagStateTrigger -from airflow.providers.standard.version_compat import ( - AIRFLOW_V_3_0_PLUS, - BaseOperator, - BaseOperatorLink, - timezone, -) +from airflow.providers.standard.version_compat import AIRFLOW_V_3_0_PLUS, BaseOperator from airflow.utils.state import DagRunState from airflow.utils.types import NOTSET, ArgNotSet, DagRunType @@ -55,17 +51,7 @@ from sqlalchemy.orm.session import Session from airflow.models.taskinstancekey import TaskInstanceKey - - try: - from airflow.sdk.definitions.context import Context - except ImportError: - # TODO: Remove once provider drops support for Airflow 2 - from airflow.utils.context import Context - -if AIRFLOW_V_3_0_PLUS: - from airflow.sdk.execution_time.xcom import XCom -else: - from airflow.models import XCom + from airflow.providers.common.compat.sdk import Context class DagIsPaused(AirflowException): diff --git a/providers/standard/src/airflow/providers/standard/operators/weekday.py b/providers/standard/src/airflow/providers/standard/operators/weekday.py index 6d42ff2f3e0a0..68c2c7cbe2a81 100644 --- a/providers/standard/src/airflow/providers/standard/operators/weekday.py +++ b/providers/standard/src/airflow/providers/standard/operators/weekday.py @@ -20,16 +20,12 @@ from collections.abc import Iterable from typing import TYPE_CHECKING +from airflow.providers.common.compat.sdk import timezone from airflow.providers.standard.operators.branch import BaseBranchOperator from airflow.providers.standard.utils.weekday import WeekDay -from airflow.utils import timezone if TYPE_CHECKING: - try: - from airflow.sdk.definitions.context import Context - except ImportError: - # TODO: Remove once provider drops support for Airflow 2 - from airflow.utils.context import Context + from airflow.providers.common.compat.sdk import Context class BranchDayOfWeekOperator(BaseBranchOperator): diff --git a/providers/standard/src/airflow/providers/standard/sensors/bash.py b/providers/standard/src/airflow/providers/standard/sensors/bash.py index 9b83b4af187a3..0294e841992a2 100644 --- a/providers/standard/src/airflow/providers/standard/sensors/bash.py +++ b/providers/standard/src/airflow/providers/standard/sensors/bash.py @@ -24,14 +24,10 @@ from typing import TYPE_CHECKING from airflow.exceptions import AirflowFailException -from airflow.providers.standard.version_compat import BaseSensorOperator +from airflow.providers.common.compat.sdk import BaseSensorOperator if TYPE_CHECKING: - try: - from airflow.sdk.definitions.context import Context - except ImportError: - # TODO: Remove once provider drops support for Airflow 2 - from airflow.utils.context import Context + from airflow.providers.common.compat.sdk import Context class BashSensor(BaseSensorOperator): diff --git a/providers/standard/src/airflow/providers/standard/sensors/date_time.py b/providers/standard/src/airflow/providers/standard/sensors/date_time.py index eae51eaacff2e..5805a7a9df924 100644 --- a/providers/standard/src/airflow/providers/standard/sensors/date_time.py +++ b/providers/standard/src/airflow/providers/standard/sensors/date_time.py @@ -22,13 +22,9 @@ from dataclasses import dataclass from typing import TYPE_CHECKING, Any, NoReturn +from airflow.providers.common.compat.sdk import BaseSensorOperator, timezone from airflow.providers.standard.triggers.temporal import DateTimeTrigger -from airflow.providers.standard.version_compat import AIRFLOW_V_3_0_PLUS, BaseSensorOperator - -try: - from airflow.sdk import timezone -except ImportError: # TODO: Remove this when min airflow version is 3.1.0 for standard provider - from airflow.utils import timezone # type: ignore[attr-defined,no-redef] +from airflow.providers.standard.version_compat import AIRFLOW_V_3_0_PLUS try: from airflow.triggers.base import StartTriggerArgs # type: ignore[no-redef] diff --git a/providers/standard/src/airflow/providers/standard/sensors/external_task.py b/providers/standard/src/airflow/providers/standard/sensors/external_task.py index 4a8cbc03f30f5..ee809caffc1a0 100644 --- a/providers/standard/src/airflow/providers/standard/sensors/external_task.py +++ b/providers/standard/src/airflow/providers/standard/sensors/external_task.py @@ -25,6 +25,7 @@ from airflow.configuration import conf from airflow.exceptions import AirflowSkipException from airflow.models.dag import DagModel +from airflow.providers.common.compat.sdk import BaseOperatorLink, BaseSensorOperator from airflow.providers.standard.exceptions import ( DuplicateStateError, ExternalDagDeletedError, @@ -42,8 +43,6 @@ AIRFLOW_V_3_0_PLUS, AIRFLOW_V_3_2_PLUS, BaseOperator, - BaseOperatorLink, - BaseSensorOperator, ) from airflow.utils.file import correct_maybe_zipped from airflow.utils.state import State, TaskInstanceState @@ -60,11 +59,7 @@ from sqlalchemy.orm import Session from airflow.models.taskinstancekey import TaskInstanceKey - - if AIRFLOW_V_3_0_PLUS: - from airflow.sdk.definitions.context import Context - else: - from airflow.utils.context import Context + from airflow.providers.common.compat.sdk import Context class ExternalDagLink(BaseOperatorLink): diff --git a/providers/standard/src/airflow/providers/standard/sensors/filesystem.py b/providers/standard/src/airflow/providers/standard/sensors/filesystem.py index d23906c3404a2..b5c4078bb182f 100644 --- a/providers/standard/src/airflow/providers/standard/sensors/filesystem.py +++ b/providers/standard/src/airflow/providers/standard/sensors/filesystem.py @@ -27,9 +27,9 @@ from airflow.configuration import conf from airflow.exceptions import AirflowException +from airflow.providers.common.compat.sdk import BaseSensorOperator from airflow.providers.standard.hooks.filesystem import FSHook from airflow.providers.standard.triggers.file import FileTrigger -from airflow.providers.standard.version_compat import BaseSensorOperator try: from airflow.triggers.base import StartTriggerArgs # type: ignore[no-redef] diff --git a/providers/standard/src/airflow/providers/standard/sensors/python.py b/providers/standard/src/airflow/providers/standard/sensors/python.py index a8bcce7017ca3..97fa865b9889f 100644 --- a/providers/standard/src/airflow/providers/standard/sensors/python.py +++ b/providers/standard/src/airflow/providers/standard/sensors/python.py @@ -20,15 +20,11 @@ from collections.abc import Callable, Mapping, Sequence from typing import TYPE_CHECKING, Any -from airflow.providers.standard.version_compat import BaseSensorOperator, PokeReturnValue, context_merge +from airflow.providers.common.compat.sdk import BaseSensorOperator, PokeReturnValue, context_merge from airflow.utils.operator_helpers import determine_kwargs if TYPE_CHECKING: - try: - from airflow.sdk.definitions.context import Context - except ImportError: - # TODO: Remove once provider drops support for Airflow 2 - from airflow.utils.context import Context # type: ignore[no-redef, attr-defined] + from airflow.providers.common.compat.sdk import Context class PythonSensor(BaseSensorOperator): diff --git a/providers/standard/src/airflow/providers/standard/sensors/time.py b/providers/standard/src/airflow/providers/standard/sensors/time.py index 129ae22fbd882..03b0dd2333fef 100644 --- a/providers/standard/src/airflow/providers/standard/sensors/time.py +++ b/providers/standard/src/airflow/providers/standard/sensors/time.py @@ -24,8 +24,8 @@ from airflow.configuration import conf from airflow.exceptions import AirflowProviderDeprecationWarning +from airflow.providers.common.compat.sdk import BaseSensorOperator, timezone from airflow.providers.standard.triggers.temporal import DateTimeTrigger -from airflow.providers.standard.version_compat import BaseSensorOperator try: from airflow.triggers.base import StartTriggerArgs # type: ignore[no-redef] @@ -42,11 +42,6 @@ class StartTriggerArgs: # type: ignore[no-redef] timeout: datetime.timedelta | None = None -try: - from airflow.sdk import timezone -except ImportError: - from airflow.utils import timezone # type: ignore[attr-defined,no-redef] - if TYPE_CHECKING: from airflow.sdk import Context diff --git a/providers/standard/src/airflow/providers/standard/sensors/time_delta.py b/providers/standard/src/airflow/providers/standard/sensors/time_delta.py index 492e9aafa1da6..bc27afd531823 100644 --- a/providers/standard/src/airflow/providers/standard/sensors/time_delta.py +++ b/providers/standard/src/airflow/providers/standard/sensors/time_delta.py @@ -27,16 +27,12 @@ from airflow.configuration import conf from airflow.exceptions import AirflowProviderDeprecationWarning, AirflowSkipException +from airflow.providers.common.compat.sdk import BaseSensorOperator, timezone from airflow.providers.standard.triggers.temporal import DateTimeTrigger, TimeDeltaTrigger -from airflow.providers.standard.version_compat import AIRFLOW_V_3_0_PLUS, BaseSensorOperator -from airflow.utils import timezone +from airflow.providers.standard.version_compat import AIRFLOW_V_3_0_PLUS if TYPE_CHECKING: - try: - from airflow.sdk.definitions.context import Context - except ImportError: - # TODO: Remove once provider drops support for Airflow 2 - from airflow.utils.context import Context + from airflow.providers.common.compat.sdk import Context def _get_airflow_version(): diff --git a/providers/standard/src/airflow/providers/standard/sensors/weekday.py b/providers/standard/src/airflow/providers/standard/sensors/weekday.py index ac23e1ac9cb6c..382230e90acc9 100644 --- a/providers/standard/src/airflow/providers/standard/sensors/weekday.py +++ b/providers/standard/src/airflow/providers/standard/sensors/weekday.py @@ -20,16 +20,11 @@ from collections.abc import Iterable from typing import TYPE_CHECKING +from airflow.providers.common.compat.sdk import BaseSensorOperator, timezone from airflow.providers.standard.utils.weekday import WeekDay -from airflow.providers.standard.version_compat import BaseSensorOperator -from airflow.utils import timezone if TYPE_CHECKING: - try: - from airflow.sdk.definitions.context import Context - except ImportError: - # TODO: Remove once provider drops support for Airflow 2 - from airflow.utils.context import Context + from airflow.providers.common.compat.sdk import Context class DayOfWeekSensor(BaseSensorOperator): diff --git a/providers/standard/src/airflow/providers/standard/triggers/temporal.py b/providers/standard/src/airflow/providers/standard/triggers/temporal.py index 599a219dbee39..cfc206ff3131d 100644 --- a/providers/standard/src/airflow/providers/standard/triggers/temporal.py +++ b/providers/standard/src/airflow/providers/standard/triggers/temporal.py @@ -23,13 +23,9 @@ import pendulum +from airflow.providers.common.compat.sdk import timezone from airflow.triggers.base import BaseTrigger, TaskSuccessEvent, TriggerEvent -try: - from airflow.sdk import timezone -except ImportError: - from airflow.utils import timezone # type: ignore[attr-defined,no-redef] - class DateTimeTrigger(BaseTrigger): """ diff --git a/providers/standard/src/airflow/providers/standard/utils/skipmixin.py b/providers/standard/src/airflow/providers/standard/utils/skipmixin.py index f00bbfcb19d24..2a00551f34f84 100644 --- a/providers/standard/src/airflow/providers/standard/utils/skipmixin.py +++ b/providers/standard/src/airflow/providers/standard/utils/skipmixin.py @@ -22,7 +22,6 @@ from typing import TYPE_CHECKING from airflow.exceptions import AirflowException -from airflow.providers.standard.version_compat import AIRFLOW_V_3_0_PLUS from airflow.utils.log.logging_mixin import LoggingMixin if TYPE_CHECKING: @@ -40,12 +39,7 @@ def _ensure_tasks(nodes: Iterable[DAGNode]) -> Sequence[Operator]: - if AIRFLOW_V_3_0_PLUS: - from airflow.sdk import BaseOperator - from airflow.sdk.definitions.mappedoperator import MappedOperator - else: - from airflow.models.baseoperator import BaseOperator # type: ignore[no-redef] - from airflow.models.mappedoperator import MappedOperator # type: ignore[assignment,no-redef] + from airflow.providers.common.compat.sdk import BaseOperator, MappedOperator return [n for n in nodes if isinstance(n, (BaseOperator, MappedOperator))] diff --git a/providers/standard/src/airflow/providers/standard/version_compat.py b/providers/standard/src/airflow/providers/standard/version_compat.py index 10820dbe1fdba..3539d070dab58 100644 --- a/providers/standard/src/airflow/providers/standard/version_compat.py +++ b/providers/standard/src/airflow/providers/standard/version_compat.py @@ -36,34 +36,16 @@ def get_base_airflow_version_tuple() -> tuple[int, int, int]: AIRFLOW_V_3_1_PLUS: bool = get_base_airflow_version_tuple() >= (3, 1, 0) AIRFLOW_V_3_2_PLUS: bool = get_base_airflow_version_tuple() >= (3, 2, 0) -# BaseOperator is not imported from SDK from 3.0 (and only done from 3.1) due to a bug with -# DecoratedOperator -- where `DecoratedOperator._handle_output` needed `xcom_push` to exist on `BaseOperator` -# even though it wasn't used. +# BaseOperator: Use 3.1+ due to xcom_push method missing in SDK BaseOperator 3.0.x +# This is needed for DecoratedOperator compatibility if AIRFLOW_V_3_1_PLUS: - from airflow.sdk import BaseHook, BaseOperator, timezone - from airflow.sdk.definitions.context import context_merge + from airflow.sdk import BaseOperator else: - from airflow.hooks.base import BaseHook # type: ignore[attr-defined,no-redef] from airflow.models.baseoperator import BaseOperator # type: ignore[no-redef] - from airflow.utils import timezone # type: ignore[attr-defined,no-redef] - from airflow.utils.context import context_merge # type: ignore[no-redef, attr-defined] - -if AIRFLOW_V_3_0_PLUS: - from airflow.sdk import BaseOperatorLink - from airflow.sdk.bases.sensor import BaseSensorOperator, PokeReturnValue -else: - from airflow.models.baseoperatorlink import BaseOperatorLink # type: ignore[no-redef] - from airflow.sensors.base import BaseSensorOperator, PokeReturnValue # type: ignore[no-redef] __all__ = [ "AIRFLOW_V_3_0_PLUS", "AIRFLOW_V_3_1_PLUS", "AIRFLOW_V_3_2_PLUS", "BaseOperator", - "BaseOperatorLink", - "BaseHook", - "BaseSensorOperator", - "PokeReturnValue", - "context_merge", - "timezone", ] diff --git a/scripts/ci/prek/check_common_compat_lazy_imports.py b/scripts/ci/prek/check_common_compat_lazy_imports.py index 8e2e85798922b..6d40eafb1c7d6 100755 --- a/scripts/ci/prek/check_common_compat_lazy_imports.py +++ b/scripts/ci/prek/check_common_compat_lazy_imports.py @@ -17,16 +17,15 @@ # under the License. """ -Check and generate sdk.pyi from sdk.py. +Validate that TYPE_CHECKING block in sdk.py matches runtime import maps. -This script can be used as: -1. Pre-commit hook - checks if .pyi is in sync with _IMPORT_MAP -2. Manual generation - generates .pyi file from _IMPORT_MAP +This pre-commit hook ensures that: +1. All items in _IMPORT_MAP, _RENAME_MAP, and _MODULE_MAP are in TYPE_CHECKING +2. All items in TYPE_CHECKING are in one of the runtime maps +3. No mismatches between type hints and runtime behavior Usage: - python scripts/ci/prek/check_common_compat_lazy_imports.py # Check only (pre-commit) - python scripts/ci/prek/check_common_compat_lazy_imports.py --generate # Generate .pyi file - python scripts/ci/prek/check_common_compat_lazy_imports.py --validate # Generate with import validation + python scripts/ci/prek/check_common_compat_lazy_imports.py """ from __future__ import annotations @@ -36,266 +35,85 @@ from pathlib import Path -def extract_import_map(py_file: Path) -> dict[str, str | tuple[str, ...]]: +def extract_runtime_maps(py_file: Path) -> tuple[set[str], set[str], set[str]]: """ - Extract _IMPORT_MAP from sdk.py. + Extract all names from _IMPORT_MAP, _RENAME_MAP, and _MODULE_MAP. - :param py_file: Path to sdk.py - :return: Dictionary mapping class names to module paths + Returns tuple of (import_names, rename_names, module_names) """ content = py_file.read_text() tree = ast.parse(content) - for node in tree.body: - if isinstance(node, ast.AnnAssign) and isinstance(node.target, ast.Name): - if node.target.id == "_IMPORT_MAP" and node.value: - return ast.literal_eval(node.value) - elif isinstance(node, ast.Assign): - for target in node.targets: - if isinstance(target, ast.Name) and target.id == "_IMPORT_MAP": - return ast.literal_eval(node.value) - - raise ValueError("Could not find _IMPORT_MAP in sdk.py") - - -def extract_rename_map(py_file: Path) -> dict[str, tuple[str, str, str]]: - """ - Extract _RENAME_MAP from sdk.py. - - :param py_file: Path to sdk.py - :return: Dictionary mapping new class names to (new_path, old_path, old_name) - """ - content = py_file.read_text() - tree = ast.parse(content) + import_map = set() + rename_map = set() + module_map = set() for node in tree.body: - if isinstance(node, ast.AnnAssign) and isinstance(node.target, ast.Name): - if node.target.id == "_RENAME_MAP" and node.value: - return ast.literal_eval(node.value) - elif isinstance(node, ast.Assign): - for target in node.targets: - if isinstance(target, ast.Name) and target.id == "_RENAME_MAP": - return ast.literal_eval(node.value) - - raise ValueError("Could not find _RENAME_MAP in sdk.py") - - -def extract_module_map(py_file: Path) -> dict[str, str | tuple[str, ...]]: - """ - Extract _MODULE_MAP from sdk.py. - - :param py_file: Path to sdk.py - :return: Dictionary mapping module names to module paths - """ - content = py_file.read_text() - tree = ast.parse(content) + # Handle both annotated assignments and regular assignments + targets = [] + value = None - for node in tree.body: if isinstance(node, ast.AnnAssign) and isinstance(node.target, ast.Name): - if node.target.id == "_MODULE_MAP" and node.value: - return ast.literal_eval(node.value) + targets = [node.target] + value = node.value elif isinstance(node, ast.Assign): - for target in node.targets: - if isinstance(target, ast.Name) and target.id == "_MODULE_MAP": - return ast.literal_eval(node.value) - - raise ValueError("Could not find _MODULE_MAP in sdk.py") + targets = node.targets + value = node.value + for target in targets: + if not isinstance(target, ast.Name): + continue -def generate_pyi_content( - rename_map: dict[str, tuple[str, str, str]], - import_map: dict[str, str | tuple[str, ...]], - module_map: dict[str, str | tuple[str, ...]], -) -> str: - """ - Generate .pyi stub content from rename, import and module maps. + if target.id == "_IMPORT_MAP" and value: + data = ast.literal_eval(value) + import_map = set(data.keys()) + elif target.id == "_RENAME_MAP" and value: + data = ast.literal_eval(value) + rename_map = set(data.keys()) + elif target.id == "_MODULE_MAP" and value: + data = ast.literal_eval(value) + module_map = set(data.keys()) - :param rename_map: Dictionary mapping new names to (new_path, old_path, old_name) - :param import_map: Dictionary mapping class names to module paths - :param module_map: Dictionary mapping module names to module paths - :return: Content for the .pyi file - """ - header = '''# 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. -""" -Type stubs for IDE autocomplete - always uses Airflow 3 paths. + return import_map, rename_map, module_map -This file is auto-generated from sdk.py. - - run scripts/ci/prek/check_common_compat_lazy_imports.py --generate instead. -""" -''' - - imports_by_module: dict[str, list[str]] = {} - - # Process renamed imports from _RENAME_MAP (use new names and new paths for type hints) - for new_name, (new_path, _old_path, _old_name) in sorted(rename_map.items()): - if new_path not in imports_by_module: - imports_by_module[new_path] = [] - imports_by_module[new_path].append(new_name) - - # Process regular imports from _IMPORT_MAP - for name, paths in sorted(import_map.items()): - module_path = paths[0] if isinstance(paths, tuple) else paths - if module_path not in imports_by_module: - imports_by_module[module_path] = [] - imports_by_module[module_path].append(name) - - lines = [] - - # Generate regular imports (always use multiline for ruff compatibility) - for module_path in sorted(imports_by_module.keys()): - names = sorted(imports_by_module[module_path]) - # Always use multiline format for ruff - lines.append(f"from {module_path} import (") - for i, name in enumerate(names): - comma = "," if i < len(names) - 1 else "" - lines.append(f" {name} as {name}{comma}") - lines.append(")") - - # Generate module imports (import the module itself) - for module_name, paths in sorted(module_map.items()): - module_path = paths[0] if isinstance(paths, tuple) else paths - lines.append(f"import {module_path} as {module_name}") - - # Generate __all__ (include renamed, attributes and modules) - all_names = sorted(list(rename_map.keys()) + list(import_map.keys()) + list(module_map.keys())) - lines.append("") - lines.append(f"__all__: list[str] = {all_names!r}") - - return header + "\n".join(lines) + "\n" - - -def validate_imports( - rename_map: dict[str, tuple[str, str, str]], - import_map: dict[str, str | tuple[str, ...]], - module_map: dict[str, str | tuple[str, ...]], - skip_on_error: bool = False, -) -> list[str]: +def extract_type_checking_names(py_file: Path) -> set[str]: """ - Validate that all imports in the maps are actually importable. - - This is optional and only runs if skip_on_error=False. It requires Airflow - and all providers to be installed, so it's meant for manual validation only. - - :param rename_map: The rename map to validate - :param import_map: The import map to validate - :param module_map: The module map to validate - :param skip_on_error: If True, skip validation and return empty list - :return: List of errors (empty if all valid or skipped) + Extract all imported names from TYPE_CHECKING block. """ - if skip_on_error: - print("\nSkipping import validation (requires full Airflow installation)") - return [] - - import importlib - - errors = [] - print("\nValidating imports (requires Airflow + providers installed)...") - - # Validate renamed imports - for new_name, (new_path, old_path, old_name) in rename_map.items(): - importable = False - last_error = None - - # Try new path with new name (Airflow 3.x) - try: - module = importlib.import_module(new_path) - if hasattr(module, new_name): - importable = True - else: - last_error = f"Module {new_path} does not have attribute {new_name}" - except (ImportError, ModuleNotFoundError) as e: - last_error = str(e) - - # Try old path with old name (Airflow 2.x) - if not importable: - try: - module = importlib.import_module(old_path) - if hasattr(module, old_name): - importable = True - else: - last_error = f"Module {old_path} does not have attribute {old_name}" - except (ImportError, ModuleNotFoundError) as e: - last_error = str(e) - - if not importable: - errors.append( - f" ✗ {new_name} (renamed from {old_name}): Could not import. Last error: {last_error}" - ) - else: - print(f" ✓ {new_name} (renamed from {old_name})") - - # Validate attribute imports - for name, paths in import_map.items(): - if isinstance(paths, str): - paths = (paths,) - - importable = False - last_error = None - - for module_path in paths: - try: - module = importlib.import_module(module_path) - if hasattr(module, name): - importable = True - break - last_error = f"Module {module_path} does not have attribute {name}" - except (ImportError, ModuleNotFoundError) as e: - last_error = str(e) - continue - - if not importable: - errors.append(f" ✗ {name}: Could not import from any path. Last error: {last_error}") - else: - print(f" ✓ {name}") - - # Validate module imports - for module_name, paths in module_map.items(): - if isinstance(paths, str): - paths = (paths,) - - importable = False - last_error = None - - for module_path in paths: - try: - importlib.import_module(module_path) - importable = True - break - except (ImportError, ModuleNotFoundError) as e: - last_error = str(e) - continue - - if not importable: - errors.append( - f" ✗ {module_name} (module): Could not import from any path. Last error: {last_error}" - ) - else: - print(f" ✓ {module_name} (module)") - - return errors - + content = py_file.read_text() + tree = ast.parse(content) -def main() -> int: - """Generate and check sdk.pyi.""" - repo_root = Path(__file__).parent.parent.parent.parent - lazy_compat_py = ( - repo_root + type_checking_names = set() + + for node in ast.walk(tree): + # Find: if TYPE_CHECKING: + if isinstance(node, ast.If): + # Check if condition is comparing TYPE_CHECKING + if isinstance(node.test, ast.Name) and node.test.id == "TYPE_CHECKING": + # Extract all imports in this block + for stmt in node.body: + if isinstance(stmt, ast.ImportFrom): + # from X import Y as Z + for alias in stmt.names: + # Use alias.asname if present, else alias.name + name = alias.asname if alias.asname else alias.name + type_checking_names.add(name) + elif isinstance(stmt, ast.Import): + # import X as Y or import X.Y as Z + for alias in stmt.names: + # For module imports like "import airflow.sdk.io as io" + # Use the alias name (io) not the full module path + name = alias.asname if alias.asname else alias.name.split(".")[-1] + type_checking_names.add(name) + + return type_checking_names + + +def main(): + """Validate TYPE_CHECKING block matches runtime maps.""" + sdk_py = ( + Path(__file__).parent.parent.parent.parent / "providers" / "common" / "compat" @@ -306,127 +124,60 @@ def main() -> int: / "compat" / "sdk.py" ) - lazy_compat_pyi = lazy_compat_py.with_suffix(".pyi") - - if not lazy_compat_py.exists(): - print(f"ERROR: Could not find {lazy_compat_py}") - return 1 - - should_generate = "--generate" in sys.argv or "--validate" in sys.argv - should_validate = "--validate" in sys.argv - - try: - rename_map = extract_rename_map(lazy_compat_py) - print(f"Found {len(rename_map)} renames in _RENAME_MAP") - except Exception as e: - print(f"ERROR: Failed to extract _RENAME_MAP: {e}") - return 1 - - try: - import_map = extract_import_map(lazy_compat_py) - print(f"Found {len(import_map)} imports in _IMPORT_MAP") - except Exception as e: - print(f"ERROR: Failed to extract _IMPORT_MAP: {e}") - return 1 - - try: - module_map = extract_module_map(lazy_compat_py) - print(f"Found {len(module_map)} modules in _MODULE_MAP") - except Exception as e: - print(f"ERROR: Failed to extract _MODULE_MAP: {e}") - return 1 - - total_imports = len(rename_map) + len(import_map) + len(module_map) - - errors = validate_imports(rename_map, import_map, module_map, skip_on_error=not should_validate) + + if not sdk_py.exists(): + print(f"❌ ERROR: {sdk_py} not found") + sys.exit(1) + + # Extract runtime maps + import_names, rename_names, module_names = extract_runtime_maps(sdk_py) + runtime_names = import_names | rename_names | module_names + + # Extract TYPE_CHECKING imports + type_checking_names = extract_type_checking_names(sdk_py) + + # Check for discrepancies + missing_in_type_checking = runtime_names - type_checking_names + extra_in_type_checking = type_checking_names - runtime_names + + errors = [] + + if missing_in_type_checking: + errors.append("\n❌ Items in runtime maps but MISSING in TYPE_CHECKING block:") + for name in sorted(missing_in_type_checking): + # Determine which map it's from + map_name = [] + if name in import_names: + map_name.append("_IMPORT_MAP") + if name in rename_names: + map_name.append("_RENAME_MAP") + if name in module_names: + map_name.append("_MODULE_MAP") + errors.append(f" - {name} (in {', '.join(map_name)})") + + if extra_in_type_checking: + errors.append("\n❌ Items in TYPE_CHECKING block but NOT in any runtime map:") + for name in sorted(extra_in_type_checking): + errors.append(f" - {name}") + if errors: - print("\n❌ Import validation failed:") - for error in errors: - print(error) - print("\nPlease fix the import paths in _IMPORT_MAP and _MODULE_MAP and try again.") - return 1 - - if should_validate: - print(f"\n✓ All {total_imports} imports validated successfully") - - # Check if .pyi exists and is in sync - if not should_generate: - # Check-only mode (pre-commit) - if not lazy_compat_pyi.exists(): - print(f"ERROR: {lazy_compat_pyi.name} does not exist") - print("Run: python scripts/ci/prek/check_common_compat_lazy_imports.py --generate") - return 1 - - pyi_content = lazy_compat_pyi.read_text() - - # Count total imports in .pyi (each "X as X" pattern + "import X as Y", excluding __all__) - import re - - # Match "import X.Y.Z as module_name" pattern (standalone module imports) - module_import_pattern = r"^import\s+[\w.]+\s+as\s+(\w+)" - pyi_module_imports = set(re.findall(module_import_pattern, pyi_content, re.MULTILINE)) - - # Remove __all__ and standalone import lines to avoid false matches - pyi_for_attr_search = pyi_content - pyi_for_attr_search = re.sub(r"__all__:.*", "", pyi_for_attr_search, flags=re.DOTALL) - pyi_for_attr_search = re.sub( - r"^import\s+[\w.]+\s+as\s+\w+.*$", "", pyi_for_attr_search, flags=re.MULTILINE + print("\n".join(errors)) + print("\n❌ FAILED: TYPE_CHECKING block and runtime maps are out of sync!") + print("\nTo fix:") + print(f"1. Add missing items to TYPE_CHECKING block in {sdk_py}") + print("2. Remove extra items from TYPE_CHECKING block") + print( + "3. Ensure every item in _IMPORT_MAP/_RENAME_MAP/_MODULE_MAP has a corresponding TYPE_CHECKING import" ) + sys.exit(1) - # Match all "Name as Name" patterns - attr_import_pattern = r"(\w+)\s+as\s+\1" - pyi_attr_imports = set(re.findall(attr_import_pattern, pyi_for_attr_search)) - - # Combine all expected imports - map_renames = set(rename_map.keys()) - map_attrs = set(import_map.keys()) - map_modules = set(module_map.keys()) - all_expected_attrs = map_renames | map_attrs - - # Check for discrepancies - missing_attrs = all_expected_attrs - pyi_attr_imports - extra_attrs = pyi_attr_imports - all_expected_attrs - missing_modules = map_modules - pyi_module_imports - extra_modules = pyi_module_imports - map_modules - - if not (missing_attrs or extra_attrs or missing_modules or extra_modules): - print(f"✓ sdk.pyi is in sync with sdk.py ({total_imports} imports)") - return 0 - - # Out of sync - if missing_attrs: - print(f"ERROR: sdk.pyi is missing {len(missing_attrs)} attributes from _RENAME_MAP/_IMPORT_MAP:") - for name in sorted(missing_attrs)[:10]: - print(f" - {name}") - if len(missing_attrs) > 10: - print(f" ... and {len(missing_attrs) - 10} more") - - if extra_attrs: - print(f"ERROR: sdk.pyi has {len(extra_attrs)} extra attributes not in _RENAME_MAP/_IMPORT_MAP:") - for name in sorted(extra_attrs)[:10]: - print(f" + {name}") - if len(extra_attrs) > 10: - print(f" ... and {len(extra_attrs) - 10} more") - - if missing_modules: - print(f"ERROR: sdk.pyi is missing {len(missing_modules)} modules from _MODULE_MAP:") - for name in sorted(missing_modules): - print(f" - {name} (module)") - - if extra_modules: - print(f"ERROR: sdk.pyi has {len(extra_modules)} extra modules not in _MODULE_MAP:") - for name in sorted(extra_modules): - print(f" + {name} (module)") - - print("\nRun: python scripts/ci/prek/check_common_compat_lazy_imports.py --generate") - return 1 - - # Generate mode - new_pyi_content = generate_pyi_content(rename_map, import_map, module_map) - lazy_compat_pyi.write_text(new_pyi_content) - print(f"✓ Generated {lazy_compat_pyi.name} with {total_imports} imports") - return 0 + print("✅ SUCCESS: TYPE_CHECKING block matches runtime maps") + print(f" - {len(import_names)} items in _IMPORT_MAP") + print(f" - {len(rename_names)} items in _RENAME_MAP") + print(f" - {len(module_names)} items in _MODULE_MAP") + print(f" - {len(type_checking_names)} items in TYPE_CHECKING") + sys.exit(0) if __name__ == "__main__": - sys.exit(main()) + main()