command_context.py 815 B

123456789101112131415161718192021222324252627282930
  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):
  6. # type: () -> None
  7. super().__init__()
  8. self._in_main_context = False
  9. self._main_context = ExitStack()
  10. @contextmanager
  11. def main_context(self):
  12. # type: () -> Iterator[None]
  13. assert not self._in_main_context
  14. self._in_main_context = True
  15. try:
  16. with self._main_context:
  17. yield
  18. finally:
  19. self._in_main_context = False
  20. def enter_context(self, context_provider):
  21. # type: (ContextManager[_T]) -> _T
  22. assert self._in_main_context
  23. return self._main_context.enter_context(context_provider)