Deprecate decorators from Core

airflow-core/src/airflow/decorators/__init__.py

@@ -16,16 +16,16 @@
 # under the License.
 from __future__ import annotations
 
-from airflow.sdk.definitions.decorators import (
-    dag as dag,
-    setup as setup,
-    task as task,
-    task_group as task_group,
-    teardown as teardown,
-)
 from airflow.utils.deprecation_tools import add_deprecated_classes
 
 __deprecated_classes = {
+    __name__: {
+        "dag": "airflow.sdk.dag",
+        "setup": "airflow.sdk.setup",
+        "task": "airflow.sdk.task",
+        "task_group": "airflow.sdk.task_group",
+        "teardown": "airflow.sdk.teardown",
+    },
     "base": {
         "DecoratedMappedOperator": "airflow.sdk.bases.decorator.DecoratedMappedOperator",
         "DecoratedOperator": "airflow.sdk.bases.decorator.DecoratedOperator",

providers/amazon/tests/system/amazon/aws/example_appflow_run.py

@@ -18,7 +18,6 @@
 
 import json
 from datetime import datetime
-from typing import TYPE_CHECKING
 
 import boto3
 
@@ -33,18 +32,14 @@
 
 from tests_common.test_utils.version_compat import AIRFLOW_V_3_0_PLUS
 
-if TYPE_CHECKING:
-    from airflow.decorators import task
-    from airflow.models.baseoperator import chain
-    from airflow.models.dag import DAG
+if AIRFLOW_V_3_0_PLUS:
+    from airflow.sdk import DAG, chain, task
 else:
-    if AIRFLOW_V_3_0_PLUS:
-        from airflow.sdk import DAG, chain, task
-    else:
-        # Airflow 2.10 compat
-        from airflow.decorators import task
-        from airflow.models.baseoperator import chain
-        from airflow.models.dag import DAG
+    # Airflow 2 path
+    from airflow.decorators import task  # type: ignore[attr-defined,no-redef]
+    from airflow.models.baseoperator import chain  # type: ignore[attr-defined,no-redef]
+    from airflow.models.dag import DAG  # type: ignore[attr-defined,no-redef,assignment]
+
 from airflow.utils.trigger_rule import TriggerRule
 
 from system.amazon.aws.utils import SystemTestContextBuilder

providers/amazon/tests/system/amazon/aws/example_athena.py

@@ -17,7 +17,6 @@
 from __future__ import annotations
 
 from datetime import datetime
-from typing import TYPE_CHECKING
 
 import boto3
 
@@ -32,18 +31,14 @@
 
 from tests_common.test_utils.version_compat import AIRFLOW_V_3_0_PLUS
 
-if TYPE_CHECKING:
-    from airflow.decorators import task
-    from airflow.models.baseoperator import chain
-    from airflow.models.dag import DAG
+if AIRFLOW_V_3_0_PLUS:
+    from airflow.sdk import DAG, chain, task
 else:
-    if AIRFLOW_V_3_0_PLUS:
-        from airflow.sdk import DAG, chain, task
-    else:
-        # Airflow 2.10 compat
-        from airflow.decorators import task
-        from airflow.models.baseoperator import chain
-        from airflow.models.dag import DAG
+    # Airflow 2 path
+    from airflow.decorators import task  # type: ignore[attr-defined,no-redef]
+    from airflow.models.baseoperator import chain  # type: ignore[attr-defined,no-redef]
+    from airflow.models.dag import DAG  # type: ignore[attr-defined,no-redef,assignment]
+
 from airflow.utils.trigger_rule import TriggerRule
 
 from system.amazon.aws.utils import SystemTestContextBuilder
@@ -193,7 +188,6 @@ def read_results_from_s3(bucket_name, query_execution_id):
     # when "tearDown" task with trigger rule is part of the DAG
     list(dag.tasks) >> watcher()
 
-
 from tests_common.test_utils.system_tests import get_test_run  # noqa: E402
 
 # Needed to run the example DAG with pytest (see: tests/system/README.md#run_via_pytest)

providers/amazon/tests/system/amazon/aws/example_batch.py

@@ -18,7 +18,6 @@
 
 import logging
 from datetime import datetime
-from typing import TYPE_CHECKING
 
 import boto3
 
@@ -32,18 +31,14 @@
 
 from tests_common.test_utils.version_compat import AIRFLOW_V_3_0_PLUS
 
-if TYPE_CHECKING:
-    from airflow.decorators import task
-    from airflow.models.baseoperator import chain
-    from airflow.models.dag import DAG
+if AIRFLOW_V_3_0_PLUS:
+    from airflow.sdk import DAG, chain, task
 else:
-    if AIRFLOW_V_3_0_PLUS:
-        from airflow.sdk import DAG, chain, task
-    else:
-        # Airflow 2.10 compat
-        from airflow.decorators import task
-        from airflow.models.baseoperator import chain
-        from airflow.models.dag import DAG
+    # Airflow 2 path
+    from airflow.decorators import task  # type: ignore[attr-defined,no-redef]
+    from airflow.models.baseoperator import chain  # type: ignore[attr-defined,no-redef]
+    from airflow.models.dag import DAG  # type: ignore[attr-defined,no-redef,assignment]
+
 from airflow.utils.trigger_rule import TriggerRule
 
 from system.amazon.aws.utils import (

providers/amazon/tests/system/amazon/aws/example_bedrock.py

@@ -19,7 +19,6 @@
 import json
 from datetime import datetime
 from os import environ
-from typing import TYPE_CHECKING
 
 import boto3
 
@@ -41,18 +40,14 @@
 
 from tests_common.test_utils.version_compat import AIRFLOW_V_3_0_PLUS
 
-if TYPE_CHECKING:
-    from airflow.decorators import task, task_group
-    from airflow.models.baseoperator import chain
-    from airflow.models.dag import DAG
+if AIRFLOW_V_3_0_PLUS:
+    from airflow.sdk import DAG, chain, task, task_group
 else:
-    if AIRFLOW_V_3_0_PLUS:
-        from airflow.sdk import DAG, chain, task, task_group
-    else:
-        # Airflow 2.10 compat
-        from airflow.decorators import task, task_group
-        from airflow.models.baseoperator import chain
-        from airflow.models.dag import DAG
+    # Airflow 2 path
+    from airflow.decorators import task, task_group  # type: ignore[attr-defined,no-redef]
+    from airflow.models.baseoperator import chain  # type: ignore[attr-defined,no-redef]
+    from airflow.models.dag import DAG  # type: ignore[attr-defined,no-redef,assignment]
+
 from airflow.utils.trigger_rule import TriggerRule
 
 from system.amazon.aws.utils import SystemTestContextBuilder
@@ -73,7 +68,6 @@
 # on the AWS System Test stack.
 SKIP_PROVISION_THROUGHPUT = environ.get("SKIP_RESTRICTED_SYSTEM_TEST_TASKS", default=True)
 
-
 LLAMA_SHORT_MODEL_ID = "meta.llama3-8b-instruct-v1:0"
 TITAN_MODEL_ID = "amazon.titan-text-express-v1:0:8k"
 TITAN_SHORT_MODEL_ID = TITAN_MODEL_ID.split(":")[0]
@@ -226,7 +220,6 @@ def should_run_provision_throughput():
     # when "tearDown" task with trigger rule is part of the DAG
     list(dag.tasks) >> watcher()
 
-
 from tests_common.test_utils.system_tests import get_test_run  # noqa: E402
 
 # Needed to run the example DAG with pytest (see: tests/system/README.md#run_via_pytest)

providers/amazon/tests/system/amazon/aws/example_bedrock_batch_inference.py

@@ -20,7 +20,6 @@
 import logging
 from datetime import datetime
 from tempfile import NamedTemporaryFile
-from typing import TYPE_CHECKING
 
 from botocore.exceptions import ClientError
 
@@ -38,25 +37,20 @@
 
 from tests_common.test_utils.version_compat import AIRFLOW_V_3_0_PLUS
 
-if TYPE_CHECKING:
-    from airflow.decorators import task
-    from airflow.models.baseoperator import chain
-    from airflow.models.dag import DAG
+if AIRFLOW_V_3_0_PLUS:
+    from airflow.sdk import DAG, chain, task
 else:
-    if AIRFLOW_V_3_0_PLUS:
-        from airflow.sdk import DAG, chain, task
-    else:
-        # Airflow 2.10 compat
-        from airflow.decorators import task
-        from airflow.models.baseoperator import chain
-        from airflow.models.dag import DAG
+    # Airflow 2 path
+    from airflow.decorators import task  # type: ignore[attr-defined,no-redef]
+    from airflow.models.baseoperator import chain  # type: ignore[attr-defined,no-redef]
+    from airflow.models.dag import DAG  # type: ignore[attr-defined,no-redef,assignment]
+
 from airflow.utils.trigger_rule import TriggerRule
 
 from system.amazon.aws.utils import SystemTestContextBuilder
 
 log = logging.getLogger(__name__)
 
-
 # Externally fetched variables:
 ROLE_ARN_KEY = "ROLE_ARN"
 sys_test_context_task = SystemTestContextBuilder().add_variable(ROLE_ARN_KEY).build()
@@ -198,7 +192,6 @@ def stop_batch_inference(job_arn: str):
     # when "tearDown" task with trigger rule is part of the DAG
     list(dag.tasks) >> watcher()
 
-
 from tests_common.test_utils.system_tests import get_test_run  # noqa: E402
 
 # Needed to run the example DAG with pytest (see: tests/system/README.md#run_via_pytest)

providers/amazon/tests/system/amazon/aws/example_bedrock_retrieve_and_generate.py

@@ -22,7 +22,6 @@
 import tempfile
 from datetime import datetime
 from time import sleep
-from typing import TYPE_CHECKING
 from urllib.request import urlretrieve
 
 import boto3
@@ -59,20 +58,14 @@
 
 from tests_common.test_utils.version_compat import AIRFLOW_V_3_0_PLUS
 
-if TYPE_CHECKING:
-    from airflow.decorators import task, task_group
-    from airflow.models.baseoperator import chain
-    from airflow.models.dag import DAG
-    from airflow.sdk import Label
+if AIRFLOW_V_3_0_PLUS:
+    from airflow.sdk import DAG, Label, chain, task, task_group
 else:
-    if AIRFLOW_V_3_0_PLUS:
-        from airflow.sdk import DAG, Label, chain, task, task_group
-    else:
-        # Airflow 2.10 compat
-        from airflow.decorators import task, task_group
-        from airflow.models.baseoperator import chain
-        from airflow.models.dag import DAG
-        from airflow.sdk import Label
+    # Airflow 2 path
+    from airflow.decorators import task, task_group  # type: ignore[attr-defined,no-redef]
+    from airflow.models.baseoperator import chain  # type: ignore[attr-defined,no-redef]
+    from airflow.models.dag import DAG  # type: ignore[attr-defined,no-redef,assignment]
+    from airflow.sdk import Label  # type: ignore[attr-defined,no-redef]
 from airflow.utils.trigger_rule import TriggerRule
 
 from system.amazon.aws.utils import SystemTestContextBuilder
@@ -615,7 +608,6 @@ def delete_opensearch_policies(collection_name: str):
     # when "tearDown" task with trigger rule is part of the DAG
     list(dag.tasks) >> watcher()
 
-
 from tests_common.test_utils.system_tests import get_test_run  # noqa: E402
 
 # Needed to run the example DAG with pytest (see: tests/system/README.md#run_via_pytest)

providers/amazon/tests/system/amazon/aws/example_comprehend.py

@@ -18,7 +18,6 @@
 
 import json
 from datetime import datetime
-from typing import TYPE_CHECKING
 
 from airflow.providers.amazon.aws.operators.comprehend import ComprehendStartPiiEntitiesDetectionJobOperator
 from airflow.providers.amazon.aws.operators.s3 import (
@@ -32,18 +31,14 @@
 
 from tests_common.test_utils.version_compat import AIRFLOW_V_3_0_PLUS
 
-if TYPE_CHECKING:
-    from airflow.decorators import task_group
-    from airflow.models.baseoperator import chain
-    from airflow.models.dag import DAG
+if AIRFLOW_V_3_0_PLUS:
+    from airflow.sdk import DAG, chain, task_group
 else:
-    if AIRFLOW_V_3_0_PLUS:
-        from airflow.sdk import DAG, chain, task_group
-    else:
-        # Airflow 2.10 compat
-        from airflow.decorators import task_group
-        from airflow.models.baseoperator import chain
-        from airflow.models.dag import DAG
+    # Airflow 2 path
+    from airflow.decorators import task_group  # type: ignore[attr-defined,no-redef]
+    from airflow.models.baseoperator import chain  # type: ignore[attr-defined,no-redef]
+    from airflow.models.dag import DAG  # type: ignore[attr-defined,no-redef,assignment]
+
 from airflow.utils.trigger_rule import TriggerRule
 
 from system.amazon.aws.utils import SystemTestContextBuilder

providers/amazon/tests/system/amazon/aws/example_comprehend_document_classifier.py

@@ -17,7 +17,6 @@
 from __future__ import annotations
 
 from datetime import datetime
-from typing import TYPE_CHECKING
 
 from airflow.providers.amazon.aws.hooks.comprehend import ComprehendHook
 from airflow.providers.amazon.aws.operators.comprehend import (
@@ -35,18 +34,14 @@
 
 from tests_common.test_utils.version_compat import AIRFLOW_V_3_0_PLUS
 
-if TYPE_CHECKING:
-    from airflow.decorators import task, task_group
-    from airflow.models.baseoperator import chain
-    from airflow.models.dag import DAG
+if AIRFLOW_V_3_0_PLUS:
+    from airflow.sdk import DAG, chain, task, task_group
 else:
-    if AIRFLOW_V_3_0_PLUS:
-        from airflow.sdk import DAG, chain, task, task_group
-    else:
-        # Airflow 2.10 compat
-        from airflow.decorators import task, task_group
-        from airflow.models.baseoperator import chain
-        from airflow.models.dag import DAG
+    # Airflow 2 path
+    from airflow.decorators import task, task_group  # type: ignore[attr-defined,no-redef]
+    from airflow.models.baseoperator import chain  # type: ignore[attr-defined,no-redef]
+    from airflow.models.dag import DAG  # type: ignore[attr-defined,no-redef,assignment]
+
 from airflow.utils.trigger_rule import TriggerRule
 
 from system.amazon.aws.utils import SystemTestContextBuilder

providers/amazon/tests/system/amazon/aws/example_datasync.py

@@ -17,7 +17,6 @@
 from __future__ import annotations
 
 from datetime import datetime
-from typing import TYPE_CHECKING
 
 import boto3
 
@@ -26,18 +25,14 @@
 
 from tests_common.test_utils.version_compat import AIRFLOW_V_3_0_PLUS
 
-if TYPE_CHECKING:
-    from airflow.decorators import task
-    from airflow.models.baseoperator import chain
-    from airflow.models.dag import DAG
+if AIRFLOW_V_3_0_PLUS:
+    from airflow.sdk import DAG, chain, task
 else:
-    if AIRFLOW_V_3_0_PLUS:
-        from airflow.sdk import DAG, chain, task
-    else:
-        # Airflow 2.10 compat
-        from airflow.decorators import task
-        from airflow.models.baseoperator import chain
-        from airflow.models.dag import DAG
+    # Airflow 2 path
+    from airflow.decorators import task  # type: ignore[attr-defined,no-redef]
+    from airflow.models.baseoperator import chain  # type: ignore[attr-defined,no-redef]
+    from airflow.models.dag import DAG  # type: ignore[attr-defined,no-redef,assignment]
+
 from airflow.utils.trigger_rule import TriggerRule
 
 from system.amazon.aws.utils import ENV_ID_KEY, SystemTestContextBuilder