feat(cli): run node inside shell command

Author: jansegre

hathor_cli/_shell_extension.py

@@ -0,0 +1,196 @@
+"""
+IPython extension that adapts logging handlers to play nicely with the interactive prompt.
+
+When loaded, all stream-based logging handlers are updated so their output is rendered
+through prompt_toolkit's ``run_in_terminal`` helper, which ensures log lines appear
+above the current input without corrupting the prompt. The original streams are restored
+when the extension is unloaded.
+"""
+
+from __future__ import annotations
+
+import io
+import logging
+import threading
+from contextlib import suppress
+from typing import Any, Callable, Iterable
+
+get_app_or_none: Callable[[], Any] | None = None
+pt_utils: Any | None = None
+
+try:  # pragma: no cover - IPython dependency
+    from prompt_toolkit.application import get_app_or_none as _get_app_or_none
+    from prompt_toolkit.shortcuts import utils as _pt_utils
+except ImportError:  # pragma: no cover - handled at runtime
+    pass
+else:
+    get_app_or_none = _get_app_or_none
+    pt_utils = _pt_utils
+
+_original_streams: dict[logging.StreamHandler, Any] = {}
+_installed = False
+
+
+class PromptToolkitLogStream(io.TextIOBase):
+    """Proxy stream that forwards writes through prompt_toolkit."""
+
+    def __init__(self, inner: Any):
+        super().__init__()
+        self._inner = inner
+        self._encoding_override: str | None = None
+        self._errors_override: str | None = None
+
+    def _run_in_terminal(self, func: Callable[[], None]) -> None:
+        if pt_utils is None:
+            func()
+            return
+
+        app = get_app_or_none() if get_app_or_none is not None else None
+        if app is None:
+            func()
+            return
+        loop = getattr(app, 'loop', None)
+        if loop is None:
+            func()
+            return
+
+        event = threading.Event()
+        handled = False
+
+        def run_and_signal() -> None:
+            nonlocal handled
+            try:
+                if not handled:
+                    pt_utils.run_in_terminal(func, in_executor=False)
+            finally:
+                handled = True
+                event.set()
+
+        loop.call_soon_threadsafe(run_and_signal)
+        if not event.wait(timeout=5):
+            handled = True
+            func()
+
+    def write(self, data: str) -> int:
+        if not data:
+            return 0
+
+        def _write() -> None:
+            self._inner.write(data)
+
+        self._run_in_terminal(_write)
+        return len(data)
+
+    def flush(self) -> None:
+        def _flush() -> None:
+            self._inner.flush()
+
+        self._run_in_terminal(_flush)
+
+    @property
+    def encoding(self) -> str:
+        if self._encoding_override is not None:
+            return self._encoding_override
+        return getattr(self._inner, 'encoding', 'utf-8') or 'utf-8'
+
+    @encoding.setter
+    def encoding(self, value: str) -> None:
+        self._encoding_override = value
+
+    @property
+    def errors(self) -> str:
+        if self._errors_override is not None:
+            return self._errors_override
+        return getattr(self._inner, 'errors', 'strict') or 'strict'
+
+    @errors.setter
+    def errors(self, value: str) -> None:
+        self._errors_override = value
+
+    def fileno(self) -> int:
+        if hasattr(self._inner, 'fileno') and callable(getattr(self._inner, 'fileno')):
+            return self._inner.fileno()
+        raise io.UnsupportedOperation('fileno not available')
+
+    def isatty(self) -> bool:
+        if hasattr(self._inner, 'isatty') and callable(getattr(self._inner, 'isatty')):
+            return self._inner.isatty()
+        return False
+
+    def close(self) -> None:
+        # Do not close the underlying stream.
+        pass
+
+    @property
+    def closed(self) -> bool:
+        return False
+
+    def readable(self) -> bool:
+        return False
+
+    def seekable(self) -> bool:
+        return False
+
+    def writable(self) -> bool:
+        return True
+
+
+def _iter_stream_handlers() -> Iterable[logging.StreamHandler]:
+    """Yield every stream handler currently registered."""
+    root = logging.getLogger()
+    for handler in root.handlers:
+        if isinstance(handler, logging.StreamHandler):
+            yield handler
+
+    for logger in logging.Logger.manager.loggerDict.values():
+        if isinstance(logger, logging.PlaceHolder):
+            continue
+        if not isinstance(logger, logging.Logger):
+            continue
+        for handler in logger.handlers:
+            if isinstance(handler, logging.StreamHandler):
+                yield handler
+
+
+def _install_prompt_toolkit_streams() -> None:
+    if pt_utils is None:
+        return
+
+    for handler in _iter_stream_handlers():
+        current_stream = getattr(handler, 'stream', None)
+        if current_stream is None or isinstance(current_stream, PromptToolkitLogStream):
+            continue
+
+        proxy = PromptToolkitLogStream(current_stream)
+        _original_streams[handler] = current_stream
+        handler.stream = proxy
+
+
+def load_ipython_extension(shell: Any) -> None:  # pragma: no cover - requires IPython
+    """Called by IPython when the extension is loaded."""
+    global _installed
+
+    if pt_utils is None:
+        shell.write_err('prompt_toolkit not available; logs will use standard output.\n')
+        return
+
+    if _installed:
+        return
+
+    _install_prompt_toolkit_streams()
+    _installed = True
+
+
+def unload_ipython_extension(shell: Any) -> None:  # pragma: no cover - requires IPython
+    """Called by IPython when the extension is unloaded."""
+    restore_logging_streams()
+
+
+def restore_logging_streams() -> None:
+    """Restore the original logging streams."""
+    global _installed
+    for handler, stream in list(_original_streams.items()):
+        with suppress(Exception):
+            handler.stream = stream
+    _original_streams.clear()
+    _installed = False

hathor_cli/shell.py

@@ -12,57 +12,160 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+import threading
 from argparse import Namespace
-from typing import Any, Callable
+from contextlib import suppress
+from typing import Any, Callable, TypeVar, cast
 
 from hathor_cli.run_node import RunNode
 
+T = TypeVar('T')
 
-def get_ipython(extra_args: list[Any], imported_objects: dict[str, Any]) -> Callable[[], None]:
+
+def get_ipython(
+        extra_args: list[Any],
+        imported_objects: dict[str, Any],
+        *,
+        config: Any | None = None,
+) -> Callable[[], None]:
     from IPython import start_ipython
 
     def run_ipython():
-        start_ipython(argv=extra_args, user_ns=imported_objects)
+        start_ipython(argv=extra_args, user_ns=imported_objects, config=config)
 
     return run_ipython
 
 
 class Shell(RunNode):
+    _reactor_thread: threading.Thread | None = None
+    _shell_run_node: bool = False
+
+    @classmethod
+    def create_parser(cls):
+        parser = super().create_parser()
+        parser.add_argument(
+            '--x-run-node',
+            action='store_true',
+            help='Start the full node in the background while keeping the interactive shell open.'
+        )
+        return parser
+
     def start_manager(self) -> None:
-        pass
+        if not self._shell_run_node:
+            return
+
+        super().start_manager()
+        self._start_reactor_thread()
 
     def register_signal_handlers(self) -> None:
         pass
 
     def prepare(self, *, register_resources: bool = True) -> None:
-        super().prepare(register_resources=False)
+        super().prepare(register_resources=self._shell_run_node)
 
         imported_objects: dict[str, Any] = {}
         imported_objects['tx_storage'] = self.tx_storage
         if self._args.wallet:
             imported_objects['wallet'] = self.wallet
         imported_objects['manager'] = self.manager
-        self.shell = get_ipython(self.extra_args, imported_objects)
+        imported_objects['reactor'] = self.reactor
+        ipy_config: Any | None = None
+
+        if self._shell_run_node:
+            import asyncio
+            from twisted.internet.defer import Deferred
+            from traitlets.config import Config
+
+            async def await_deferred(deferred: Deferred[T]) -> T:
+                loop = asyncio.get_running_loop()
+                return await deferred.asFuture(loop)
+
+            imported_objects['await_deferred'] = await_deferred
+            imported_objects['asyncio'] = asyncio
+            ipy_config = Config()
+            ipy_config.InteractiveShellApp.extra_extensions = ['hathor_cli._shell_extension']
+
+        self.shell = get_ipython(self.extra_args, imported_objects, config=ipy_config)
 
         print()
         print('--- Injected globals ---')
         for name, obj in imported_objects.items():
             print(name, obj)
         print('------------------------')
         print()
+        if self._shell_run_node:
+            print('Node reactor started in background. Use await_deferred() for Deferreds.')
 
     def parse_args(self, argv: list[str]) -> Namespace:
         # TODO: add help for the `--` extra argument separator
+        argv = list(argv)
         extra_args: list[str] = []
         if '--' in argv:
             idx = argv.index('--')
             extra_args = argv[idx + 1:]
             argv = argv[:idx]
         self.extra_args = extra_args
-        return self.parser.parse_args(argv)
+        namespace = self.parser.parse_args(argv)
+        self._shell_run_node = bool(getattr(namespace, 'x_run_node', False))
+        return namespace
 
     def run(self) -> None:
-        self.shell()
+        try:
+            self.shell()
+        finally:
+            if self._shell_run_node:
+                self._shutdown_background()
+
+    def _start_reactor_thread(self) -> None:
+        if self._reactor_thread and self._reactor_thread.is_alive():
+            return
+
+        def run_reactor() -> None:
+            self.log.info('reactor thread starting')
+            try:
+                run = getattr(self.reactor, 'run')
+                try:
+                    run(installSignalHandlers=False)
+                except TypeError:
+                    run()
+            finally:
+                self.log.info('reactor thread finished')
+
+        self._reactor_thread = threading.Thread(
+            target=run_reactor,
+            name='hathor-reactor',
+            daemon=True,
+        )
+        self._reactor_thread.start()
+
+    def _shutdown_background(self) -> None:
+        thread = self._reactor_thread
+        if thread and thread.is_alive():
+            try:
+                from twisted.internet.interfaces import IReactorFromThreads
+
+                threaded_reactor = cast(IReactorFromThreads, self.reactor)
+                threaded_reactor.callFromThread(self.reactor.stop)
+            except Exception:
+                self.log.exception('failed to schedule reactor shutdown from shell')
+            thread.join(timeout=30)
+            if thread.is_alive():
+                self.log.warning('reactor thread did not finish cleanly')
+        self._reactor_thread = None
+        if self._shell_run_node:
+            self._restore_logging_streams()
+
+    def _restore_logging_streams(self) -> None:
+        try:
+            from hathor_cli import _shell_extension
+        except ImportError:
+            return
+        with suppress(Exception):
+            _shell_extension.restore_logging_streams()
+
+    def __del__(self):
+        with suppress(Exception):
+            self._restore_logging_streams()
 
 
 def main():