From b474666787aa8d495030f9b2080c8a6b95168ca1 Mon Sep 17 00:00:00 2001
From: Ash Berlin-Taylor <ash_github@firemirror.com>
Date: Tue, 18 Jan 2022 21:17:06 +0000
Subject: [PATCH] Correctly specify overloads for TaskFlow API for type-hinting
 (#20933)

The position of `@overload` matters -- at least to MyPy, and all
function overloads have to be next to each other.

I also removed the overload for `@task.docker`, as that one _requires_
at least an image to be passed, so isn't valid with no parameters
---
 airflow/decorators/__init__.pyi                | 18 ++++++++----------
 .../howto/create-custom-decorator.rst          |  2 +-
 2 files changed, 9 insertions(+), 11 deletions(-)

diff --git a/airflow/decorators/__init__.pyi b/airflow/decorators/__init__.pyi
index 2989e2e1a9314..00b79c977252e 100644
--- a/airflow/decorators/__init__.pyi
+++ b/airflow/decorators/__init__.pyi
@@ -58,6 +58,10 @@ class TaskDecoratorFactory:
             such as transmission a large amount of XCom to TaskAPI.
         :type show_return_value_in_logs: bool
         """
+    # [START mixin_for_typing]
+    @overload
+    def python(self, python_callable: F) -> F: ...
+    # [END mixin_for_typing]
     @overload
     def __call__(
         self,
@@ -87,6 +91,8 @@ class TaskDecoratorFactory:
         :type show_return_value_in_logs: bool
         """
     @overload
+    def __call__(self, python_callable: F) -> F: ...
+    @overload
     def virtualenv(
         self,
         *,
@@ -133,6 +139,8 @@ class TaskDecoratorFactory:
             such as transmission a large amount of XCom to TaskAPI.
         :type show_return_value_in_logs: bool
         """
+    @overload
+    def virtualenv(self, python_callable: F) -> F: ...
     # [START decorator_signature]
     @overload
     def docker(
@@ -263,15 +271,5 @@ class TaskDecoratorFactory:
         :type cap_add: list[str]
         """
         # [END decorator_signature]
-    # [START mixin_for_typing]
-    @overload
-    def __call__(self, python_callable: F) -> F: ...
-    @overload
-    def python(self, python_callable: F) -> F: ...
-    @overload
-    def virtualenv(self, python_callable: F) -> F: ...
-    @overload
-    def docker(self, python_callable: F) -> F: ...
-    # [END mixin_for_typing]
 
 task = TaskDecoratorFactory()
diff --git a/docs/apache-airflow/howto/create-custom-decorator.rst b/docs/apache-airflow/howto/create-custom-decorator.rst
index de609d7998755..b156c9bc23b4c 100644
--- a/docs/apache-airflow/howto/create-custom-decorator.rst
+++ b/docs/apache-airflow/howto/create-custom-decorator.rst
@@ -105,7 +105,7 @@ automatically provided by default. All other arguments should be copied directly
 and we recommend adding a comment to explain what arguments are filled automatically by FooDecoratedOperator
 and thus not included.
 
-You should also add an overload at the end of the class similar to this so mypy can recognize the function as
+You should also add an overload that takes a single callable immediately after the "real" definition so mypy can recognize the function as
 a decorator:
 
 .. exampleinclude:: ../../../airflow/decorators/__init__.pyi