Fix typing of task decorator for retry_condition_fn argument

[peterbygrave]

In migrating to prefect v3, we are seeing issues with mypy checking of the task decorators. For example when making the decorator a call e.g. @task() or @task(name="foo") you end up with:

error: Argument 1 has incompatible type "Callable[[int, int], int]"; expected "Callable[[VarArg(Never), KwArg(Never)], int]"  [arg-type]

I found that it was the typing of retry_condition_fn that flakes out because the input typing should connected to the return type, but nothing can be inferred because nothing was given.

In this PR I add an overload to differentiate between None and an actual function is passed.

Checklist

• This pull request references any related issue by including "closes <link to issue>"
  • If no issue exists and your change is not a small fix, please create an issue first.
• If this pull request adds new functionality, it includes unit tests that cover the changes
• If this pull request removes docs files, it includes redirect settings in mint.json.
• If this pull request adds functions or classes, it includes helpful docstrings.

[peterbygrave]

Suggested change

	retry_condition_fn: Optional[Callable[[Task[P, R], TaskRun, State], bool]] = None,
	retry_condition_fn: Optional[Callable[[Task[P, R], TaskRun, State], bool]] = None,

Note this previously was Task[P, Any] but now has R because it should exactly match. The overload above now matches first when None is given (or default is used).

[peterbygrave]

This is the new thing that matches first when it is None and doesn't interact with returned P and R.

[peterbygrave]

Works on main branch

[peterbygrave]

These 3 all fail on main branch

[codspeed-hq]

CodSpeed Performance Report

Merging #16621 will not alter performance

_{Comparing peterbygrave:fix-tasks-overload-retry-condition-fn (77968bf) with main (8b8eea2)}

Summary

✅ 2 untouched benchmarks

[peterbygrave]

Works on main branch

[zzstoatzz]

hi @peterbygrave - thanks for the PR! can you provide your version of prefect and a specific code example?

for example, I am not seeing pyright have any issues with this on main, though I can see mypy warns as you mention

from typing import Any, reveal_type

from prefect import Task, task
from prefect.cache_policies import NONE
from prefect.client.schemas.objects import State, TaskRun


def state_hook(task: Task[Any, Any], task_run: TaskRun, state: State):
    print(f"Task run {task_run.name!r} from {task.name!r} entered state {state.name!r}")


@task(retries=1, persist_result=False, on_completion=[state_hook])
def identity(x: int) -> int:
    return x


@task(persist_result=False)
def identity_no_cache(x: int) -> int:
    return x


@task(cache_policy=NONE)
def returns_something(x: int) -> int:
    return x


@task(
    name="conditionally_retries",
    retry_condition_fn=lambda task, task_run, state: True,
    retry_delay_seconds=[1, 2],
)
def conditionally_retries(x: int) -> int:
    return x


@task
def also_returns_something(x: int) -> int:
    return x


if __name__ == "__main__":
    value = returns_something(1)
    something = identity(value)
    something_else = identity_no_cache(value)
    other_value = also_returns_something(1)
    result = conditionally_retries(1)

    reveal_type(value)
    reveal_type(something)
    reveal_type(something_else)
    reveal_type(other_value)
    reveal_type(result)

when I checkout your branch, this result (from the task with a retry_condition_fn) becomes Unknown

[peterbygrave]

hi @peterbygrave - thanks for the PR! can you provide your version of prefect and a specific code example?

We are just doing our v2 to v3 migration so first saw this on v3.1.11, but has issues on main branch.

The code example I have are in the tests. Its all isolated to retry_condition_fn because it is capable of narrowing the P and R.

for example, I am not seeing pyright have any issues with this on main, though I can see mypy warns as you mention

from typing import Any, reveal_type

from prefect import Task, task
from prefect.cache_policies import NONE
from prefect.client.schemas.objects import State, TaskRun


def state_hook(task: Task[Any, Any], task_run: TaskRun, state: State):
    print(f"Task run {task_run.name!r} from {task.name!r} entered state {state.name!r}")


@task(retries=1, persist_result=False, on_completion=[state_hook])
def identity(x: int) -> int:
    return x


@task(persist_result=False)
def identity_no_cache(x: int) -> int:
    return x


@task(cache_policy=NONE)
def returns_something(x: int) -> int:
    return x


@task(
    name="conditionally_retries",
    retry_condition_fn=lambda task, task_run, state: True,
    retry_delay_seconds=[1, 2],
)
def conditionally_retries(x: int) -> int:
    return x


@task
def also_returns_something(x: int) -> int:
    return x


if __name__ == "__main__":
    value = returns_something(1)
    something = identity(value)
    something_else = identity_no_cache(value)
    other_value = also_returns_something(1)
    result = conditionally_retries(1)

    reveal_type(value)
    reveal_type(something)
    reveal_type(something_else)
    reveal_type(other_value)
    reveal_type(result)

when I checkout your branch, this result (from the task with a retry_condition_fn) becomes Unknown

So the on_completion arg type does not try to infer the P and R generics (its just Task, no generics) for the returned Task type. But retry_condition_fn does (because it has Task[P, R]. So in your example you pass a lambda in which has no typing at all, which on first pass through the decorator destroys the typing in the returned partial. So when the function is passed through the decorator/partial again it can't infer the P and R of the return Task anymore. My guess would be if you want to support lambdas, then you'd need to modify the typing to something like:

retry_condition_fn: Optional[Callable[[Task, TaskRun, State], bool]]

I'm not sure if that would need to be in additional or replacement of. If replacement it would probably exclude the ability to have stronger type checking - in that one wouldn't be able to check you are handling a task with the right args and return type (if you wanted to depend on them).