Commands - dot-prefixed, stdlib only

.pre-commit-config.yaml

@@ -19,6 +19,7 @@ repos:
     -   id: mypy
         args: [--show-error-codes, --strict, --no-warn-return-any, --no-warn-unused-ignores]
         files: ^(drgn/.*\.py|_drgn.pyi|_drgn_util/.*\.py|tools/.*\.py)$
+        additional_dependencies: ["types-setuptools"]
 -   repo: https://github.com/pre-commit/pre-commit-hooks
     rev: v4.6.0
     hooks:

contrib/ptdrgn.py

@@ -11,19 +11,25 @@
 
 Requires: "pip install ptpython" which brings in pygments and prompt_toolkit
 """
+import builtins
 import functools
 import os
+import re
 import shutil
 from typing import Any, Dict, Set
 
-from prompt_toolkit.completion import Completer
+import ptpython.repl
+from prompt_toolkit.completion import Completer, Completion
+from prompt_toolkit.document import Document
 from prompt_toolkit.formatted_text import PygmentsTokens
 from ptpython import embed
-from ptpython.repl import run_config
+from ptpython.repl import PythonRepl, run_config
+from ptpython.validator import PythonValidator
 from pygments.lexers.c_cpp import CLexer
 
 import drgn
 import drgn.cli
+from drgn.cli import Command
 
 
 class DummyForRepr:
@@ -68,10 +74,18 @@ def _object_fields() -> Set[str]:
 class ReorderDrgnObjectCompleter(Completer):
     """A completer which puts Object member fields above Object defaults"""
 
-    def __init__(self, c: Completer):
+    def __init__(self, c: Completer, commands: Dict[str, Command]):
         self.c = c
+        self.__commands = [f".{cmd}" for cmd in commands.keys()]
+        self.__command_re = re.compile(r"\.[\S]+")
 
     def get_completions(self, document, complete_event):
+        text = document.text_before_cursor
+        if self.__command_re.fullmatch(text):
+            return [
+                Completion(cmd, start_position=-len(text))
+                for cmd in self.__commands if cmd.startswith(text)
+            ]
         completions = list(self.c.get_completions(document, complete_event))
         if not completions:
             return completions
@@ -119,18 +133,47 @@ def _format_result_output(result: object):
 
     repl._format_result_output = _format_result_output
     run_config(repl)
-    repl._completer = ReorderDrgnObjectCompleter(repl._completer)
-    repl.completer = ReorderDrgnObjectCompleter(repl.completer)
+    repl._completer = ReorderDrgnObjectCompleter(repl._completer, repl._commands)
+    repl.completer = ReorderDrgnObjectCompleter(repl.completer, repl._commands)
 
 
-def interact(local: Dict[str, Any], banner: str):
+class DrgnPythonValidator(PythonValidator):
+
+    def validate(self, document: Document) -> None:
+        if document.text.lstrip().startswith("."):
+            return
+        return super().validate(document)
+
+
+class DrgnPythonRepl(PythonRepl):
+
+    def __init__(self, *args: Any, **kwargs: Any):
+        super().__init__(*args, **kwargs, _validator=DrgnPythonValidator())
+
+    def __run_command(self, line: str) -> object:
+        cmd_name = line.split(maxsplit=1)[0][1:]
+        if cmd_name not in self._commands:
+            print(f"{cmd_name}: drgn command not found")
+            return None
+        cmd = self._commands[cmd_name]
+        locals = self.get_locals()
+        prog = locals["prog"]
+        setattr(builtins, "_", cmd(prog, line, locals))
+
+    def eval(self, line: str) -> object:
+        if line.lstrip().startswith('.'):
+            return self.__run_command(line)
+        return super().eval(line)
+
+
+def interact(local: Dict[str, Any], banner: str, commands: Dict[str, Command]):
+    DrgnPythonRepl._commands = commands
     histfile = os.path.expanduser("~/.drgn_history.ptpython")
     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.
+    ptpython.repl.PythonRepl = DrgnPythonRepl
     drgn.cli.interact = interact
     drgn.cli._main()

docs/user_guide.rst

@@ -375,6 +375,7 @@ adds a few nice features, including:
 * Tab completion
 * Automatic import of relevant helpers
 * Pretty printing of objects and types
+* Helper commands
 
 The default behavior of the Python `REPL
 <https://en.wikipedia.org/wiki/Read%E2%80%93eval%E2%80%93print_loop>`_ is to
@@ -409,6 +410,30 @@ explicitly::
             int counter;
     } atomic_t
 
+Interactive Commands
+^^^^^^^^^^^^^^^^^^^^
+
+When running in interactive mode, the drgn CLI accepts an extensible set of
+commands. Any input which starts with a ``.`` will be interpreted as a command,
+rather than as Python code. For example, you can use the ``.x`` command to
+execute a script, similar to the :func:`execscript()` function::
+
+    $ cat myscript.py
+    import sys
+    print("Executing myscript.py with arguments: " + str(sys.argv[1:]))
+    $ sudo drgn
+    >>> .x myscript.py argument
+    Executing myscript.py with arguments ['argument']
+    >>> execscript("myscript.py", "argument")
+    Executing myscript.py with arguments ['argument']
+
+Not only are the interactive commands less verbose than their equivalent Python
+code, but they are also user-extensible. You can define a
+:class:`Command <drgn.cli.Command>` function and either register it with the
+:func:`@command <drgn.cli.command>` decorator, or provide it to
+:func:`run_interactive <drgn.cli.run_interactive>` using the argument
+``commands_func``.
+
 Next Steps
 ----------
 

drgn/__init__.py

@@ -44,7 +44,7 @@
 import pkgutil
 import sys
 import types
-from typing import Union
+from typing import Any, BinaryIO, Dict, Optional, Union
 
 from _drgn import (
     NULL,
@@ -160,7 +160,6 @@
 if sys.version_info >= (3, 8):
     _open_code = io.open_code  # novermin
 else:
-    from typing import BinaryIO
 
     def _open_code(path: str) -> BinaryIO:
         return open(path, "rb")
@@ -180,7 +179,7 @@ def _open_code(path: str) -> BinaryIO:
 )
 
 
-def execscript(path: str, *args: str) -> None:
+def execscript(path: str, *args: str, globals: Optional[Dict[str, Any]] = None) -> None:
     """
     Execute a script.
 
@@ -222,6 +221,7 @@ def task_exe_path(task):
     :param path: File path of the script.
     :param args: Zero or more additional arguments to pass to the script. This
         is a :ref:`variable argument list <python:tut-arbitraryargs>`.
+    :param globals: If provided, globals to use instead of the caller's.
     """
     # This is based on runpy.run_path(), which we can't use because we want to
     # update globals even if the script throws an exception.
@@ -246,7 +246,10 @@ def task_exe_path(task):
         module.__file__ = path
         module.__cached__ = None  # type: ignore[attr-defined]
 
-        caller_globals = sys._getframe(1).f_globals
+        if globals is not None:
+            caller_globals = globals
+        else:
+            caller_globals = sys._getframe(1).f_globals
         caller_special_globals = {
             name: caller_globals[name]
             for name in _special_globals

drgn/cli.py

@@ -14,9 +14,11 @@
 import runpy
 import shutil
 import sys
-from typing import Any, Callable, Dict, Optional
+import textwrap
+from typing import Any, Callable, Dict, Iterable, Optional
 
 import drgn
+from drgn import Program
 from drgn.internal.repl import interact, readline
 from drgn.internal.rlcompleter import Completer
 from drgn.internal.sudohelper import open_via_sudo
@@ -25,6 +27,128 @@
 
 logger = logging.getLogger("drgn")
 
+Command = Callable[[Program, str, Dict[str, Any]], Any]
+"""
+A command which can be executed in the drgn CLI
+
+The drgn CLI allows for shell-like commands to be executed. Any input to the CLI
+which begins with a ``.`` is interpreted as a command rather than a Python
+statement. Commands are simply callables which take three arguments:
+
+- a :class:`drgn.Program`
+- a ``str`` which contains the command line, and
+- a dictionary of local variables in the CLI (``Dict[str, Any]``)
+
+For example, the following is a command function::
+
+    def hello_world(prog, cmdline, locals_):
+        print("hello world!")
+        print(f"your command: {cmdline}")
+        print(f"kernel command line: {prog['saved_command_line']}")
+        locals_["secret"] = 42
+
+The command, if registered with the drgn CLI, might be used like so:
+
+    >>> .hello
+    hello world!
+    your command: .hello
+    kernel command line: (char *)0xffff9ea9cf7c3600 = "quiet splash"
+
+User-defined commands may be provided to :func:`run_interactive()` via the
+``commands_func`` argument. Commands can also be registered so that they are
+included in drgn's default command set using the :func:`command` decorator.
+"""
+
+_COMMANDS: Dict[str, Command] = {}
+
+
+def command(name: str) -> Callable[[Command], Command]:
+    """
+    A decorator for registering a command function
+
+    Example usage:
+
+    >>> @drgn.cli.command("hello")
+    ... def hello(prog, line, locals_):
+    ...     print("hello world")
+
+    The decorator will be added to Drgn's default command set. Please keep in
+    mind that the decorator is evaluated when your module is imported. If you'd
+    like to extend drgn's default set of commands, then you should ensure your
+    module is imported before the CLI starts.
+    """
+
+    def decorator(cmd: Command) -> Command:
+        _COMMANDS[name] = cmd
+        return cmd
+
+    return decorator
+
+
+def all_commands() -> Dict[str, Command]:
+    """
+    Returns all registered drgn CLI commands
+
+    By default, only the commands which are built-in to drgn, or registered via
+    :func:`command`, are returned. Since decorators are evaluated at module load
+    time, any command defined in a module which is not imported prior to the
+    drgn CLI being run, will not be loaded.
+
+    However, drgn can allow user command modules to be loaded and registered by
+    using the ``drgn.command.v1`` `entry point
+    <https://setuptools.pypa.io/en/latest/pkg_resources.html#entry-points>`_.
+    Third-party packages can define a module as an entry point which should be
+    imported, typically in the setup.py::
+
+        setup(
+            ...,
+            entry_points={
+                "drgn.command.v1": {
+                    "my_module = fully_qualified.module_path",
+                },
+            }
+        )
+
+    In the above example, a function defined within the module
+    ``fully_qualified.module_path`` and registered with :func:`command`, would
+    always be included in the drgn CLI and this functon if the relevant package
+    is installed.
+    """
+    import drgn.helpers.common.commands  # noqa
+
+    # The importlib.metadata API is included in Python 3.8+. Normally, one might
+    # simply try to import it, catching the ImportError and falling back to the
+    # older API. However, the API was _transitional_ in 3.8 and 3.9, and it is
+    # different enough to break callers compared to the non-transitional API. So
+    # here we are, using sys.version_info like heathens.
+    if sys.version_info >= (3, 10):
+        from importlib.metadata import entry_points  # novermin
+    else:
+        import pkg_resources
+
+        def entry_points(group: str) -> Iterable[pkg_resources.EntryPoint]:
+            return pkg_resources.iter_entry_points(group)
+
+    # Drgn command "entry points" are simply modules. The act of loading /
+    # importing them will result in their @command decorators being executed,
+    # and _COMMANDS will be updated properly.
+    for entry_point in entry_points(group="drgn.command.v1"):  # type: ignore
+        entry_point.load()  # type: ignore
+
+    return _COMMANDS.copy()
+
+
+def help_command(commands: Dict[str, Command]) -> Command:
+    def help(prog: drgn.Program, line: str, locals_: Dict[str, Any]) -> None:
+        try:
+            width = os.get_terminal_size().columns
+        except OSError:
+            width = 80
+        print("Drgn CLI commands:\n")
+        print(textwrap.fill(" ".join(commands.keys()), width=width))
+
+    return help
+
 
 class _LogFormatter(logging.Formatter):
     _LEVELS = (
@@ -333,6 +457,7 @@ def run_interactive(
     prog: drgn.Program,
     banner_func: Optional[Callable[[str], str]] = None,
     globals_func: Optional[Callable[[Dict[str, Any]], Dict[str, Any]]] = None,
+    commands_func: Optional[Callable[[Dict[str, Command]], Dict[str, Command]]] = None,
     quiet: bool = False,
 ) -> None:
     """
@@ -350,6 +475,9 @@ def run_interactive(
     :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 commands_func: Optional function to modify the command list which is
+        used for the session. Called with a dictionary of commands, and must
+        return a dictionary to use instead.
     :param quiet: Ignored. Will be removed in the future.
 
     .. note::
@@ -407,6 +535,11 @@ def run_interactive(
     if globals_func:
         init_globals = globals_func(init_globals)
 
+    commands = all_commands()
+    if commands_func:
+        commands = commands_func(commands)
+    commands["help"] = help_command(commands)
+
     old_path = list(sys.path)
     old_displayhook = sys.displayhook
     old_history_length = readline.get_history_length()
@@ -426,15 +559,15 @@ def run_interactive(
 
         readline.set_history_length(1000)
         readline.parse_and_bind("tab: complete")
-        readline.set_completer(Completer(init_globals).complete)
+        readline.set_completer(Completer(init_globals, commands).complete)
 
         sys.path.insert(0, "")
         sys.displayhook = _displayhook
 
         drgn.set_default_prog(prog)
 
         try:
-            interact(init_globals, banner)
+            interact(init_globals, banner, commands)
         finally:
             try:
                 readline.write_history_file(histfile)