Make ExitStack, AbstractContextManager and AsyncAbstractContextManager generic in return type of __exit__ (#11048)

stdlib/contextlib.pyi

@@ -31,32 +31,33 @@ if sys.version_info >= (3, 11):
 _T = TypeVar("_T")
 _T_co = TypeVar("_T_co", covariant=True)
 _T_io = TypeVar("_T_io", bound=IO[str] | None)
+_ExitT_co = TypeVar("_ExitT_co", covariant=True, bound=bool | None, default=bool | None)
 _F = TypeVar("_F", bound=Callable[..., Any])
 _P = ParamSpec("_P")
 
 _ExitFunc: TypeAlias = Callable[[type[BaseException] | None, BaseException | None, TracebackType | None], bool | None]
-_CM_EF = TypeVar("_CM_EF", bound=AbstractContextManager[Any] | _ExitFunc)
+_CM_EF = TypeVar("_CM_EF", bound=AbstractContextManager[Any, Any] | _ExitFunc)
 
 @runtime_checkable
-class AbstractContextManager(Protocol[_T_co]):
+class AbstractContextManager(Protocol[_T_co, _ExitT_co]):
     def __enter__(self) -> _T_co: ...
     @abstractmethod
     def __exit__(
         self, exc_type: type[BaseException] | None, exc_value: BaseException | None, traceback: TracebackType | None, /
-    ) -> bool | None: ...
+    ) -> _ExitT_co: ...
 
 @runtime_checkable
-class AbstractAsyncContextManager(Protocol[_T_co]):
+class AbstractAsyncContextManager(Protocol[_T_co, _ExitT_co]):
     async def __aenter__(self) -> _T_co: ...
     @abstractmethod
     async def __aexit__(
         self, exc_type: type[BaseException] | None, exc_value: BaseException | None, traceback: TracebackType | None, /
-    ) -> bool | None: ...
+    ) -> _ExitT_co: ...
 
 class ContextDecorator:
     def __call__(self, func: _F) -> _F: ...
 
-class _GeneratorContextManager(AbstractContextManager[_T_co], ContextDecorator):
+class _GeneratorContextManager(AbstractContextManager[_T_co, bool | None], ContextDecorator):
     # __init__ and all instance attributes are actually inherited from _GeneratorContextManagerBase
     # _GeneratorContextManagerBase is more trouble than it's worth to include in the stub; see #6676
     def __init__(self, func: Callable[..., Iterator[_T_co]], args: tuple[Any, ...], kwds: dict[str, Any]) -> None: ...
@@ -81,7 +82,7 @@ if sys.version_info >= (3, 10):
     class AsyncContextDecorator:
         def __call__(self, func: _AF) -> _AF: ...
 
-    class _AsyncGeneratorContextManager(AbstractAsyncContextManager[_T_co], AsyncContextDecorator):
+    class _AsyncGeneratorContextManager(AbstractAsyncContextManager[_T_co, bool | None], AsyncContextDecorator):
         # __init__ and these attributes are actually defined in the base class _GeneratorContextManagerBase,
         # which is more trouble than it's worth to include in the stub (see #6676)
         def __init__(self, func: Callable[..., AsyncIterator[_T_co]], args: tuple[Any, ...], kwds: dict[str, Any]) -> None: ...
@@ -94,7 +95,7 @@ if sys.version_info >= (3, 10):
         ) -> bool | None: ...
 
 else:
-    class _AsyncGeneratorContextManager(AbstractAsyncContextManager[_T_co]):
+    class _AsyncGeneratorContextManager(AbstractAsyncContextManager[_T_co, bool | None]):
         def __init__(self, func: Callable[..., AsyncIterator[_T_co]], args: tuple[Any, ...], kwds: dict[str, Any]) -> None: ...
         gen: AsyncGenerator[_T_co, Any]
         func: Callable[..., AsyncGenerator[_T_co, Any]]
@@ -111,7 +112,7 @@ class _SupportsClose(Protocol):
 
 _SupportsCloseT = TypeVar("_SupportsCloseT", bound=_SupportsClose)
 
-class closing(AbstractContextManager[_SupportsCloseT]):
+class closing(AbstractContextManager[_SupportsCloseT, None]):
     def __init__(self, thing: _SupportsCloseT) -> None: ...
     def __exit__(self, *exc_info: Unused) -> None: ...
 
@@ -121,17 +122,17 @@ if sys.version_info >= (3, 10):
 
     _SupportsAcloseT = TypeVar("_SupportsAcloseT", bound=_SupportsAclose)
 
-    class aclosing(AbstractAsyncContextManager[_SupportsAcloseT]):
+    class aclosing(AbstractAsyncContextManager[_SupportsAcloseT, None]):
         def __init__(self, thing: _SupportsAcloseT) -> None: ...
         async def __aexit__(self, *exc_info: Unused) -> None: ...
 
-class suppress(AbstractContextManager[None]):
+class suppress(AbstractContextManager[None, bool]):
     def __init__(self, *exceptions: type[BaseException]) -> None: ...
     def __exit__(
         self, exctype: type[BaseException] | None, excinst: BaseException | None, exctb: TracebackType | None
     ) -> bool: ...
 
-class _RedirectStream(AbstractContextManager[_T_io]):
+class _RedirectStream(AbstractContextManager[_T_io, None]):
     def __init__(self, new_target: _T_io) -> None: ...
     def __exit__(
         self, exctype: type[BaseException] | None, excinst: BaseException | None, exctb: TracebackType | None
@@ -142,8 +143,8 @@ class redirect_stderr(_RedirectStream[_T_io]): ...
 
 # In reality this is a subclass of `AbstractContextManager`;
 # see #7961 for why we don't do that in the stub
-class ExitStack(metaclass=abc.ABCMeta):
-    def enter_context(self, cm: AbstractContextManager[_T]) -> _T: ...
+class ExitStack(Generic[_ExitT_co], metaclass=abc.ABCMeta):
+    def enter_context(self, cm: AbstractContextManager[_T, _ExitT_co]) -> _T: ...
     def push(self, exit: _CM_EF) -> _CM_EF: ...
     def callback(self, callback: Callable[_P, _T], /, *args: _P.args, **kwds: _P.kwargs) -> Callable[_P, _T]: ...
     def pop_all(self) -> Self: ...
@@ -151,18 +152,18 @@ class ExitStack(metaclass=abc.ABCMeta):
     def __enter__(self) -> Self: ...
     def __exit__(
         self, exc_type: type[BaseException] | None, exc_value: BaseException | None, traceback: TracebackType | None, /
-    ) -> bool: ...
+    ) -> _ExitT_co: ...
 
 _ExitCoroFunc: TypeAlias = Callable[
     [type[BaseException] | None, BaseException | None, TracebackType | None], Awaitable[bool | None]
 ]
-_ACM_EF = TypeVar("_ACM_EF", bound=AbstractAsyncContextManager[Any] | _ExitCoroFunc)
+_ACM_EF = TypeVar("_ACM_EF", bound=AbstractAsyncContextManager[Any, Any] | _ExitCoroFunc)
 
 # In reality this is a subclass of `AbstractAsyncContextManager`;
 # see #7961 for why we don't do that in the stub
-class AsyncExitStack(metaclass=abc.ABCMeta):
-    def enter_context(self, cm: AbstractContextManager[_T]) -> _T: ...
-    async def enter_async_context(self, cm: AbstractAsyncContextManager[_T]) -> _T: ...
+class AsyncExitStack(Generic[_ExitT_co], metaclass=abc.ABCMeta):
+    def enter_context(self, cm: AbstractContextManager[_T, _ExitT_co]) -> _T: ...
+    async def enter_async_context(self, cm: AbstractAsyncContextManager[_T, _ExitT_co]) -> _T: ...
     def push(self, exit: _CM_EF) -> _CM_EF: ...
     def push_async_exit(self, exit: _ACM_EF) -> _ACM_EF: ...
     def callback(self, callback: Callable[_P, _T], /, *args: _P.args, **kwds: _P.kwargs) -> Callable[_P, _T]: ...
@@ -177,7 +178,7 @@ class AsyncExitStack(metaclass=abc.ABCMeta):
     ) -> bool: ...
 
 if sys.version_info >= (3, 10):
-    class nullcontext(AbstractContextManager[_T], AbstractAsyncContextManager[_T]):
+    class nullcontext(AbstractContextManager[_T, None], AbstractAsyncContextManager[_T, None]):
         enter_result: _T
         @overload
         def __init__(self: nullcontext[None], enter_result: None = None) -> None: ...
@@ -189,7 +190,7 @@ if sys.version_info >= (3, 10):
         async def __aexit__(self, *exctype: Unused) -> None: ...
 
 else:
-    class nullcontext(AbstractContextManager[_T]):
+    class nullcontext(AbstractContextManager[_T, None]):
         enter_result: _T
         @overload
         def __init__(self: nullcontext[None], enter_result: None = None) -> None: ...
@@ -201,7 +202,7 @@ else:
 if sys.version_info >= (3, 11):
     _T_fd_or_any_path = TypeVar("_T_fd_or_any_path", bound=FileDescriptorOrPath)
 
-    class chdir(AbstractContextManager[None], Generic[_T_fd_or_any_path]):
+    class chdir(AbstractContextManager[None, None], Generic[_T_fd_or_any_path]):
         path: _T_fd_or_any_path
         def __init__(self, path: _T_fd_or_any_path) -> None: ...
         def __enter__(self) -> None: ...