From 9c0a59a75a376ffdbc269c808f1ea0bbfa76f26a Mon Sep 17 00:00:00 2001 From: Brett Jia Date: Wed, 6 Mar 2024 10:55:28 -0500 Subject: [PATCH] server: graceful repl exit + multi-repl support (#1362) * server: graceful repl exit + hacky multi-repl support * remove prints * more multi-repl layout fixes, ignore benign CancelledError * add missing import * disable input function for safety --- server/python/plugin_repl.py | 127 +++++++++++++++++++++++++++++++---- 1 file changed, 115 insertions(+), 12 deletions(-) diff --git a/server/python/plugin_repl.py b/server/python/plugin_repl.py index 292ea4ada..4ecdb9945 100644 --- a/server/python/plugin_repl.py +++ b/server/python/plugin_repl.py @@ -1,11 +1,25 @@ import asyncio import concurrent.futures +from functools import partial +import inspect +import prompt_toolkit from prompt_toolkit import print_formatted_text -from prompt_toolkit.contrib.telnet.server import TelnetServer +from prompt_toolkit.application import Application +import prompt_toolkit.application.current +import prompt_toolkit.key_binding.key_processor +import prompt_toolkit.contrib.telnet.server +from prompt_toolkit.contrib.telnet.server import TelnetServer, TelnetConnection +from prompt_toolkit.shortcuts import clear_title, set_title from ptpython.repl import embed, PythonRepl +import ptpython.key_bindings +import ptpython.python_input +import ptpython.history_browser +import ptpython.layout import socket import telnetlib import threading +import traceback +import types from typing import List, Dict, Any from scrypted_python.scrypted_sdk import ScryptedStatic, ScryptedDevice @@ -13,10 +27,90 @@ from scrypted_python.scrypted_sdk import ScryptedStatic, ScryptedDevice from rpc import maybe_await +# This section is a bit of a hack - prompt_toolkit has many assumptions +# that there is only one global Application, so multiple REPLs will confuse +# the library. The patches here allow us to scope a particular call stack +# to a particular REPL, and to get the current Application from the stack. +default_get_app = prompt_toolkit.application.current.get_app +def get_app_patched() -> Application[Any]: + stack = inspect.stack() + for frame in stack: + self_var = frame.frame.f_locals.get("self") + if self_var is not None and isinstance(self_var, Application): + return self_var + return default_get_app() +prompt_toolkit.application.current.get_app = get_app_patched +prompt_toolkit.key_binding.key_processor.get_app = get_app_patched +prompt_toolkit.contrib.telnet.server.get_app = get_app_patched +ptpython.python_input.get_app = get_app_patched +ptpython.key_bindings.get_app = get_app_patched +ptpython.history_browser.get_app = get_app_patched +ptpython.layout.get_app = get_app_patched + + +async def run_async_patched(self: PythonRepl) -> None: + # This is a patched version of PythonRepl.run_async to handle an + # AssertionError raised by prompt_toolkit when the TelnetServer exits. + # Original: https://github.com/prompt-toolkit/ptpython/blob/3.0.26/ptpython/repl.py#L215 + + """ + Run the REPL loop, but run the blocking parts in an executor, so that + we don't block the event loop. Both the input and output (which can + display a pager) will run in a separate thread with their own event + loop, this way ptpython's own event loop won't interfere with the + asyncio event loop from where this is called. + + The "eval" however happens in the current thread, which is important. + (Both for control-C to work, as well as for the code to see the right + thread in which it was embedded). + """ + loop = asyncio.get_running_loop() + + if self.terminal_title: + set_title(self.terminal_title) + + self._add_to_namespace() + + try: + while True: + try: + # Read. + try: + text = await loop.run_in_executor(None, self.read) + except EOFError: + return + except asyncio.CancelledError: + return + except AssertionError: + return + except BaseException: + # Something went wrong while reading input. + # (E.g., a bug in the completer that propagates. Don't + # crash the REPL.) + traceback.print_exc() + continue + + # Eval. + await self.run_and_show_expression_async(text) + + except KeyboardInterrupt as e: + # XXX: This does not yet work properly. In some situations, + # `KeyboardInterrupt` exceptions can end up in the event + # loop selector. + self._handle_keyboard_interrupt(e) + except SystemExit: + return + finally: + if self.terminal_title: + clear_title() + self._remove_from_namespace() + + def configure(repl: PythonRepl) -> None: repl.confirm_exit = False repl.enable_system_bindings = False repl.enable_mouse_support = False + repl.run_async = types.MethodType(run_async_patched, repl) async def createREPLServer(sdk: ScryptedStatic, plugin: ScryptedDevice) -> int: @@ -29,9 +123,11 @@ async def createREPLServer(sdk: ScryptedStatic, plugin: ScryptedDevice) -> int: sock.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1) sock.settimeout(None) sock.bind(('localhost', 0)) - sock.listen(1) + sock.listen() - async def start_telnet_repl(future, filter) -> None: + loop: asyncio.AbstractEventLoop = asyncio.get_event_loop() + + async def start_telnet_repl(future: concurrent.futures.Future, filter: str) -> None: if filter == "undefined": filter = None @@ -58,11 +154,13 @@ async def createREPLServer(sdk: ScryptedStatic, plugin: ScryptedDevice) -> int: telnet_port = s.getsockname()[1] s.close() - async def interact(connection) -> None: + async def interact(connection: TelnetConnection) -> None: + repl_print = partial(print_formatted_text, output=connection.vt100_output) global_dict = { **globals(), - "print": print_formatted_text, - "help": lambda *args, **kwargs: print_formatted_text("Help is not available in this environment"), + "print": repl_print, + "help": lambda *args, **kwargs: repl_print("Help is not available in this environment"), + "input": lambda *args, **kwargs: repl_print("Input is not available in this environment"), } locals_dict = { "device": device, @@ -77,13 +175,13 @@ async def createREPLServer(sdk: ScryptedStatic, plugin: ScryptedDevice) -> int: print_formatted_text(banner) await embed(return_asyncio_coroutine=True, globals=global_dict, locals=locals_dict, configure=configure) + server_task: asyncio.Task = None + def ready_cb(): + future.set_result((telnet_port, lambda: loop.call_soon_threadsafe(server_task.cancel))) + # Start the REPL server telnet_server = TelnetServer(interact=interact, port=telnet_port, enable_cpr=False) - telnet_server.start() - - future.set_result(telnet_port) - - loop = asyncio.get_event_loop() + server_task = asyncio.create_task(telnet_server.run(ready_cb=ready_cb)) def handle_connection(conn: socket.socket): conn.settimeout(None) @@ -91,7 +189,7 @@ async def createREPLServer(sdk: ScryptedStatic, plugin: ScryptedDevice) -> int: future = concurrent.futures.Future() loop.call_soon_threadsafe(loop.create_task, start_telnet_repl(future, filter)) - telnet_port = future.result() + telnet_port, exit_server = future.result() telnet_client = telnetlib.Telnet('localhost', telnet_port, timeout=None) @@ -111,6 +209,9 @@ async def createREPLServer(sdk: ScryptedStatic, plugin: ScryptedDevice) -> int: if not data: break telnet_client.write(data) + telnet_client.close() + exit_server() + def forward_to_socket(): prompt_count = 0 while True: @@ -128,6 +229,8 @@ async def createREPLServer(sdk: ScryptedStatic, plugin: ScryptedDevice) -> int: if prompt_count < 5: data = data.replace(b">>>", b" ") conn.sendall(data) + conn.close() + exit_server() threading.Thread(target=forward_to_telnet).start() threading.Thread(target=forward_to_socket).start()