Pass the prompt-toolkit ssh session to the interact coroutine

Author: vxgmichel

examples/ssh/asyncssh-server.py

@@ -9,7 +9,7 @@
 from pygments.lexers.html import HtmlLexer
 
 from prompt_toolkit.completion import WordCompleter
-from prompt_toolkit.contrib.ssh import PromptToolkitSSHServer
+from prompt_toolkit.contrib.ssh import PromptToolkitSSHServer, PromptToolkitSSHSession
 from prompt_toolkit.lexers import PygmentsLexer
 from prompt_toolkit.shortcuts import ProgressBar, print_formatted_text
 from prompt_toolkit.shortcuts.dialogs import input_dialog, yes_no_dialog
@@ -54,7 +54,7 @@
 )
 
 
-async def interact() -> None:
+async def interact(ssh_session: PromptToolkitSSHSession) -> None:
  """
  The application interaction.
 

prompt_toolkit/contrib/ssh/__init__.py

@@ -1,6 +1,6 @@
-from .server import PromptToolkitSession, PromptToolkitSSHServer
+from .server import PromptToolkitSSHServer, PromptToolkitSSHSession
 
 __all__ = [
- "PromptToolkitSession",
+ "PromptToolkitSSHSession",
  "PromptToolkitSSHServer",
 ]

prompt_toolkit/contrib/ssh/server.py

@@ -12,14 +12,13 @@
 from prompt_toolkit.input import create_pipe_input
 from prompt_toolkit.output.vt100 import Vt100_Output
 
-__all__ = [
- "PromptToolkitSession",
- "PromptToolkitSSHServer",
-]
+__all__ = ["PromptToolkitSSHSession", "PromptToolkitSSHServer"]
 
 
-class PromptToolkitSession(asyncssh.SSHServerSession):
- def __init__(self, interact: Callable[[], Awaitable[None]]) -> None:
+class PromptToolkitSSHSession(asyncssh.SSHServerSession):
+ def __init__(
+ self, interact: Callable[["PromptToolkitSSHSession"], Awaitable[None]]
+ ) -> None:
  self.interact = interact
  self._chan = None
  self.app_session: Optional[AppSession] = None
@@ -77,7 +76,7 @@ async def _interact(self) -> None:
  with create_app_session(input=self._input, output=self._output) as session:
  self.app_session = session
  try:
- await self.interact()
+ await self.interact(self)
  except BaseException:
  traceback.print_exc()
  finally:
@@ -106,7 +105,7 @@ class PromptToolkitSSHServer(asyncssh.SSHServer):
 
  .. code:: python
 
- async def interact() -> None:
+ async def interact(ssh_session: PromptToolkitSSHSession) -> None:
  await yes_no_dialog("my title", "my text").run_async()
 
  prompt_session = PromptSession()
@@ -126,12 +125,14 @@ async def interact() -> None:
  loop.run_forever()
  """
 
- def __init__(self, interact: Callable[[], Awaitable[None]]) -> None:
+ def __init__(
+ self, interact: Callable[[PromptToolkitSSHSession], Awaitable[None]]
+ ) -> None:
  self.interact = interact
 
  def begin_auth(self, username):
  # No authentication.
  return False
 
- def session_requested(self) -> PromptToolkitSession:
- return PromptToolkitSession(self.interact)
+ def session_requested(self) -> PromptToolkitSSHSession:
+ return PromptToolkitSSHSession(self.interact)