Add stricter types for TaskGroup.start()

src/anyio/abc/_tasks.py

@@ -7,14 +7,14 @@
 from typing import TYPE_CHECKING, Any, Protocol, TypeVar, overload
 
 if sys.version_info >= (3, 11):
-    from typing import TypeVarTuple, Unpack
+    from typing import Never, TypeVarTuple, Unpack
 else:
-    from typing_extensions import TypeVarTuple, Unpack
+    from typing_extensions import Never, TypeVarTuple, Unpack
 
 if TYPE_CHECKING:
     from .._core._tasks import CancelScope
 
-T_Retval = TypeVar("T_Retval")
+T_Retval = TypeVar("T_Retval", covariant=True)
 T_contra = TypeVar("T_contra", contravariant=True)
 PosArgsT = TypeVarTuple("PosArgsT")
 
@@ -36,6 +36,13 @@ def started(self, value: T_contra | None = None) -> None:
         """
 
 
+class _StartFunc(Protocol[Unpack[PosArgsT], T_Retval]):
+    async def __call__(
+        self, *args: Unpack[PosArgsT], task_status: TaskStatus[T_Retval]
+    ) -> None:
+        ...
+
+
 class TaskGroup(metaclass=ABCMeta):
     """
     Groups several asynchronous tasks together.
@@ -63,13 +70,37 @@ def start_soon(
         .. versionadded:: 3.0
         """
 
+    @overload
+    async def start(
+        self,
+        func: _StartFunc[Unpack[PosArgsT], object],
+        *args: Unpack[PosArgsT],
+        name: object = None,
+    ) -> Never:
+        # Overload added for when the returned value is never captured.
+        # In that case, it doesn't matter which kind of function is given.
+        # Without this overload, mypy would fail:
+        #   Argument 1 to "start" of "TaskGroup" has incompatible type
+        #   "Callable[[TaskStatus[None]], Coroutine[Any, Any, None]]";
+        #   expected "_StartFunc[Never]"
+        ...
+
+    @overload
+    async def start(
+        self,
+        func: _StartFunc[Unpack[PosArgsT], T_Retval],
+        *args: Unpack[PosArgsT],
+        name: object = None,
+    ) -> T_Retval:
+        ...
+
     @abstractmethod
     async def start(
         self,
-        func: Callable[..., Awaitable[Any]],
-        *args: object,
+        func: _StartFunc[Unpack[PosArgsT], T_Retval],
+        *args: Unpack[PosArgsT],
         name: object = None,
-    ) -> Any:
+    ) -> T_Retval:
         """
         Start a new task and wait until it signals for readiness.
 

src/anyio/streams/stapled.py

@@ -120,7 +120,9 @@ def __post_init__(self) -> None:
         self.listeners = listeners
 
     async def serve(
-        self, handler: Callable[[T_Stream], Any], task_group: TaskGroup | None = None
+        self,
+        handler: Callable[[T_Stream], Any],
+        task_group: TaskGroup | None = None,
     ) -> None:
         from .. import create_task_group
 

tests/test_debugging.py

@@ -62,7 +62,7 @@ async def main() -> None:
 async def test_non_main_task_name(
     name_input: bytes | str | None, expected: str
 ) -> None:
-    async def non_main(*, task_status: TaskStatus) -> None:
+    async def non_main(*, task_status: TaskStatus[str | None]) -> None:
         task_status.started(anyio.get_current_task().name)
 
     async with anyio.create_task_group() as tg:

tests/test_taskgroups.py

@@ -101,16 +101,16 @@ async def test_start_soon_after_error() -> None:
 
 
 async def test_start_no_value() -> None:
-    async def taskfunc(*, task_status: TaskStatus) -> None:
+    async def taskfunc(*, task_status: TaskStatus[None]) -> None:
         task_status.started()
 
     async with create_task_group() as tg:
-        value = await tg.start(taskfunc)
+        value: None = await tg.start(taskfunc)
         assert value is None
 
 
 async def test_start_called_twice() -> None:
-    async def taskfunc(*, task_status: TaskStatus) -> None:
+    async def taskfunc(*, task_status: TaskStatus[None]) -> None:
         task_status.started()
 
         with pytest.raises(
@@ -119,12 +119,12 @@ async def taskfunc(*, task_status: TaskStatus) -> None:
             task_status.started()
 
     async with create_task_group() as tg:
-        value = await tg.start(taskfunc)
+        value: None = await tg.start(taskfunc)
         assert value is None
 
 
 async def test_start_with_value() -> None:
-    async def taskfunc(*, task_status: TaskStatus) -> None:
+    async def taskfunc(*, task_status: TaskStatus[str]) -> None:
         task_status.started("foo")
 
     async with create_task_group() as tg:
@@ -144,7 +144,7 @@ async def taskfunc(*, task_status: TaskStatus) -> NoReturn:
 
 
 async def test_start_crash_after_started_call() -> None:
-    async def taskfunc(*, task_status: TaskStatus) -> NoReturn:
+    async def taskfunc(*, task_status: TaskStatus[int]) -> NoReturn:
         task_status.started(2)
         raise Exception("foo")
 
@@ -250,7 +250,7 @@ async def taskfunc() -> None:
 
 
 async def test_start_exception_delivery(anyio_backend_name: str) -> None:
-    def task_fn(*, task_status: TaskStatus = TASK_STATUS_IGNORED) -> None:
+    def task_fn(*, task_status: TaskStatus[str] = TASK_STATUS_IGNORED) -> None:
         task_status.started("hello")
 
     if anyio_backend_name == "trio":