command_context.py 760 B

123456789101112131415161718192021222324252627
  1. from contextlib import ExitStack, contextmanager
  2. from typing import ContextManager, Iterator, TypeVar
  3. _T = TypeVar("_T", covariant=True)
  4. class CommandContextMixIn:
  5. def __init__(self) -> None:
  6. super().__init__()
  7. self._in_main_context = False
  8. self._main_context = ExitStack()
  9. @contextmanager
  10. def main_context(self) -> Iterator[None]:
  11. assert not self._in_main_context
  12. self._in_main_context = True
  13. try:
  14. with self._main_context:
  15. yield
  16. finally:
  17. self._in_main_context = False
  18. def enter_context(self, context_provider: ContextManager[_T]) -> _T:
  19. assert self._in_main_context
  20. return self._main_context.enter_context(context_provider)