From 906fa90202f668e25ce5f5dd5910d69c8add3e63 Mon Sep 17 00:00:00 2001
From: Omar Sandoval <osandov@osandov.com>
Date: Fri, 28 Jun 2019 12:48:45 -0700
Subject: [PATCH] python: add execscript()

On many occasions, I've wanted to build and iterate on a script while
exploring in an interactive session. This has involved various
workarounds, like copy-and-paste of the work-in-progress script into the
REPL, using importlib.reload(), or adding code.interact() at the end of
the script. These workflows aren't very convenient.

Add execscript(), which runs a script as if it was run from the REPL,
adding all defined functions and variables to the global scope so that
iterating on the script further is easy.
---
 docs/api_reference.rst |  5 ++++
 drgn/__init__.py       | 64 ++++++++++++++++++++++++++++++++++++++++++
 drgn/internal/cli.py   |  2 +-
 3 files changed, 70 insertions(+), 1 deletion(-)

diff --git a/docs/api_reference.rst b/docs/api_reference.rst
index ff0318ac..bf679338 100644
--- a/docs/api_reference.rst
+++ b/docs/api_reference.rst
@@ -1167,6 +1167,11 @@ can be used just like types obtained from :meth:`Program.type()`.
     :type qualifiers: Qualifiers or None
     :rtype: Type
 
+Miscellaneous
+-------------
+
+.. autofunction:: execscript
+
 Exceptions
 ----------
 
diff --git a/drgn/__init__.py b/drgn/__init__.py
index 375f33d0..6761fbd6 100644
--- a/drgn/__init__.py
+++ b/drgn/__init__.py
@@ -40,6 +40,9 @@ that package should be considered implementation details and should not
 be used.
 """
 
+import runpy
+import sys
+
 from _drgn import (
     Architecture,
     FaultError,
@@ -113,3 +116,64 @@ __all__ = [
     'union_type',
     'void_type',
 ]
+
+
+def execscript(path, *args):
+    """
+    Execute a script.
+
+    The script is executed in the same context as the caller: currently defined
+    globals are available to the script, and globals defined by the script are
+    added back to the calling context.
+
+    This is most useful for executing scripts from interactive mode. For
+    example, you could have a script named ``tasks.py``:
+
+    .. code-block:: python3
+
+        import sys
+
+        \"\"\"
+        Get all tasks in a given state.
+        \"\"\"
+
+        # From include/linux/sched.h.
+        def task_state_index(task):
+            task_state = task.state.value_()
+            if task_state == 0x402:  # TASK_IDLE
+                return 8
+            else:
+                state = (task_state | task.exit_state.value_()) & 0x7f
+                return state.bit_length()
+
+        def task_state_to_char(task):
+            return 'RSDTtXZPI'[task_state_index(task)]
+
+        tasks = [
+            task for task in for_each_task(prog)
+            if task_state_to_char(task) == sys.argv[1]
+        ]
+
+    Then, you could execute it and use the defined variables and functions:
+
+    >>> execscript('tasks.py', 'R')
+    >>> tasks[0].comm
+    (char [16])"python3"
+    >>> task_state_to_char(find_task(prog, 1))
+    'S'
+
+    :param str path: File path of the script.
+    :param str \*args: Zero or more additional arguments to pass to the script.
+        This is a :ref:`variable argument list <python:tut-arbitraryargs>`.
+    """
+    old_argv = sys.argv
+    sys.argv = [path]
+    sys.argv.extend(args)
+    try:
+        old_globals = sys._getframe(1).f_globals
+        new_globals = runpy.run_path(path, init_globals=old_globals,
+                                     run_name='__main__')
+        old_globals.clear()
+        old_globals.update(new_globals)
+    finally:
+        sys.argv = old_argv
diff --git a/drgn/internal/cli.py b/drgn/internal/cli.py
index afb06df2..06dc6554 100644
--- a/drgn/internal/cli.py
+++ b/drgn/internal/cli.py
@@ -110,7 +110,7 @@ def main() -> None:
         from drgn.internal.rlcompleter import Completer
 
         init_globals['drgn'] = drgn
-        drgn_globals = ['cast', 'container_of', 'NULL', 'Object',
+        drgn_globals = ['cast', 'container_of', 'execscript', 'NULL', 'Object',
                         'reinterpret']
         for attr in drgn_globals:
             init_globals[attr] = getattr(drgn, attr)