Typing issues for @task after change between 3.1.14 and 3.1.15

[ogenstad]

Bug summary

Between Prefect 3.1.14 and 3.1.15 there was a change in the overload definitions for @task which resulted mypy issues within our codebase.

The code looks like this:

class InfrahubRepositoryIntegrator(InfrahubRepositoryBase):


    @flow(name="import-python-files", flow_run_name="Import Python file")
    async def import_all_python_files(
        self, branch_name: str, commit: str, config_file: InfrahubRepositoryConfig
    ) -> None:
        await add_tags(branches=[branch_name], nodes=[str(self.id)])

        await self.import_python_check_definitions(branch_name=branch_name, commit=commit, config_file=config_file)



    @task(name="import-python-check-definitions", task_run_name="Import Python Check Definitions", cache_policy=NONE)
    async def import_python_check_definitions(
        self, branch_name: str, commit: str, config_file: InfrahubRepositoryConfig
    ) -> None:

The mypy errors we are are seeing from 3.1.15 and onward look like this:

backend/infrahub/git/integrator.py:1073: error: Incompatible types in "await" (actual type "State[Coroutine[Any, Any, None]]", expected type "Awaitable[Any]")  [misc]
            await self.import_python_check_definitions(branch_name=branch_name, commit=commit, config_file=config_file)
            ^~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~

The change was introduced in this PR: #16891.

git diff 3.1.14 3.1.15 -- src/prefect/tasks.py

diff --git a/src/prefect/tasks.py b/src/prefect/tasks.py
index 7763722573..9b291f7a66 100644
--- a/src/prefect/tasks.py
+++ b/src/prefect/tasks.py
@@ -965,7 +965,7 @@ class Task(Generic[P, R]):
     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 @@ class Task(Generic[P, R]):
     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]:
         ...

If I switch the order of these overloads it works as expected for us, however I don't know if that would break anything else. There's also a TODO comment here indicating that additional work needs to be done.

    # 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[True] = True,
        wait_for: Optional[OneOrManyFutureOrResult[Any]] = None,
        **kwargs: P.kwargs,
    ) -> State[R]: ...

    @overload
    def __call__(
        self: "Task[P, R]",
        *args: P.args,
        return_state: Literal[False] = False,
        wait_for: Optional[OneOrManyFutureOrResult[Any]] = None,
        **kwargs: P.kwargs,
    ) -> R: ...

Version info

❯ prefect version
Version:             3.2.9
API version:         0.8.4
Python version:      3.12.8
Git commit:          27eb408c
Built:               Fri, Feb 28, 2025 8:12 PM
OS/Arch:             darwin/arm64
Profile:             ephemeral
Server type:         server
Pydantic version:    2.10.6
Integrations:
  prefect-redis:     0.2.2

Additional context

No response