Python 3.13 REPL support

contrib/ptdrgn.py

@@ -12,17 +12,13 @@
 Requires: "pip install ptpython" which brings in pygments and prompt_toolkit
 """
 import functools
-import importlib
 import os
 import shutil
-import sys
-from typing import Any, Callable, Dict, Optional, Set
+from typing import Any, Dict, Set
 
-from prompt_toolkit.completion import Completion, Completer
+from prompt_toolkit.completion import Completer
 from prompt_toolkit.formatted_text import PygmentsTokens
-from prompt_toolkit.formatted_text import fragment_list_to_text, to_formatted_text
 from ptpython import embed
-from ptpython.completer import DictionaryCompleter
 from ptpython.repl import run_config
 from pygments.lexers.c_cpp import CLexer
 
@@ -127,95 +123,14 @@ def _format_result_output(result: object):
     repl.completer = ReorderDrgnObjectCompleter(repl.completer)
 
 
-def run_interactive(
-    prog: drgn.Program,
-    banner_func: Optional[Callable[[str], str]] = None,
-    globals_func: Optional[Callable[[Dict[str, Any]], Dict[str, Any]]] = None,
-    quiet: bool = False,
-) -> None:
-    """
-    Run drgn's :ref:`interactive-mode` via ptpython
-
-    :param prog: Pre-configured program to run against. Available as a global
-        named ``prog`` in the CLI.
-    :param banner_func: Optional function to modify the printed banner. Called
-        with the default banner, and must return a string to use as the new
-        banner. The default banner does not include the drgn version, which can
-        be retrieved via :func:`version_header()`.
-    :param globals_func: Optional function to modify globals provided to the
-        session. Called with a dictionary of default globals, and must return a
-        dictionary to use instead.
-    :param quiet: Whether to suppress non-fatal warnings.
-    """
-    init_globals: Dict[str, Any] = {
-        "prog": prog,
-        "drgn": drgn,
-        "__name__": "__main__",
-        "__doc__": None,
-    }
-    drgn_globals = [
-        "NULL",
-        "Object",
-        "cast",
-        "container_of",
-        "execscript",
-        "offsetof",
-        "reinterpret",
-        "sizeof",
-        "stack_trace",
-    ]
-    for attr in drgn_globals:
-        init_globals[attr] = getattr(drgn, attr)
-
-    banner = f"""\
-For help, type help(drgn).
->>> import drgn
->>> from drgn import {", ".join(drgn_globals)}
->>> from drgn.helpers.common import *"""
-
-    module = importlib.import_module("drgn.helpers.common")
-    for name in module.__dict__["__all__"]:
-        init_globals[name] = getattr(module, name)
-    if prog.flags & drgn.ProgramFlags.IS_LINUX_KERNEL:
-        banner += "\n>>> from drgn.helpers.linux import *"
-        module = importlib.import_module("drgn.helpers.linux")
-        for name in module.__dict__["__all__"]:
-            init_globals[name] = getattr(module, name)
-
-    if banner_func:
-        banner = banner_func(banner)
-    if globals_func:
-        init_globals = globals_func(init_globals)
-
-    old_path = list(sys.path)
-    try:
-        old_default_prog = drgn.get_default_prog()
-    except drgn.NoDefaultProgramError:
-        old_default_prog = None
-    # The ptpython history file format is different from a standard readline
-    # history file since it must handle multi-line input, and it includes some
-    # metadata as well. Use a separate history format, even though it would be
-    # nice to share.
+def interact(local: Dict[str, Any], banner: str):
     histfile = os.path.expanduser("~/.drgn_history.ptpython")
-    try:
-        sys.path.insert(0, "")
-
-        drgn.set_default_prog(prog)
-
-        print(banner)
-        embed(
-            globals=init_globals,
-            history_filename=histfile,
-            title="drgn",
-            configure=configure,
-        )
-    finally:
-        drgn.set_default_prog(old_default_prog)
-        sys.path[:] = old_path
+    print(banner)
+    embed(globals=local, history_filename=histfile, title="drgn", configure=configure)
 
 
 if __name__ == "__main__":
     # Muck around with the internals of drgn: swap out run_interactive() with our
     # ptpython version, and then call main as if nothing happened.
-    drgn.cli.run_interactive = run_interactive
+    drgn.cli.interact = interact
     drgn.cli._main()

docs/advanced_usage.rst

@@ -114,6 +114,13 @@ Some of drgn's behavior can be modified through environment variables:
     Whether drgn should use libkdumpfile for ELF vmcores (0 or 1). The default
     is 0. This functionality will be removed in the future.
 
+``DRGN_USE_PYREPL``
+    Whether drgn should attempt to use the improved REPL (pyrepl) from Python
+    3.13. This provides colored output and multiline editing, among other
+    features. The default is 1.  Unfortunately, Python has no public API to use
+    these features, so drgn must rely on internal implementation details. Set
+    this to 0 to disable this feature.
+
 ``DRGN_USE_SYS_MODULE``
     Whether drgn should use ``/sys/module`` to find information about loaded
     kernel modules for the running kernel instead of getting them from the core

drgn/cli.py

@@ -6,19 +6,18 @@
 
 import argparse
 import builtins
-import code
 import importlib
 import logging
 import os
 import os.path
 import pkgutil
-import readline
 import runpy
 import shutil
 import sys
 from typing import Any, Callable, Dict, Optional
 
 import drgn
+from drgn.internal.repl import interact, readline
 from drgn.internal.rlcompleter import Completer
 from drgn.internal.sudohelper import open_via_sudo
 
@@ -435,7 +434,7 @@ def run_interactive(
         drgn.set_default_prog(prog)
 
         try:
-            code.interact(banner=banner, exitmsg="", local=init_globals)
+            interact(init_globals, banner)
         finally:
             try:
                 readline.write_history_file(histfile)

drgn/internal/repl.py

@@ -0,0 +1,47 @@
+# Copyright (c) 2024, Oracle and/or its affiliates.
+# SPDX-License-Identifier: LGPL-2.1-or-later
+
+"""Compatibility shim between drgn and the pyrepl/code modules"""
+
+import os
+import sys
+from typing import Any, Dict
+
+__all__ = ("interact", "readline")
+
+# Python 3.13 introduces a new REPL implemented by the "_pyrepl" internal
+# module. It includes features such as colored output and multiline editing.
+# Unfortunately, there is no public API exposing these abilities to users, even
+# in the "code" module. We'd like to give the best experience possible, so we'll
+# detect _pyrepl and try to use it where possible.
+try:
+    # Since this mucks with internals, add a knob that can be used to disable it
+    # and use the traditional REPL.
+    if os.environ.get("DRGN_USE_PYREPL") in ("0", "n", "N", "false", "False"):
+        raise ModuleNotFoundError()
+
+    # Unfortunately, the typeshed library behind mypy explicitly removed type
+    # stubs for these modules. This makes sense as they are private APIs, but it
+    # means we need to disable mypy checks.
+    from _pyrepl import readline  # type: ignore
+    from _pyrepl.console import InteractiveColoredConsole  # type: ignore
+    from _pyrepl.simple_interact import (  # type: ignore
+        run_multiline_interactive_console,
+    )
+
+    # This _setup() function clobbers the readline completer, but it is
+    # protected so it only runs once. Call it early so that our overridden
+    # completer doesn't get clobbered.
+    readline._setup({})
+
+    def interact(local: Dict[str, Any], banner: str) -> None:
+        console = InteractiveColoredConsole(local)
+        print(banner, file=sys.stderr)
+        run_multiline_interactive_console(console)
+
+except (ModuleNotFoundError, ImportError):
+    import code
+    import readline
+
+    def interact(local: Dict[str, Any], banner: str) -> None:
+        code.interact(banner=banner, exitmsg="", local=local)

drgn/internal/rlcompleter.py

@@ -6,9 +6,10 @@
 import builtins
 import keyword
 import re
-import readline
 from typing import Any, Dict, List, Optional
 
+from drgn.internal.repl import readline
+
 _EXPR_RE = re.compile(
     r"""
 (