_contextmanagers.py 7.0 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200
  1. from __future__ import annotations
  2. from abc import abstractmethod
  3. from contextlib import AbstractAsyncContextManager, AbstractContextManager
  4. from inspect import isasyncgen, iscoroutine, isgenerator
  5. from types import TracebackType
  6. from typing import Protocol, TypeVar, cast, final
  7. _T_co = TypeVar("_T_co", covariant=True)
  8. _ExitT_co = TypeVar("_ExitT_co", covariant=True, bound="bool | None")
  9. class _SupportsCtxMgr(Protocol[_T_co, _ExitT_co]):
  10. def __contextmanager__(self) -> AbstractContextManager[_T_co, _ExitT_co]: ...
  11. class _SupportsAsyncCtxMgr(Protocol[_T_co, _ExitT_co]):
  12. def __asynccontextmanager__(
  13. self,
  14. ) -> AbstractAsyncContextManager[_T_co, _ExitT_co]: ...
  15. class ContextManagerMixin:
  16. """
  17. Mixin class providing context manager functionality via a generator-based
  18. implementation.
  19. This class allows you to implement a context manager via :meth:`__contextmanager__`
  20. which should return a generator. The mechanics are meant to mirror those of
  21. :func:`@contextmanager <contextlib.contextmanager>`.
  22. .. note:: Classes using this mix-in are not reentrant as context managers, meaning
  23. that once you enter it, you can't re-enter before first exiting it.
  24. .. seealso:: :doc:`contextmanagers`
  25. """
  26. __cm: AbstractContextManager[object, bool | None] | None = None
  27. @final
  28. def __enter__(self: _SupportsCtxMgr[_T_co, bool | None]) -> _T_co:
  29. # Needed for mypy to assume self still has the __cm member
  30. assert isinstance(self, ContextManagerMixin)
  31. if self.__cm is not None:
  32. raise RuntimeError(
  33. f"this {self.__class__.__qualname__} has already been entered"
  34. )
  35. cm = self.__contextmanager__()
  36. if not isinstance(cm, AbstractContextManager):
  37. if isgenerator(cm):
  38. raise TypeError(
  39. "__contextmanager__() returned a generator object instead of "
  40. "a context manager. Did you forget to add the @contextmanager "
  41. "decorator?"
  42. )
  43. raise TypeError(
  44. f"__contextmanager__() did not return a context manager object, "
  45. f"but {cm.__class__!r}"
  46. )
  47. if cm is self:
  48. raise TypeError(
  49. f"{self.__class__.__qualname__}.__contextmanager__() returned "
  50. f"self. Did you forget to add the @contextmanager decorator and a "
  51. f"'yield' statement?"
  52. )
  53. value = cm.__enter__()
  54. self.__cm = cm
  55. return value
  56. @final
  57. def __exit__(
  58. self: _SupportsCtxMgr[object, _ExitT_co],
  59. exc_type: type[BaseException] | None,
  60. exc_val: BaseException | None,
  61. exc_tb: TracebackType | None,
  62. ) -> _ExitT_co:
  63. # Needed for mypy to assume self still has the __cm member
  64. assert isinstance(self, ContextManagerMixin)
  65. if self.__cm is None:
  66. raise RuntimeError(
  67. f"this {self.__class__.__qualname__} has not been entered yet"
  68. )
  69. # Prevent circular references
  70. cm = self.__cm
  71. del self.__cm
  72. return cast(_ExitT_co, cm.__exit__(exc_type, exc_val, exc_tb))
  73. @abstractmethod
  74. def __contextmanager__(self) -> AbstractContextManager[object, bool | None]:
  75. """
  76. Implement your context manager logic here.
  77. This method **must** be decorated with
  78. :func:`@contextmanager <contextlib.contextmanager>`.
  79. .. note:: Remember that the ``yield`` will raise any exception raised in the
  80. enclosed context block, so use a ``finally:`` block to clean up resources!
  81. :return: a context manager object
  82. """
  83. class AsyncContextManagerMixin:
  84. """
  85. Mixin class providing async context manager functionality via a generator-based
  86. implementation.
  87. This class allows you to implement a context manager via
  88. :meth:`__asynccontextmanager__`. The mechanics are meant to mirror those of
  89. :func:`@asynccontextmanager <contextlib.asynccontextmanager>`.
  90. .. note:: Classes using this mix-in are not reentrant as context managers, meaning
  91. that once you enter it, you can't re-enter before first exiting it.
  92. .. seealso:: :doc:`contextmanagers`
  93. """
  94. __cm: AbstractAsyncContextManager[object, bool | None] | None = None
  95. @final
  96. async def __aenter__(self: _SupportsAsyncCtxMgr[_T_co, bool | None]) -> _T_co:
  97. # Needed for mypy to assume self still has the __cm member
  98. assert isinstance(self, AsyncContextManagerMixin)
  99. if self.__cm is not None:
  100. raise RuntimeError(
  101. f"this {self.__class__.__qualname__} has already been entered"
  102. )
  103. cm = self.__asynccontextmanager__()
  104. if not isinstance(cm, AbstractAsyncContextManager):
  105. if isasyncgen(cm):
  106. raise TypeError(
  107. "__asynccontextmanager__() returned an async generator instead of "
  108. "an async context manager. Did you forget to add the "
  109. "@asynccontextmanager decorator?"
  110. )
  111. elif iscoroutine(cm):
  112. cm.close()
  113. raise TypeError(
  114. "__asynccontextmanager__() returned a coroutine object instead of "
  115. "an async context manager. Did you forget to add the "
  116. "@asynccontextmanager decorator and a 'yield' statement?"
  117. )
  118. raise TypeError(
  119. f"__asynccontextmanager__() did not return an async context manager, "
  120. f"but {cm.__class__!r}"
  121. )
  122. if cm is self:
  123. raise TypeError(
  124. f"{self.__class__.__qualname__}.__asynccontextmanager__() returned "
  125. f"self. Did you forget to add the @asynccontextmanager decorator and a "
  126. f"'yield' statement?"
  127. )
  128. value = await cm.__aenter__()
  129. self.__cm = cm
  130. return value
  131. @final
  132. async def __aexit__(
  133. self: _SupportsAsyncCtxMgr[object, _ExitT_co],
  134. exc_type: type[BaseException] | None,
  135. exc_val: BaseException | None,
  136. exc_tb: TracebackType | None,
  137. ) -> _ExitT_co:
  138. assert isinstance(self, AsyncContextManagerMixin)
  139. if self.__cm is None:
  140. raise RuntimeError(
  141. f"this {self.__class__.__qualname__} has not been entered yet"
  142. )
  143. # Prevent circular references
  144. cm = self.__cm
  145. del self.__cm
  146. return cast(_ExitT_co, await cm.__aexit__(exc_type, exc_val, exc_tb))
  147. @abstractmethod
  148. def __asynccontextmanager__(
  149. self,
  150. ) -> AbstractAsyncContextManager[object, bool | None]:
  151. """
  152. Implement your async context manager logic here.
  153. This method **must** be decorated with
  154. :func:`@asynccontextmanager <contextlib.asynccontextmanager>`.
  155. .. note:: Remember that the ``yield`` will raise any exception raised in the
  156. enclosed context block, so use a ``finally:`` block to clean up resources!
  157. :return: an async context manager object
  158. """