From 7d8e810ab19877e32b33fb92c9f8e51b7bdf4e00 Mon Sep 17 00:00:00 2001 From: Saran Tunyasuvunakool Date: Wed, 8 Feb 2023 06:36:52 -0800 Subject: [PATCH] Fix IPython history corruption when using `launch_repl`. PiperOrigin-RevId: 508064761 Change-Id: I58fa2e4f35b57a1ab68c3172ba5c0d4495e40ced --- doc/changelog.rst | 14 ++++++- python/mujoco/simulate.cc | 5 ++- python/mujoco/viewer.py | 83 +++++++++++++++++++++++++++++++++------ 3 files changed, 86 insertions(+), 16 deletions(-) diff --git a/doc/changelog.rst b/doc/changelog.rst index 31211c83..91205484 100644 --- a/doc/changelog.rst +++ b/doc/changelog.rst @@ -2,9 +2,19 @@ Changelog ========= -Version 2.3.2 (February 7, 2023) +Upcoming version (not yet released) ----------------------------------- +Python bindings +^^^^^^^^^^^^^^^ + +- Fixed IPython history corruption when using ``launch_repl``. The ``launch_repl`` function now provides seamless + continuation of an IPython interactive shell session, and is no longer considered experimental feature. + + +Version 2.3.2 (February 7, 2023) +-------------------------------- + General ^^^^^^^ @@ -54,7 +64,7 @@ Plugins Version 2.3.1 (December 6, 2022) ------------------------------------ +-------------------------------- Python bindings ^^^^^^^^^^^^^^^ diff --git a/python/mujoco/simulate.cc b/python/mujoco/simulate.cc index 45f84b81..8f778a6c 100644 --- a/python/mujoco/simulate.cc +++ b/python/mujoco/simulate.cc @@ -88,11 +88,14 @@ PYBIND11_MODULE(_simulate, pymodule) { .def_readonly("run", &mujoco::Simulate::run, py::call_guard()) - .def_property_readonly( + .def_property( "exitrequest", [](mujoco::Simulate& simulate) { return simulate.exitrequest.load(); }, + [](mujoco::Simulate& simulate, bool exitrequest) { + simulate.exitrequest.store(exitrequest); + }, py::call_guard()) .def_property_readonly( diff --git a/python/mujoco/viewer.py b/python/mujoco/viewer.py index f4201e40..89ab655b 100644 --- a/python/mujoco/viewer.py +++ b/python/mujoco/viewer.py @@ -19,6 +19,7 @@ import code import inspect import math import os +import sys import threading import time from typing import Callable, Optional, Tuple, Union @@ -213,7 +214,8 @@ def _launch_internal(model: Optional[mujoco.MjModel] = None, data: Optional[mujoco.MjData] = None, *, run_physics_thread: bool = True, - loader: Optional[_InternalLoaderType] = None) -> None: + loader: Optional[_InternalLoaderType] = None, + simulate: Optional[Simulate] = None) -> None: """Internal API, so that the public API has more readable type annotations.""" if model is None and data is not None: raise ValueError('mjData is specified but mjModel is not') @@ -233,7 +235,8 @@ def _launch_internal(model: Optional[mujoco.MjModel] = None, loader = _loader # The simulate object encapsulates the UI. - simulate = Simulate() + if simulate is None: + simulate = Simulate() # Initialize GLFW. if not glfw.init(): @@ -269,25 +272,79 @@ def launch_from_path(path: str) -> None: def launch_repl(model: mujoco.MjModel, data: mujoco.MjData) -> None: - """EXPERIMENTAL FEATURE: Launches the Simulate GUI in REPL mode.""" + """Launches the Simulate GUI in REPL mode.""" + ipython_shell = None try: import IPython # pylint: disable=g-import-not-at-top - has_ipython = True + ipython_shell = IPython.get_ipython() + ipython_is_terminal_interactive_shell = isinstance( + ipython_shell, + IPython.terminal.interactiveshell.TerminalInteractiveShell) except ImportError: - has_ipython = False + ipython_is_terminal_interactive_shell = False + + simulate = Simulate() + viewer_is_running = True def start_shell(global_variables): - if has_ipython and IPython.get_ipython() is not None: - locals().update(global_variables) - IPython.embed() + if ipython_is_terminal_interactive_shell: + ipython_shell.execution_count += 1 + + # A SQLite connection can only be used on the same thread that opened it. + # We cache the existing connection and reopen on the current thread. + old_db = ipython_shell.history_manager.db + ipython_shell.history_manager.init_db() + ipython_shell.history_manager.new_session() + + try: + # Replicate IPython main loop without exiting on keyboard interrupt, + # unless the viewer window has already been closed. + # (https://github.com/ipython/ipython/blob/8.9.0/IPython/terminal/interactiveshell.py#L701) + while viewer_is_running and ipython_shell.keep_running: + print(ipython_shell.separate_in, end='') + try: + c = ipython_shell.prompt_for_code() + except EOFError: + if not ipython_shell.confirm_exit or ipython_shell.ask_yes_no( + 'Do you really want to exit ([y]/n)?', 'y', 'n'): + ipython_shell.ask_exit() + if not ipython_shell.keep_running and simulate is not None: + simulate.exitrequest = True + else: + if c: + ipython_shell.run_cell(c, store_history=True) + finally: + # Close the temporary history DB connection and restore the old one. + ipython_shell.history_manager.end_session() + ipython_shell.history_manager.db.close() + ipython_shell.history_manager.db = old_db + ipython_shell.execution_count -= 1 else: code.InteractiveConsole(locals=global_variables).interact() - repl_thread = threading.Thread( - target=start_shell, args=(inspect.stack()[1][0].f_globals,)) - repl_thread.start() - launch(model, data, run_physics_thread=False) - repl_thread.join() + # End IPython history session on the main thread. We will need to open + # a new session in the REPL thread. + if ipython_is_terminal_interactive_shell: + ipython_shell.history_manager.end_session() + + try: + # Continue the IPython REPL session in a separate thread. + repl_thread = threading.Thread( + target=start_shell, args=(inspect.stack()[1][0].f_globals,)) + repl_thread.start() + + # Launch the viewer on the main thread. + _launch_internal( + model, data, run_physics_thread=False, simulate=simulate) + simulate = None + + # Wait until the REPL thread quits, then restore the IPython history + # DB session on the main thread. + viewer_is_running = False + repl_thread.join() + finally: + if ipython_is_terminal_interactive_shell: + ipython_shell.history_manager.new_session() if __name__ == '__main__':