current.py 6.1 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195
  1. from __future__ import annotations
  2. from contextlib import contextmanager
  3. from contextvars import ContextVar
  4. from typing import TYPE_CHECKING, Any, Generator
  5. if TYPE_CHECKING:
  6. from prompt_toolkit.input.base import Input
  7. from prompt_toolkit.output.base import Output
  8. from .application import Application
  9. __all__ = [
  10. "AppSession",
  11. "get_app_session",
  12. "get_app",
  13. "get_app_or_none",
  14. "set_app",
  15. "create_app_session",
  16. "create_app_session_from_tty",
  17. ]
  18. class AppSession:
  19. """
  20. An AppSession is an interactive session, usually connected to one terminal.
  21. Within one such session, interaction with many applications can happen, one
  22. after the other.
  23. The input/output device is not supposed to change during one session.
  24. Warning: Always use the `create_app_session` function to create an
  25. instance, so that it gets activated correctly.
  26. :param input: Use this as a default input for all applications
  27. running in this session, unless an input is passed to the `Application`
  28. explicitly.
  29. :param output: Use this as a default output.
  30. """
  31. def __init__(
  32. self, input: Input | None = None, output: Output | None = None
  33. ) -> None:
  34. self._input = input
  35. self._output = output
  36. # The application will be set dynamically by the `set_app` context
  37. # manager. This is called in the application itself.
  38. self.app: Application[Any] | None = None
  39. def __repr__(self) -> str:
  40. return f"AppSession(app={self.app!r})"
  41. @property
  42. def input(self) -> Input:
  43. if self._input is None:
  44. from prompt_toolkit.input.defaults import create_input
  45. self._input = create_input()
  46. return self._input
  47. @property
  48. def output(self) -> Output:
  49. if self._output is None:
  50. from prompt_toolkit.output.defaults import create_output
  51. self._output = create_output()
  52. return self._output
  53. _current_app_session: ContextVar[AppSession] = ContextVar(
  54. "_current_app_session", default=AppSession()
  55. )
  56. def get_app_session() -> AppSession:
  57. return _current_app_session.get()
  58. def get_app() -> Application[Any]:
  59. """
  60. Get the current active (running) Application.
  61. An :class:`.Application` is active during the
  62. :meth:`.Application.run_async` call.
  63. We assume that there can only be one :class:`.Application` active at the
  64. same time. There is only one terminal window, with only one stdin and
  65. stdout. This makes the code significantly easier than passing around the
  66. :class:`.Application` everywhere.
  67. If no :class:`.Application` is running, then return by default a
  68. :class:`.DummyApplication`. For practical reasons, we prefer to not raise
  69. an exception. This way, we don't have to check all over the place whether
  70. an actual `Application` was returned.
  71. (For applications like pymux where we can have more than one `Application`,
  72. we'll use a work-around to handle that.)
  73. """
  74. session = _current_app_session.get()
  75. if session.app is not None:
  76. return session.app
  77. from .dummy import DummyApplication
  78. return DummyApplication()
  79. def get_app_or_none() -> Application[Any] | None:
  80. """
  81. Get the current active (running) Application, or return `None` if no
  82. application is running.
  83. """
  84. session = _current_app_session.get()
  85. return session.app
  86. @contextmanager
  87. def set_app(app: Application[Any]) -> Generator[None, None, None]:
  88. """
  89. Context manager that sets the given :class:`.Application` active in an
  90. `AppSession`.
  91. This should only be called by the `Application` itself.
  92. The application will automatically be active while its running. If you want
  93. the application to be active in other threads/coroutines, where that's not
  94. the case, use `contextvars.copy_context()`, or use `Application.context` to
  95. run it in the appropriate context.
  96. """
  97. session = _current_app_session.get()
  98. previous_app = session.app
  99. session.app = app
  100. try:
  101. yield
  102. finally:
  103. session.app = previous_app
  104. @contextmanager
  105. def create_app_session(
  106. input: Input | None = None, output: Output | None = None
  107. ) -> Generator[AppSession, None, None]:
  108. """
  109. Create a separate AppSession.
  110. This is useful if there can be multiple individual `AppSession`s going on.
  111. Like in the case of an Telnet/SSH server.
  112. """
  113. # If no input/output is specified, fall back to the current input/output,
  114. # if there was one that was set/created for the current session.
  115. # (Note that we check `_input`/`_output` and not `input`/`output`. This is
  116. # because we don't want to accidently create a new input/output objects
  117. # here and store it in the "parent" `AppSession`. Especially, when
  118. # combining pytest's `capsys` fixture and `create_app_session`, sys.stdin
  119. # and sys.stderr are patched for every test, so we don't want to leak
  120. # those outputs object across `AppSession`s.)
  121. if input is None:
  122. input = get_app_session()._input
  123. if output is None:
  124. output = get_app_session()._output
  125. # Create new `AppSession` and activate.
  126. session = AppSession(input=input, output=output)
  127. token = _current_app_session.set(session)
  128. try:
  129. yield session
  130. finally:
  131. _current_app_session.reset(token)
  132. @contextmanager
  133. def create_app_session_from_tty() -> Generator[AppSession, None, None]:
  134. """
  135. Create `AppSession` that always prefers the TTY input/output.
  136. Even if `sys.stdin` and `sys.stdout` are connected to input/output pipes,
  137. this will still use the terminal for interaction (because `sys.stderr` is
  138. still connected to the terminal).
  139. Usage::
  140. from prompt_toolkit.shortcuts import prompt
  141. with create_app_session_from_tty():
  142. prompt('>')
  143. """
  144. from prompt_toolkit.input.defaults import create_input
  145. from prompt_toolkit.output.defaults import create_output
  146. input = create_input(always_prefer_tty=True)
  147. output = create_output(always_prefer_tty=True)
  148. with create_app_session(input=input, output=output) as app_session:
  149. yield app_session