From e95b5df87afbf88bdc5bf8076102e1cb86c32a77 Mon Sep 17 00:00:00 2001 From: Daniel Standish <15932138+dstandish@users.noreply.github.com> Date: Mon, 7 Feb 2022 12:24:12 -0800 Subject: [PATCH 1/7] Fix type annotations re task decorator map --- airflow/decorators/__init__.pyi | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/airflow/decorators/__init__.pyi b/airflow/decorators/__init__.pyi index 49e0cfe58a570..e2fc635b6b81e 100644 --- a/airflow/decorators/__init__.pyi +++ b/airflow/decorators/__init__.pyi @@ -226,4 +226,5 @@ class TaskDecoratorCollection: """ # [END decorator_signature] -task: TaskDecoratorCollection +T = TypeVar("T", bound=TaskDecoratorCollection) +task: Callable[[Function], T] From 67e2f0f5a180261d67c7c67872f0beec4cbb8517 Mon Sep 17 00:00:00 2001 From: Daniel Standish <15932138+dstandish@users.noreply.github.com> Date: Mon, 7 Feb 2022 12:52:03 -0800 Subject: [PATCH 2/7] try fix --- airflow/decorators/__init__.pyi | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/airflow/decorators/__init__.pyi b/airflow/decorators/__init__.pyi index e2fc635b6b81e..3d73e7b03920b 100644 --- a/airflow/decorators/__init__.pyi +++ b/airflow/decorators/__init__.pyi @@ -226,5 +226,5 @@ class TaskDecoratorCollection: """ # [END decorator_signature] -T = TypeVar("T", bound=TaskDecoratorCollection) -task: Callable[[Function], T] +DecoratedType = Callable[[Function], TaskDecoratorCollection] +task: Union[DecoratedType, Callable[[Any], DecoratedType]] From 5c227715dd6bead2177bbf53f8df5ea5c9ae452a Mon Sep 17 00:00:00 2001 From: Daniel Standish <15932138+dstandish@users.noreply.github.com> Date: Mon, 7 Feb 2022 13:02:41 -0800 Subject: [PATCH 3/7] try other fix --- airflow/decorators/__init__.pyi | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/airflow/decorators/__init__.pyi b/airflow/decorators/__init__.pyi index 3d73e7b03920b..b8e750e493841 100644 --- a/airflow/decorators/__init__.pyi +++ b/airflow/decorators/__init__.pyi @@ -226,5 +226,5 @@ class TaskDecoratorCollection: """ # [END decorator_signature] -DecoratedType = Callable[[Function], TaskDecoratorCollection] -task: Union[DecoratedType, Callable[[Any], DecoratedType]] +DecoratedTaskType = Callable[[Function], TaskDecoratorCollection] +task: Union[DecoratedTaskType, Callable[[..., Any], DecoratedTaskType]] From 3bc17bd2e76446998a2e2e5f32821261177be868 Mon Sep 17 00:00:00 2001 From: Daniel Standish <15932138+dstandish@users.noreply.github.com> Date: Mon, 7 Feb 2022 13:24:54 -0800 Subject: [PATCH 4/7] try kwarg --- airflow/decorators/__init__.pyi | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/airflow/decorators/__init__.pyi b/airflow/decorators/__init__.pyi index b8e750e493841..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 @@ -227,4 +229,4 @@ class TaskDecoratorCollection: # [END decorator_signature] DecoratedTaskType = Callable[[Function], TaskDecoratorCollection] -task: Union[DecoratedTaskType, Callable[[..., Any], DecoratedTaskType]] +task: Union[DecoratedTaskType, Callable[[..., KwArg()], DecoratedTaskType]] From c68fa59736c8afb0637b36c6324685ddd90e8287 Mon Sep 17 00:00:00 2001 From: Daniel Standish <15932138+dstandish@users.noreply.github.com> Date: Mon, 7 Feb 2022 13:53:08 -0800 Subject: [PATCH 5/7] try using protocols --- airflow/decorators/__init__.pyi | 10 ++++++---- 1 file changed, 6 insertions(+), 4 deletions(-) diff --git a/airflow/decorators/__init__.pyi b/airflow/decorators/__init__.pyi index 3a33dd254cb23..390359a71052e 100644 --- a/airflow/decorators/__init__.pyi +++ b/airflow/decorators/__init__.pyi @@ -20,9 +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, TypeVar, Union, overload - -from mypy_extensions import KwArg +from typing import Any, Callable, Dict, Iterable, List, Mapping, Optional, Protocol, TypeVar, Union, overload from airflow.decorators.base import TaskDecorator from airflow.decorators.python import python_task @@ -229,4 +227,8 @@ class TaskDecoratorCollection: # [END decorator_signature] DecoratedTaskType = Callable[[Function], TaskDecoratorCollection] -task: Union[DecoratedTaskType, Callable[[..., KwArg()], DecoratedTaskType]] + +class DecoratorFuncType(Protocol): + def __call__(self, *args, **kwargs) -> DecoratedTaskType: ... + +task: Union[DecoratedTaskType, DecoratorFuncType] From 31ea94ec7e7999ea7886c4d2bed5fb5dc9d08fdc Mon Sep 17 00:00:00 2001 From: Daniel Standish <15932138+dstandish@users.noreply.github.com> Date: Mon, 7 Feb 2022 14:03:54 -0800 Subject: [PATCH 6/7] revert protocol change --- airflow/decorators/__init__.pyi | 10 ++++------ 1 file changed, 4 insertions(+), 6 deletions(-) diff --git a/airflow/decorators/__init__.pyi b/airflow/decorators/__init__.pyi index 390359a71052e..3a33dd254cb23 100644 --- a/airflow/decorators/__init__.pyi +++ b/airflow/decorators/__init__.pyi @@ -20,7 +20,9 @@ # necessarily exist at run time. See "Creating Custom @task Decorators" # documentation for more details. -from typing import Any, Callable, Dict, Iterable, List, Mapping, Optional, Protocol, TypeVar, Union, overload +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 @@ -227,8 +229,4 @@ class TaskDecoratorCollection: # [END decorator_signature] DecoratedTaskType = Callable[[Function], TaskDecoratorCollection] - -class DecoratorFuncType(Protocol): - def __call__(self, *args, **kwargs) -> DecoratedTaskType: ... - -task: Union[DecoratedTaskType, DecoratorFuncType] +task: Union[DecoratedTaskType, Callable[[..., KwArg()], DecoratedTaskType]] From d7b73dbe35039aecb0d3c6cbd8262bbb6901211b Mon Sep 17 00:00:00 2001 From: Daniel Standish <15932138+dstandish@users.noreply.github.com> Date: Mon, 7 Feb 2022 14:06:20 -0800 Subject: [PATCH 7/7] ignore type --- airflow/providers/influxdb/example_dags/example_influxdb.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) 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()