diff --git a/airflow/decorators/__init__.py b/airflow/decorators/__init__.py index 40d9921ec0169..f375fbf7167ba 100644 --- a/airflow/decorators/__init__.py +++ b/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``.""" diff --git a/airflow/decorators/__init__.pyi b/airflow/decorators/__init__.pyi index e970f61657c64..726798f1341e3 100644 --- a/airflow/decorators/__init__.pyi +++ b/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]: ... diff --git a/airflow/decorators/branch_python.py b/airflow/decorators/branch_python.py index ac83132574402..b5876f49c75b7 100644 --- a/airflow/decorators/branch_python.py +++ b/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: diff --git a/airflow/decorators/python.py b/airflow/decorators/python.py index c7c1a629a4dfb..3f00681ccfdde 100644 --- a/airflow/decorators/python.py +++ b/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, diff --git a/airflow/decorators/python_virtualenv.py b/airflow/decorators/python_virtualenv.py index e8fd681b8f3af..dba306776c83d 100644 --- a/airflow/decorators/python_virtualenv.py +++ b/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, diff --git a/airflow/providers/docker/decorators/docker.py b/airflow/providers/docker/decorators/docker.py index db3293a5f569f..b07bb6cdff59d 100644 --- a/airflow/providers/docker/decorators/docker.py +++ b/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,