Simplify version-specific imports in the Standard provider

.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

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"),
     # ============================================================================

providers/standard/pyproject.toml

@@ -58,13 +58,15 @@ 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]
 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",
 ]

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):

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):

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):

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):

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):

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):

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):