Mypy issue with task_runner argument of flow

[david-gang]

Bug summary

Mypy does not pass on an example from the prefect website

from prefect import flow, task
from prefect.futures import wait
from prefect.task_runners import ThreadPoolTaskRunner
import time


@task
def stop_at_floor(floor):
    print(f"elevator moving to floor {floor}")
    time.sleep(floor)
    print(f"elevator stops on floor {floor}")


@flow(task_runner=ThreadPoolTaskRunner(max_workers=3))
def elevator():
    floors = []

    for floor in range(10, 0, -1):
        floors.append(stop_at_floor.submit(floor))

    wait(floors)

This is the pyproject.toml

[project]
name = "prefect-playground"
version = "0.1.0"
description = ""
readme = "README.md"
requires-python = ">=3.12"
dependencies = [
    "prefect (>=3.1.15,<4.0.0)"
]


[build-system]
requires = ["poetry-core>=2.0.0,<3.0.0"]
build-backend = "poetry.core.masonry.api"

[tool.poetry.group.dev.dependencies]
mypy = "^1.14.1"

poetry run mypy .
example.py:14: error: Argument "task_runner" to "call" of "FlowDecorator" has incompatible type "ThreadPoolTaskRunner[Never]"; >expected "TaskRunner[PrefectFuture[Never]] | None" [arg-type]
Found 1 error in 1 file (checked 1 source file)

This is python 3.12

Version info

poetry run prefect version
Version:             3.1.15
API version:         0.8.4
Python version:      3.12.8
Git commit:          3ac3d548
Built:               Thu, Jan 30, 2025 11:31 AM
OS/Arch:             darwin/arm64
Profile:             ephemeral
Server type:         ephemeral
Pydantic version:    2.11.0a1
Server:
  Database:          sqlite
  SQLite version:    3.43.2

Additional context

No response

[david-gang]

Looks connected to: #16876