Make @task import from airflow.sdk (#48896)

closes #48811

Author: kaxil

.pre-commit-config.yaml

@@ -646,7 +646,7 @@ repos:
         language: python
         entry: ./scripts/ci/pre_commit/check_init_decorator_arguments.py
         pass_filenames: false
-        files: ^airflow-core/src/airflow/models/dag\.py$|^airflow-core/src/airflow/(?:decorators|utils)/task_group\.py$
+        files: ^task-sdk/src/airflow/sdk/definitions/dag\.py$|^task-sdk/src/airflow/sdk/definitions/decorators/task_group\.py$
         additional_dependencies: ['rich>=12.4.4']
       - id: check-template-context-variable-in-sync
         name: Sync template context variable refs

airflow-core/docs/authoring-and-scheduling/dynamic-task-mapping.rst

@@ -254,7 +254,7 @@ In this example, you have a regular data delivery to an S3 bucket and want to ap
     from datetime import datetime
 
     from airflow.sdk import DAG
-    from airflow.decorators import task
+    from airflow.sdk import task
     from airflow.providers.amazon.aws.hooks.s3 import S3Hook
     from airflow.providers.amazon.aws.operators.s3 import S3ListOperator
 

airflow-core/docs/best-practices.rst

@@ -125,7 +125,7 @@ Not avoiding top-level DAG code:
   import pendulum
 
   from airflow.sdk import DAG
-  from airflow.decorators import task
+  from airflow.sdk import task
 
 
   def expensive_api_call():
@@ -154,7 +154,7 @@ Avoiding top-level DAG code:
   import pendulum
 
   from airflow.sdk import DAG
-  from airflow.decorators import task
+  from airflow.sdk import task
 
 
   def expensive_api_call():
@@ -536,7 +536,7 @@ It's easier to grab the concept with an example. Let's say that we have the foll
     from datetime import datetime
 
     from airflow.sdk import DAG
-    from airflow.decorators import task
+    from airflow.sdk import task
     from airflow.exceptions import AirflowException
     from airflow.providers.standard.operators.bash import BashOperator
     from airflow.utils.trigger_rule import TriggerRule

airflow-core/docs/core-concepts/dags.rst

@@ -83,7 +83,7 @@ Or, you can use the ``@dag`` decorator to :ref:`turn a function into a DAG gener
 
     import datetime
 
-    from airflow.decorators import dag
+    from airflow.sdk import dag
     from airflow.providers.standard.operators.empty import EmptyOperator
 
 
@@ -458,7 +458,7 @@ You can also combine this with the :ref:`concepts:depends-on-past` functionality
         # dags/branch_without_trigger.py
         import pendulum
 
-        from airflow.decorators import task
+        from airflow.sdk import task
         from airflow.sdk import DAG
         from airflow.providers.standard.operators.empty import EmptyOperator
 
@@ -557,7 +557,7 @@ Dependency relationships can be applied across all tasks in a TaskGroup with the
 .. code-block:: python
    :emphasize-lines: 4,12
 
-    from airflow.decorators import task_group
+    from airflow.sdk import task_group
 
 
     @task_group()
@@ -578,7 +578,7 @@ TaskGroup also supports ``default_args`` like DAG, it will overwrite the ``defau
     import datetime
 
     from airflow.sdk import DAG
-    from airflow.decorators import task_group
+    from airflow.sdk import task_group
     from airflow.providers.standard.operators.bash import BashOperator
     from airflow.providers.standard.operators.empty import EmptyOperator
 

airflow-core/docs/core-concepts/params.rst

@@ -39,7 +39,7 @@ Use a dictionary that maps Param names to either a :class:`~airflow.sdk.definiti
    :emphasize-lines: 7-10
 
     from airflow.sdk import DAG
-    from airflow.decorators import task
+    from airflow.sdk import task
     from airflow.sdk import Param
 
     with DAG(

airflow-core/docs/core-concepts/taskflow.rst

@@ -24,7 +24,7 @@ If you write most of your dags using plain Python code rather than Operators, th
 
 TaskFlow takes care of moving inputs and outputs between your Tasks using XComs for you, as well as automatically calculating dependencies - when you call a TaskFlow function in your DAG file, rather than executing it, you will get an object representing the XCom for the result (an ``XComArg``), that you can then use as inputs to downstream tasks or operators. For example::
 
-    from airflow.decorators import task
+    from airflow.sdk import task
     from airflow.providers.email import EmailOperator
 
     @task
@@ -112,7 +112,7 @@ a ``Asset``, which is ``@attr.define`` decorated, together with TaskFlow.
     import requests
 
     from airflow import Asset
-    from airflow.decorators import dag, task
+    from airflow.sdk import dag, task
 
     SRC = Asset(
         "https://www.ncei.noaa.gov/access/monitoring/climate-at-a-glance/global/time-series/globe/land_ocean/ytd/12/1880-2022.json"

airflow-core/docs/core-concepts/tasks.rst

@@ -200,7 +200,7 @@ If you'd like to reproduce task instance heartbeat timeouts for development/test
 
 .. code-block:: python
 
-    from airflow.decorators import dag
+    from airflow.sdk import dag
     from airflow.providers.standard.operators.bash import BashOperator
     from datetime import datetime
 

airflow-core/docs/faq.rst

@@ -182,7 +182,7 @@ until ``min_file_process_interval`` is reached since DAG Parser will look for mo
    :name: dag_loader.py
 
     from airflow.sdk import DAG
-    from airflow.decorators import task
+    from airflow.sdk import task
 
     import pendulum
 
@@ -413,7 +413,7 @@ upstream task.
 
     import pendulum
 
-    from airflow.decorators import dag, task
+    from airflow.sdk import dag, task
     from airflow.exceptions import AirflowException
     from airflow.utils.trigger_rule import TriggerRule
 

airflow-core/docs/howto/create-custom-decorator.rst

@@ -43,10 +43,10 @@ tasks. The steps to create and register ``@task.foo`` are:
     .. code-block:: python
 
         from typing import TYPE_CHECKING
-        from airflow.decorators.base import task_decorator_factory
+        from airflow.sdk.bases.decorator import task_decorator_factory
 
         if TYPE_CHECKING:
-            from airflow.decorators.base import TaskDecorator
+            from airflow.sdk.bases.decorator import TaskDecorator
 
 
         def foo_task(
@@ -96,11 +96,11 @@ tasks. The steps to create and register ``@task.foo`` are:
 For better or worse, Python IDEs can not auto-complete dynamically
 generated methods (see `JetBrain's write up on the subject <https://intellij-support.jetbrains.com/hc/en-us/community/posts/115000665110-auto-completion-for-dynamic-module-attributes-in-python>`_).
 
-To hack around this problem, a type stub ``airflow/decorators/__init__.pyi`` is provided to statically declare
+To hack around this problem, a type stub ``airflow/sdk/definitions/decorators/__init__.pyi`` is provided to statically declare
 the type signature of each task decorator. A newly added task decorator should declare its signature stub
 like this:
 
-.. exampleinclude:: /../src/airflow/decorators/__init__.pyi
+.. exampleinclude:: ../../../task-sdk/src/airflow/sdk/definitions/decorators/__init__.pyi
     :language: python
     :start-after: [START decorator_signature]
     :end-before: [END decorator_signature]
@@ -114,7 +114,7 @@ If the new decorator can be used without arguments (e.g. ``@task.python`` instea
 You should also add an overload that takes a single callable immediately after the "real" definition so mypy
 can recognize the function as a "bare decorator":
 
-.. exampleinclude:: /../src/airflow/decorators/__init__.pyi
+.. exampleinclude:: ../../../task-sdk/src/airflow/sdk/definitions/decorators/__init__.pyi
     :language: python
     :start-after: [START mixin_for_typing]
     :end-before: [END mixin_for_typing]

airflow-core/docs/howto/dynamic-dag-generation.rst

@@ -127,7 +127,7 @@ and Airflow will automatically register them.
 .. code-block:: python
 
     from datetime import datetime
-    from airflow.decorators import dag, task
+    from airflow.sdk import dag, task
 
     configs = {
         "config1": {"message": "first DAG will receive this message"},