diff --git a/airflow/decorators/__init__.pyi b/airflow/decorators/__init__.pyi index 49e0cfe58a570..3a33dd254cb23 100644 --- a/airflow/decorators/__init__.pyi +++ b/airflow/decorators/__init__.pyi @@ -22,6 +22,8 @@ from typing import Any, Callable, Dict, Iterable, List, Mapping, Optional, TypeVar, Union, overload +from mypy_extensions import KwArg + from airflow.decorators.base import TaskDecorator from airflow.decorators.python import python_task from airflow.decorators.python_virtualenv import virtualenv_task @@ -226,4 +228,5 @@ class TaskDecoratorCollection: """ # [END decorator_signature] -task: TaskDecoratorCollection +DecoratedTaskType = Callable[[Function], TaskDecoratorCollection] +task: Union[DecoratedTaskType, Callable[[..., KwArg()], DecoratedTaskType]] diff --git a/airflow/providers/influxdb/example_dags/example_influxdb.py b/airflow/providers/influxdb/example_dags/example_influxdb.py index a6c160ea2f57a..387cf24d1d9a7 100644 --- a/airflow/providers/influxdb/example_dags/example_influxdb.py +++ b/airflow/providers/influxdb/example_dags/example_influxdb.py @@ -22,7 +22,7 @@ from airflow.providers.influxdb.hooks.influxdb import InfluxDBHook -@task(task_id="influxdb_task") +@task(task_id="influxdb_task") # type: ignore def test_influxdb_hook(): bucket_name = 'test-influx' influxdb_hook = InfluxDBHook()