Clean up task decorator type hints and docstrings (#24667)

airflow/decorators/__init__.py

@@ -41,11 +41,11 @@
 class TaskDecoratorCollection:
     """Implementation to provide the ``@task`` syntax."""
 
-    python: Any = staticmethod(python_task)
+    python = staticmethod(python_task)
     virtualenv = staticmethod(virtualenv_task)
     branch = staticmethod(branch_task)
 
-    __call__ = python  # Alias '@task' to '@task.python'.
+    __call__: Any = python  # Alias '@task' to '@task.python'.
 
     def __getattr__(self, name: str) -> TaskDecorator:
         """Dynamically get provider-registered task decorators, e.g. ``@task.docker``."""

airflow/decorators/__init__.pyi

@@ -20,7 +20,7 @@
 # necessarily exist at run time. See "Creating Custom @task Decorators"
 # documentation for more details.
 
-from typing import Any, Callable, Dict, Iterable, List, Mapping, Optional, Union, overload
+from typing import Any, Dict, Iterable, List, Mapping, Optional, Union, overload
 
 from airflow.decorators.base import Function, Task, TaskDecorator
 from airflow.decorators.branch_python import branch_task
@@ -60,7 +60,7 @@ class TaskDecoratorCollection:
         :param templates_dict: a dictionary where the values are templates that
             will get templated by the Airflow engine sometime between
             ``__init__`` and ``execute`` takes place and are made available
-            in your callable's context after the template has been applied
+            in your callable's context after the template has been applied.
         :param show_return_value_in_logs: a bool value whether to show return_value
             logs. Defaults to True, which allows return value log output.
             It can be set to False to prevent log output of return value when you return huge data
@@ -115,7 +115,7 @@ class TaskDecoratorCollection:
         :param templates_dict: a dictionary where the values are templates that
             will get templated by the Airflow engine sometime between
             ``__init__`` and ``execute`` takes place and are made available
-            in your callable's context after the template has been applied
+            in your callable's context after the template has been applied.
         :param show_return_value_in_logs: a bool value whether to show return_value
             logs. Defaults to True, which allows return value log output.
             It can be set to False to prevent log output of return value when you return huge data
@@ -124,20 +124,14 @@ class TaskDecoratorCollection:
     @overload
     def virtualenv(self, python_callable: Function) -> Task[Function]: ...
     @overload
-    def branch(
-        self, python_callable: Optional[Callable] = None, multiple_outputs: Optional[bool] = None, **kwargs
-    ) -> TaskDecorator:
-        """Wraps a python function into a BranchPythonOperator
+    def branch(self, *, multiple_outputs: Optional[bool] = None, **kwargs) -> TaskDecorator:
+        """Create a decorator to wrap the decorated callable into a BranchPythonOperator.
+
+        For more information on how to use this decorator, see :ref:`howto/operator:BranchPythonOperator`.
+        Accepts arbitrary for operator kwarg. Can be reused in a single DAG.
 
-        For more information on how to use this operator, take a look at the guide:
-        :ref:`howto/operator:BranchPythonOperator`
-        Accepts kwargs for operator kwarg. Can be reused in a single DAG.
-        :param python_callable: Function to decorate
-        :type python_callable: Optional[Callable]
-        :param multiple_outputs: if set, function return value will be
-            unrolled to multiple XCom values. Dict will unroll to xcom values with keys as XCom keys.
-            Defaults to False.
-        :type multiple_outputs: bool
+        :param multiple_outputs: If set, function return value will be unrolled to multiple XCom values.
+            Dict will unroll to XCom values with keys as XCom keys. Defaults to False.
         """
     @overload
     def branch(self, python_callable: Function) -> Task[Function]: ...

airflow/decorators/branch_python.py

@@ -17,7 +17,7 @@
 
 import inspect
 from textwrap import dedent
-from typing import Callable, Optional, Sequence, TypeVar
+from typing import Callable, Optional, Sequence
 
 from airflow.decorators.base import DecoratedOperator, TaskDecorator, task_decorator_factory
 from airflow.operators.python import BranchPythonOperator
@@ -63,9 +63,6 @@ def get_python_source(self):
         return res
 
 
-T = TypeVar("T", bound=Callable)
-
-
 def branch_task(
     python_callable: Optional[Callable] = None, multiple_outputs: Optional[bool] = None, **kwargs
 ) -> TaskDecorator:

airflow/decorators/python.py

@@ -15,7 +15,7 @@
 # specific language governing permissions and limitations
 # under the License.
 
-from typing import Callable, Optional, Sequence, TypeVar
+from typing import Callable, Optional, Sequence
 
 from airflow.decorators.base import DecoratedOperator, TaskDecorator, task_decorator_factory
 from airflow.operators.python import PythonOperator
@@ -56,9 +56,6 @@ def __init__(self, *, python_callable, op_args, op_kwargs, **kwargs) -> None:
         )
 
 
-T = TypeVar("T", bound=Callable)
-
-
 def python_task(
     python_callable: Optional[Callable] = None,
     multiple_outputs: Optional[bool] = None,

airflow/decorators/python_virtualenv.py

@@ -17,7 +17,7 @@
 
 import inspect
 from textwrap import dedent
-from typing import Callable, Optional, Sequence, TypeVar
+from typing import Callable, Optional, Sequence
 
 from airflow.decorators.base import DecoratedOperator, TaskDecorator, task_decorator_factory
 from airflow.operators.python import PythonVirtualenvOperator
@@ -65,9 +65,6 @@ def get_python_source(self):
         return res
 
 
-T = TypeVar("T", bound=Callable)
-
-
 def virtualenv_task(
     python_callable: Optional[Callable] = None,
     multiple_outputs: Optional[bool] = None,

airflow/providers/docker/decorators/docker.py

@@ -21,7 +21,7 @@
 import pickle
 from tempfile import TemporaryDirectory
 from textwrap import dedent
-from typing import TYPE_CHECKING, Callable, Optional, Sequence, TypeVar
+from typing import TYPE_CHECKING, Callable, Optional, Sequence
 
 import dill
 
@@ -129,9 +129,6 @@ def _get_python_source(self):
         return res
 
 
-T = TypeVar("T", bound=Callable)
-
-
 def docker_task(
     python_callable: Optional[Callable] = None,
     multiple_outputs: Optional[bool] = None,