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

@@ -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()

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