Add missing overload for Task.__call__ (#16891)

src/prefect/tasks.py

@@ -965,7 +965,7 @@ async def create_local_run(
     def __call__(
         self: "Task[P, NoReturn]",
         *args: P.args,
-        return_state: Literal[False],
+        return_state: Literal[False] = False,
         wait_for: Optional[OneOrManyFutureOrResult[Any]] = None,
         **kwargs: P.kwargs,
     ) -> None:
@@ -977,20 +977,22 @@ def __call__(
     def __call__(
         self: "Task[P, R]",
         *args: P.args,
-        return_state: Literal[True],
-        wait_for: Optional[OneOrManyFutureOrResult[Any]] = None,
         **kwargs: P.kwargs,
-    ) -> State[R]:
+    ) -> R:
         ...
 
+    # Keyword parameters `return_state` and `wait_for` aren't allowed after the
+    # ParamSpec `*args` parameter, so we lose return type typing when either of
+    # those are provided.
+    # TODO: Find a way to expose this functionality without losing type information
     @overload
     def __call__(
         self: "Task[P, R]",
         *args: P.args,
-        return_state: Literal[False],
+        return_state: Literal[True] = True,
         wait_for: Optional[OneOrManyFutureOrResult[Any]] = None,
         **kwargs: P.kwargs,
-    ) -> R:
+    ) -> State[R]:
         ...
 
     @overload

tests/typesafety/test_flows.yml

@@ -34,4 +34,16 @@
     reveal_type(foo)
   out: "main:5: note: Revealed type is \"\
       prefect.flows.Flow[[bar: builtins.str], builtins.int]\
-    \""
+    \""
+
+- case: prefect_flow_call
+  main: |
+    from prefect import flow
+    @flow
+    def foo(bar: str) -> int:
+        return 42
+    ret = foo(bar="baz")
+    reveal_type(ret)
+  out: "main:6: note: Revealed type is \"\
+      builtins.int\
+    \""

tests/typesafety/test_tasks.yml

@@ -56,4 +56,16 @@
     reveal_type(foo)
   out: "main:9: note: Revealed type is \"\
       prefect.tasks.Task[[bar: builtins.str], builtins.int]\
-    \""
+    \""
+
+- case: prefect_task_call
+  main: |
+    from prefect import task
+    @task
+    def foo(bar: str) -> int:
+        return 42
+    ret = foo(bar="baz")
+    reveal_type(ret)
+  out: "main:6: note: Revealed type is \"\
+      builtins.int\
+    \""